From ec7cd6ef31c30e37bc26cc3aba31596ac6644d21 Mon Sep 17 00:00:00 2001
From: Daniel McDonald <dmcdonald@infoblox.com>
Date: Sun, 15 Apr 2018 10:51:45 -0700
Subject: [PATCH] Support multiple metadata annotators (#586)

---
 runtime/context.go      |  4 ++--
 runtime/context_test.go | 20 ++++++++++++++++++++
 runtime/mux.go          |  4 ++--
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/runtime/context.go b/runtime/context.go
index c492b9732c9..a745074c961 100644
--- 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
index 847fc6c747d..e78a037d033 100644
--- 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
index 4ead83fc256..1d4c75760fe 100644
--- 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)
 	}
 }