multi-service: preserve context values

Author: jessepeterson

service/multi/context.go

@@ -0,0 +1,41 @@
+package multi
+
+import (
+	"context"
+	"time"
+)
+
+// ContextWithoutCancel returns a derived context that points to the parent context
+// and is not canceled when parent is canceled.
+// The returned context returns no Deadline or Err, and its Done channel is nil.
+// Calling Cause on the returned context returns nil.
+func ContextWithoutCancel(parent context.Context) context.Context {
+	if parent == nil {
+		panic("cannot create context from nil parent")
+	}
+	return withoutCancelCtx{parent}
+}
+
+type withoutCancelCtx struct {
+	c context.Context
+}
+
+func (withoutCancelCtx) Deadline() (deadline time.Time, ok bool) {
+	return
+}
+
+func (withoutCancelCtx) Done() <-chan struct{} {
+	return nil
+}
+
+func (withoutCancelCtx) Err() error {
+	return nil
+}
+
+func (c withoutCancelCtx) Value(key interface{}) interface{} {
+	return c.c.Value(key)
+}
+
+// func (c withoutCancelCtx) String() string {
+// 	return c.c.String() + ".WithoutCancel"
+// }

service/multi/context_test.go

@@ -0,0 +1,57 @@
+package multi
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/micromdm/nanolib/log"
+	"github.com/micromdm/nanomdm/mdm"
+	"github.com/micromdm/nanomdm/service"
+	"github.com/micromdm/nanomdm/test"
+)
+
+type ctxTest1 struct{}
+
+type testSvc struct {
+	wg      *sync.WaitGroup
+	capture string
+	service.CheckinAndCommandService
+}
+
+func (ts *testSvc) Authenticate(r *mdm.Request, _ *mdm.Authenticate) error {
+	ts.capture, _ = r.Context.Value(&ctxTest1{}).(string)
+	ts.wg.Done()
+	return nil
+}
+
+func TestContextPassthru(t *testing.T) {
+	nopSvc1 := &test.NopService{}
+
+	var ctx context.Context = context.Background()
+
+	ctx = context.WithValue(ctx, &ctxTest1{}, "test-ctx-val")
+
+	r := &mdm.Request{Context: ctx}
+
+	var wg sync.WaitGroup
+
+	wg.Add(1)
+	ts := &testSvc{
+		wg:                       &wg,
+		CheckinAndCommandService: &test.NopService{},
+	}
+
+	multi := New(log.NopLogger, nopSvc1, ts)
+
+	err := multi.Authenticate(r, &mdm.Authenticate{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	wg.Wait()
+
+	if have, want := ts.capture, "test-ctx-val"; have != want {
+		t.Errorf("have: %v, want: %v", have, want)
+	}
+}

service/multi/multi.go

@@ -19,7 +19,6 @@ import (
 type MultiService struct {
 	logger log.Logger
 	svcs   []service.CheckinAndCommandService
-	ctx    context.Context
 }
 
 func New(logger log.Logger, svcs ...service.CheckinAndCommandService) *MultiService {
@@ -29,7 +28,6 @@ func New(logger log.Logger, svcs ...service.CheckinAndCommandService) *MultiServ
 	return &MultiService{
 		logger: logger,
 		svcs:   svcs,
-		ctx:    context.Background(),
 	}
 }
 
@@ -52,7 +50,7 @@ func (ms *MultiService) runOthers(ctx context.Context, r errorRunner) {
 // RequestWithContext returns a clone of r and sets its context to ctx.
 func (ms *MultiService) RequestWithContext(r *mdm.Request) *mdm.Request {
 	r2 := r.Clone()
-	r2.Context = ms.ctx
+	r2.Context = ContextWithoutCancel(r.Context)
 	return r2
 }