From d1d3bc022f64089ad7f3f1c7dc124e9053a8be32 Mon Sep 17 00:00:00 2001
From: Sadlil Rhythom <sadlil.rhythm@gmail.com>
Date: Fri, 10 Feb 2017 18:36:44 +0600
Subject: [PATCH] Added auth provider token to context (#6)

---
 pkg/tiller/server.go | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/pkg/tiller/server.go b/pkg/tiller/server.go
index 1b83770a1..2b018ffa0 100644
--- a/pkg/tiller/server.go
+++ b/pkg/tiller/server.go
@@ -121,11 +121,29 @@ func newStreamInterceptor() grpc.StreamServerInterceptor {
 			log.Println(err)
 			return err
 		}
-		// TODO: How to pass modified ctx?
-		return handler(srv, ss)
+
+		newStream := serverStreamWrapper{
+			ss: ss,
+			ctx: ctx,
+		}
+		return handler(srv, newStream)
 	}
 }
 
+// serverStreamWrapper wraps original ServerStream but uses modified context.
+// this modified context will be available inside handler()
+type serverStreamWrapper struct {
+	ss  grpc.ServerStream
+	ctx context.Context
+}
+
+func (w serverStreamWrapper) Context() context.Context        { return w.ctx }
+func (w serverStreamWrapper) RecvMsg(msg interface{}) error   { return w.ss.RecvMsg(msg) }
+func (w serverStreamWrapper) SendMsg(msg interface{}) error   { return w.ss.SendMsg(msg) }
+func (w serverStreamWrapper) SendHeader(md metadata.MD) error { return w.ss.SendHeader(md) }
+func (w serverStreamWrapper) SetHeader(md metadata.MD) error  { return w.ss.SetHeader(md) }
+func (w serverStreamWrapper) SetTrailer(md metadata.MD)       { w.ss.SetTrailer(md) }
+
 func splitMethod(fullMethod string) (string, string) {
 	if frags := strings.Split(fullMethod, "/"); len(frags) == 3 {
 		return frags[1], frags[2]