Allow storing alternate transport.ServerStream implementations in context (#1904)

stream.go

@@ -608,6 +608,7 @@ type ServerStream interface {
 
 // serverStream implements a server side Stream.
 type serverStream struct {
+	ctx   context.Context
 	t     transport.ServerTransport
 	s     *transport.Stream
 	p     *parser
@@ -628,7 +629,7 @@ type serverStream struct {
 }
 
 func (ss *serverStream) Context() context.Context {
-	return ss.s.Context()
+	return ss.ctx
 }
 
 func (ss *serverStream) SetHeader(md metadata.MD) error {
@@ -731,9 +732,9 @@ func (ss *serverStream) RecvMsg(m interface{}) (err error) {
 // MethodFromServerStream returns the method string for the input stream.
 // The returned string is in the format of "/service/method".
 func MethodFromServerStream(stream ServerStream) (string, bool) {
-	s, ok := transport.StreamFromContext(stream.Context())
-	if !ok {
-		return "", ok
+	s := serverTransportStreamFromContext(stream.Context())
+	if s == nil {
+		return "", false
 	}
-	return s.Method(), ok
+	return s.Method(), true
 }