From 21ea8b8e0425d421b88bb1d0cff85dbf217c4c85 Mon Sep 17 00:00:00 2001
From: jammerful <jammerful@gmail.com>
Date: Mon, 23 Sep 2019 22:59:04 -0400
Subject: [PATCH] Add Expand Headers Feature

See discussion of the feature here:
https://github.com/fullstorydev/grpcurl/issues/116
---
 cmd/grpcurl/grpcurl.go | 31 ++++++++++++++++++++++---------
 grpcurl.go             | 33 +++++++++++++++++++++++++++++++++
 grpcurl_test.go        | 17 +++++++++++++++++
 3 files changed, 72 insertions(+), 9 deletions(-)

diff --git a/cmd/grpcurl/grpcurl.go b/cmd/grpcurl/grpcurl.go
index 013e2e1..f40b1b5 100644
--- a/cmd/grpcurl/grpcurl.go
+++ b/cmd/grpcurl/grpcurl.go
@@ -52,13 +52,19 @@ var (
 	key = flags.String("key", "", prettify(`
 		File containing client private key, to present to the server. Not valid
 		with -plaintext option. Must also provide -cert option.`))
-	protoset    multiString
-	protoFiles  multiString
-	importPaths multiString
-	addlHeaders multiString
-	rpcHeaders  multiString
-	reflHeaders multiString
-	authority   = flags.String("authority", "", prettify(`
+	protoset      multiString
+	protoFiles    multiString
+	importPaths   multiString
+	addlHeaders   multiString
+	rpcHeaders    multiString
+	reflHeaders   multiString
+	expandHeaders = flags.Bool("expand-headers", false, prettify(`
+		If set any environmental variables contained contained in the
+		header string will be substituted for by its corresponding environmental
+		variable. For instance, for the header 'key: ${VALUE}' where VALUE="foo"
+		will be evaluated to 'key: foo'. Note if no corresponding environmental
+		variable is found the header will be unchanged.`))
+	authority = flags.String("authority", "", prettify(`
 		Value of :authority pseudo-header to be use with underlying HTTP/2
 		requests. It defaults to the given address.`))
 	data = flags.String("d", "", prettify(`
@@ -313,6 +319,13 @@ func main() {
 		return cc
 	}
 
+	var headers []string
+	if *expandHeaders {
+		headers = grpcurl.ExpandHeaders(addlHeaders)
+	} else {
+		headers = addlHeaders
+	}
+
 	var cc *grpc.ClientConn
 	var descSource grpcurl.DescriptorSource
 	var refClient *grpcreflect.Client
@@ -329,7 +342,7 @@ func main() {
 			fail(err, "Failed to process proto source files.")
 		}
 	} else {
-		md := grpcurl.MetadataFromHeaders(append(addlHeaders, reflHeaders...))
+		md := grpcurl.MetadataFromHeaders(append(headers, reflHeaders...))
 		refCtx := metadata.NewOutgoingContext(ctx, md)
 		cc = dial()
 		refClient = grpcreflect.NewClient(refCtx, reflectpb.NewServerReflectionClient(cc))
@@ -502,7 +515,7 @@ func main() {
 		}
 		h := grpcurl.NewDefaultEventHandler(os.Stdout, descSource, formatter, *verbose)
 
-		err = grpcurl.InvokeRPC(ctx, descSource, cc, symbol, append(addlHeaders, rpcHeaders...), h, rf.Next)
+		err = grpcurl.InvokeRPC(ctx, descSource, cc, symbol, append(headers, rpcHeaders...), h, rf.Next)
 		if err != nil {
 			fail(err, "Error invoking method %q", symbol)
 		}
diff --git a/grpcurl.go b/grpcurl.go
index 64947de..d3e74eb 100644
--- a/grpcurl.go
+++ b/grpcurl.go
@@ -15,6 +15,8 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net"
+	"os"
+	"regexp"
 	"sort"
 	"strings"
 
@@ -161,6 +163,37 @@ func MetadataFromHeaders(headers []string) metadata.MD {
 	return md
 }
 
+/* Expands environmental variables contained in the header string
+ * If no corresponding environmental variable is found, the header
+ * string is not changed. Hence, if the regex matches accidentally
+ * no changes are made.
+ */
+func ExpandHeaders(headers []string) []string {
+	expandedHeaders := make([]string, len(headers))
+	for idx, header := range headers {
+		if header != "" {
+			regex := regexp.MustCompile(`\${\w+}`)
+			results := regex.FindAllString(header, -1)
+			if results != nil {
+				expandedHeader := header
+				for _, result := range results {
+					envVarValue := os.Getenv(result[2 : len(result)-1])
+					replacementValue := envVarValue
+					// If no corresponding env var is found, leave the header as is.
+					if len(envVarValue) == 0 {
+						replacementValue = result
+					}
+					expandedHeader = strings.Replace(expandedHeader, result, replacementValue, -1)
+				}
+				expandedHeaders[idx] = expandedHeader
+			} else {
+				expandedHeaders[idx] = headers[idx]
+			}
+		}
+	}
+	return expandedHeaders
+}
+
 var base64Codecs = []*base64.Encoding{base64.StdEncoding, base64.URLEncoding, base64.RawStdEncoding, base64.RawURLEncoding}
 
 func decode(val string) (string, error) {
diff --git a/grpcurl_test.go b/grpcurl_test.go
index a0847f3..5c65478 100644
--- a/grpcurl_test.go
+++ b/grpcurl_test.go
@@ -300,6 +300,23 @@ func TestGetAllFiles(t *testing.T) {
 	}
 }
 
+func TestExpandHeaders(t *testing.T) {
+	inHeaders := []string{"key1: ${value}", "key2: bar", "key3: ${woo", "key4: woo}", "key5: ${TEST}",
+		"key6: ${TEST_VAR}", "${TEST}: ${TEST_VAR}"}
+	os.Setenv("value", "value")
+	os.Setenv("TEST", "value5")
+	os.Setenv("TEST_VAR", "value6")
+	expectedHeaders := map[string]bool{"key1: value": true, "key2: bar": true, "key3: ${woo": true, "key4: woo}": true,
+		"key5: value5": true, "key6: value6": true, "value5: value6": true}
+
+	outHeaders := ExpandHeaders(inHeaders)
+	for _, expandedHeader := range outHeaders {
+		if _, ok := expectedHeaders[expandedHeader]; !ok {
+			t.Errorf("The ExpandHeaders function has generated an unexpected header. Recieved unexpected header %s", expandedHeader)
+		}
+	}
+}
+
 func fileNames(files []*desc.FileDescriptor) []string {
 	names := make([]string, len(files))
 	for i, f := range files {