From e5d20af097e6e145eaaee85961dd61b01c0d66cd Mon Sep 17 00:00:00 2001
From: James Crasta <jcrasta@underarmour.com>
Date: Tue, 9 Jun 2015 13:15:17 -0600
Subject: [PATCH] Receive GRPC metadata from HTTP headers.

* Add support for Grpc-Metadata-(varname) headers
* Modify codegen to do metadata annotation
---
 .../gengateway/template.go                    |  2 +-
 runtime/context.go                            | 31 +++++++++++++++++
 runtime/context_test.go                       | 34 +++++++++++++++++++
 3 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 runtime/context.go
 create mode 100644 runtime/context_test.go

diff --git a/protoc-gen-grpc-gateway/gengateway/template.go b/protoc-gen-grpc-gateway/gengateway/template.go
index 8d1d3e66a85..8f4efb7eca9 100644
--- a/protoc-gen-grpc-gateway/gengateway/template.go
+++ b/protoc-gen-grpc-gateway/gengateway/template.go
@@ -220,7 +220,7 @@ func Register{{$svc.GetName}}Handler(ctx context.Context, mux *runtime.ServeMux,
 	{{range $m := $svc.Methods}}
 	{{range $b := $m.Bindings}}
 	mux.Handle({{$b.HTTPMethod | printf "%q"}}, pattern_{{$svc.GetName}}_{{$m.GetName}}_{{$b.Index}}, func(w http.ResponseWriter, req *http.Request, pathParams map[string]string) {
-		resp, err := request_{{$svc.GetName}}_{{$m.GetName}}_{{$b.Index}}(ctx, client, req, pathParams)
+		resp, err := request_{{$svc.GetName}}_{{$m.GetName}}_{{$b.Index}}(runtime.AnnotateContext(ctx, req), client, req, pathParams)
 		if err != nil {
 			runtime.HTTPError(w, err)
 			return
diff --git a/runtime/context.go b/runtime/context.go
new file mode 100644
index 00000000000..b8692b798ef
--- /dev/null
+++ b/runtime/context.go
@@ -0,0 +1,31 @@
+package runtime
+
+import (
+	"net/http"
+	"strings"
+
+	"golang.org/x/net/context"
+	"google.golang.org/grpc/metadata"
+)
+
+const metadataHeaderPrefix = "Grpc-Metadata-"
+
+/*
+AnnotateContext adds context information such as metadata from the request.
+
+If there are no metadata headers in the request, then the context returned
+will be the same context.
+*/
+func AnnotateContext(ctx context.Context, req *http.Request) context.Context {
+	var pairs []string
+	for key, val := range req.Header {
+		if strings.HasPrefix(key, metadataHeaderPrefix) {
+			pairs = append(pairs, key[len(metadataHeaderPrefix):], val[0])
+		}
+	}
+
+	if len(pairs) != 0 {
+		ctx = metadata.NewContext(ctx, metadata.Pairs(pairs...))
+	}
+	return ctx
+}
diff --git a/runtime/context_test.go b/runtime/context_test.go
new file mode 100644
index 00000000000..4e1bec641f6
--- /dev/null
+++ b/runtime/context_test.go
@@ -0,0 +1,34 @@
+package runtime_test
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/gengo/grpc-gateway/runtime"
+	"golang.org/x/net/context"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestAnnotateContext(t *testing.T) {
+	ctx := context.Background()
+
+	request, _ := http.NewRequest("GET", "http://localhost", nil)
+	request.Header = http.Header{}
+	annotated := runtime.AnnotateContext(ctx, request)
+	if annotated != ctx {
+		t.Errorf("AnnotateContext(ctx, request) = %v; want %v", annotated, ctx)
+	}
+	request.Header.Add("Grpc-Metadata-FooBar", "Value1")
+	request.Header.Add("Grpc-Metadata-Foo-BAZ", "Value2")
+	annotated = runtime.AnnotateContext(ctx, request)
+	md, ok := metadata.FromContext(annotated)
+	if !ok || len(md) != 2 {
+		t.Errorf("Expected 2 metadata items in context; got %v", md)
+	}
+	if md["Foobar"] != "Value1" {
+		t.Errorf("md[\"Foobar\"] = %v; want %v", md["Foobar"], "Value1")
+	}
+	if md["Foo-Baz"] != "Value2" {
+		t.Errorf("md[\"Foo-Baz\"] = %v want %v", md["Foo-Baz"], "Value2")
+	}
+}