From dfd784d4fa6cbe0bcf2a7e97e15434ee080fad5d Mon Sep 17 00:00:00 2001
From: Ben Toogood <ben@micro.mu>
Date: Thu, 16 Apr 2020 15:01:16 +0100
Subject: [PATCH] Set authorization header on grpc stream

---
 grpc.go | 46 +++++++++++++++++++++++++++++++---------------
 1 file changed, 31 insertions(+), 15 deletions(-)

diff --git a/grpc.go b/grpc.go
index 9d2fd59..30dd97f 100644
--- a/grpc.go
+++ b/grpc.go
@@ -131,21 +131,10 @@ func (g *grpcClient) call(ctx context.Context, node *registry.Node, req client.R
 	// set the content type for the request
 	header["x-content-type"] = req.ContentType()
 
-	// if the caller specifies using service token or no token
-	// was passed with the request, set the service token
-	var srvToken string
-	if g.opts.Auth != nil && g.opts.Auth.Options().Token != nil {
-		srvToken = g.opts.Auth.Options().Token.AccessToken
-	}
-	if (opts.ServiceToken || len(header["authorization"]) == 0) && len(srvToken) > 0 {
-		header["authorization"] = auth.BearerScheme + srvToken
-	}
-
-	// fall back to using the authorization token set in config,
-	// this enables the CLI to provide a token
-	if len(header["authorization"]) == 0 {
-		if token, err := config.Get("micro", "auth", "token"); err == nil && len(token) > 0 {
-			header["authorization"] = auth.BearerScheme + token
+	// set the authorization header
+	if opts.ServiceToken || len(header["authorization"]) == 0 {
+		if h := g.authorizationHeader(); len(h) > 0 {
+			header["authorization"] = h
 		}
 	}
 
@@ -227,6 +216,13 @@ func (g *grpcClient) stream(ctx context.Context, node *registry.Node, req client
 	// set the content type for the request
 	header["x-content-type"] = req.ContentType()
 
+	// set the authorization header
+	if opts.ServiceToken || len(header["authorization"]) == 0 {
+		if h := g.authorizationHeader(); len(h) > 0 {
+			header["authorization"] = h
+		}
+	}
+
 	md := gmetadata.New(header)
 	ctx = gmetadata.NewOutgoingContext(ctx, md)
 
@@ -313,6 +309,26 @@ func (g *grpcClient) stream(ctx context.Context, node *registry.Node, req client
 	}, nil
 }
 
+func (g *grpcClient) authorizationHeader() string {
+	// if the caller specifies using service token or no token
+	// was passed with the request, set the service token
+	var srvToken string
+	if g.opts.Auth != nil && g.opts.Auth.Options().Token != nil {
+		srvToken = g.opts.Auth.Options().Token.AccessToken
+	}
+	if len(srvToken) > 0 {
+		return auth.BearerScheme + srvToken
+	}
+
+	// fall back to using the authorization token set in config,
+	// this enables the CLI to provide a token
+	if token, err := config.Get("micro", "auth", "token"); err == nil && len(token) > 0 {
+		return auth.BearerScheme + token
+	}
+
+	return ""
+}
+
 func (g *grpcClient) poolMaxStreams() int {
 	if g.opts.Context == nil {
 		return DefaultPoolMaxStreams