From a4bf341022f076582fc2e9a802ce170a8938f81d Mon Sep 17 00:00:00 2001
From: Victor Vrantchan <groob@users.noreply.github.com>
Date: Thu, 30 Nov 2017 17:37:49 -0500
Subject: [PATCH] use the request context with net/http handler (#1696)

---
 transport/go16.go           | 6 ++++++
 transport/go17.go           | 6 ++++++
 transport/handler_server.go | 6 +++---
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/transport/go16.go b/transport/go16.go
index 7cffee11e..5babcf9b8 100644
--- a/transport/go16.go
+++ b/transport/go16.go
@@ -22,6 +22,7 @@ package transport
 
 import (
 	"net"
+	"net/http"
 
 	"google.golang.org/grpc/codes"
 
@@ -43,3 +44,8 @@ func ContextErr(err error) StreamError {
 	}
 	return streamErrorf(codes.Internal, "Unexpected error from context packet: %v", err)
 }
+
+// contextFromRequest returns a background context.
+func contextFromRequest(r *http.Request) context.Context {
+	return context.Background()
+}
diff --git a/transport/go17.go b/transport/go17.go
index 2464e69fa..b7fa6bdb9 100644
--- a/transport/go17.go
+++ b/transport/go17.go
@@ -23,6 +23,7 @@ package transport
 import (
 	"context"
 	"net"
+	"net/http"
 
 	"google.golang.org/grpc/codes"
 
@@ -44,3 +45,8 @@ func ContextErr(err error) StreamError {
 	}
 	return streamErrorf(codes.Internal, "Unexpected error from context packet: %v", err)
 }
+
+// contextFromRequest returns a context from the HTTP Request.
+func contextFromRequest(r *http.Request) context.Context {
+	return r.Context()
+}
diff --git a/transport/handler_server.go b/transport/handler_server.go
index 7e0fdb359..27c4ebb5f 100644
--- a/transport/handler_server.go
+++ b/transport/handler_server.go
@@ -284,12 +284,12 @@ func (ht *serverHandlerTransport) WriteHeader(s *Stream, md metadata.MD) error {
 func (ht *serverHandlerTransport) HandleStreams(startStream func(*Stream), traceCtx func(context.Context, string) context.Context) {
 	// With this transport type there will be exactly 1 stream: this HTTP request.
 
-	var ctx context.Context
+	ctx := contextFromRequest(ht.req)
 	var cancel context.CancelFunc
 	if ht.timeoutSet {
-		ctx, cancel = context.WithTimeout(context.Background(), ht.timeout)
+		ctx, cancel = context.WithTimeout(ctx, ht.timeout)
 	} else {
-		ctx, cancel = context.WithCancel(context.Background())
+		ctx, cancel = context.WithCancel(ctx)
 	}
 
 	// requestOver is closed when either the request's context is done