Support multiple metadata annotators (grpc-ecosystem#586)

adasari · Apr 15, 2018 · ec7cd6e · ec7cd6e
1 parent bdfb07f
commit ec7cd6e
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 4 deletions.
diff --git a/runtime/context.go b/runtime/context.go
@@ -91,8 +91,8 @@ func AnnotateContext(ctx context.Context, mux *ServeMux, req *http.Request) (con
 		return ctx, nil
 	}
 	md := metadata.Pairs(pairs...)
-	if mux.metadataAnnotator != nil {
-		md = metadata.Join(md, mux.metadataAnnotator(ctx, req))
+	for _, mda := range mux.metadataAnnotators {
+		md = metadata.Join(md, mda(ctx, req))
 	}
 	return metadata.NewOutgoingContext(ctx, md), nil
 }

diff --git a/runtime/context_test.go b/runtime/context_test.go
@@ -170,3 +170,23 @@ func TestAnnotateContext_SupportsTimeouts(t *testing.T) {
 		}
 	}
 }
+func TestAnnotateContext_SupportsCustomAnnotators(t *testing.T) {
+	md1 := func(context.Context, *http.Request) metadata.MD { return metadata.New(map[string]string{"foo": "bar"}) }
+	md2 := func(context.Context, *http.Request) metadata.MD { return metadata.New(map[string]string{"baz": "qux"}) }
+	expected := metadata.New(map[string]string{"foo": "bar", "baz": "qux"})
+	request, err := http.NewRequest("GET", "http://example.com", nil)
+	if err != nil {
+		t.Fatalf(`http.NewRequest("GET", "http://example.com", nil failed with %v; want success`, err)
+	}
+	annotated, err := runtime.AnnotateContext(context.Background(), runtime.NewServeMux(runtime.WithMetadata(md1), runtime.WithMetadata(md2)), request)
+	if err != nil {
+		t.Errorf("runtime.AnnotateContext(ctx, %#v) failed with %v; want success", request, err)
+		return
+	}
+	actual, _ := metadata.FromOutgoingContext(annotated)
+	for key, e := range expected {
+		if a, ok := actual[key]; !ok || !reflect.DeepEqual(e, a) {
+			t.Errorf("metadata.MD[%s] = %v; want %v", key, a, e)
+		}
+	}
+}
diff --git a/runtime/mux.go b/runtime/mux.go
@@ -25,7 +25,7 @@ type ServeMux struct {
 	marshalers             marshalerRegistry
 	incomingHeaderMatcher  HeaderMatcherFunc
 	outgoingHeaderMatcher  HeaderMatcherFunc
-	metadataAnnotator      func(context.Context, *http.Request) metadata.MD
+	metadataAnnotators     []func(context.Context, *http.Request) metadata.MD
 	protoErrorHandler      ProtoErrorHandlerFunc
 }
 
@@ -87,7 +87,7 @@ func WithOutgoingHeaderMatcher(fn HeaderMatcherFunc) ServeMuxOption {
 // is reading token from cookie and adding it in gRPC context.
 func WithMetadata(annotator func(context.Context, *http.Request) metadata.MD) ServeMuxOption {
 	return func(serveMux *ServeMux) {
-		serveMux.metadataAnnotator = annotator
+		serveMux.metadataAnnotators = append(serveMux.metadataAnnotators, annotator)
 	}
 }