Make the forwarder function customizable

grpc-ecosystem · Jul 13, 2015 · d45f379 · d45f379
1 parent d430203
commit d45f379
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 18 deletions.
diff --git a/examples/a_bit_of_everything.pb.gw.go b/examples/a_bit_of_everything.pb.gw.go
diff --git a/examples/echo_service.pb.gw.go b/examples/echo_service.pb.gw.go
diff --git a/protoc-gen-grpc-gateway/gengateway/template.go b/protoc-gen-grpc-gateway/gengateway/template.go
@@ -228,9 +228,9 @@ func Register{{$svc.GetName}}Handler(ctx context.Context, mux *runtime.ServeMux,
 			return
 		}
 		{{if $m.GetServerStreaming}}
-		runtime.ForwardResponseStream(w, func() (proto.Message, error) { return resp.Recv() })
+		forward_{{$svc.GetName}}_{{$m.GetName}}_{{$b.Index}}(ctx, w, req, func() (proto.Message, error) { return resp.Recv() })
 		{{else}}
-		runtime.ForwardResponseMessage(ctx, w, resp)
+		forward_{{$svc.GetName}}_{{$m.GetName}}_{{$b.Index}}(ctx, w, req, resp)
 		{{end}}
 	})
 	{{end}}
@@ -245,5 +245,13 @@ var (
 	{{end}}
 	{{end}}
 )
+
+var (
+	{{range $m := $svc.Methods}}
+	{{range $b := $m.Bindings}}
+	forward_{{$svc.GetName}}_{{$m.GetName}}_{{$b.Index}} = {{if $m.GetServerStreaming}}runtime.ForwardResponseStream{{else}}runtime.ForwardResponseMessage{{end}}
+	{{end}}
+	{{end}}
+)
 {{end}}`))
 )
diff --git a/runtime/handler.go b/runtime/handler.go
@@ -17,7 +17,7 @@ type responseStreamChunk struct {
 }
 
 // ForwardResponseStream forwards the stream from gRPC server to REST client.
-func ForwardResponseStream(w http.ResponseWriter, recv func() (proto.Message, error)) {
+func ForwardResponseStream(ctx context.Context, w http.ResponseWriter, req *http.Request, recv func() (proto.Message, error)) {
 	f, ok := w.(http.Flusher)
 	if !ok {
 		glog.Errorf("Flush not supported in %T", w)
@@ -58,8 +58,8 @@ func ForwardResponseStream(w http.ResponseWriter, recv func() (proto.Message, er
 	}
 }
 
-// ForwardResponseMessage forwards the message from gRPC server to REST client.
-func ForwardResponseMessage(ctx context.Context, w http.ResponseWriter, resp proto.Message) {
+// ForwardResponseMessage forwards the message "resp" from gRPC server to REST client.
+func ForwardResponseMessage(ctx context.Context, w http.ResponseWriter, req *http.Request, resp proto.Message) {
 	buf, err := json.Marshal(resp)
 	if err != nil {
 		glog.Errorf("Marshal error: %v", err)