Add Expand Headers Feature

cmd/grpcurl/grpcurl.go

@@ -52,13 +52,21 @@ 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, headers may use '${NAME}' syntax to reference environment variables.
+		These will be expanded to the actual environment variable value before
+		sending to the server. For example, if there is an environment variable
+		defined like FOO=bar, then a header of 'key: ${FOO}' would expand to 'key: bar'.
+		This applies to -H, -rpc-header, and -reflect-header options. No other
+		expansion/escaping is performed. This can be used to supply
+		credentials/secrets without having to put them in command-line arguments.`))
+	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 +321,22 @@ func main() {
 		return cc
 	}
 
+	if *expandHeaders {
+		var err error
+		addlHeaders, err = grpcurl.ExpandHeaders(addlHeaders)
+		if err != nil {
+			fail(err, "Failed to expand additional headers")
+		}
+		rpcHeaders, err = grpcurl.ExpandHeaders(rpcHeaders)
+		if err != nil {
+			fail(err, "Failed to expand rpc headers")
+		}
+		reflHeaders, err = grpcurl.ExpandHeaders(reflHeaders)
+		if err != nil {
+			fail(err, "Failed to expand reflection headers")
+		}
+	}
+
 	var cc *grpc.ClientConn
 	var descSource grpcurl.DescriptorSource
 	var refClient *grpcreflect.Client

grpcurl.go

@@ -15,6 +15,8 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net"
+	"os"
+	"regexp"
 	"sort"
 	"strings"
 
@@ -161,6 +163,36 @@ func MetadataFromHeaders(headers []string) metadata.MD {
 	return md
 }
 
+var envVarRegex = regexp.MustCompile(`\${\w+}`)
+
+// ExpandHeaders expands environment variables contained in the header string.
+// If no corresponding environment variable is found an error is returned.
+// TODO: Add escaping for `${`
+func ExpandHeaders(headers []string) ([]string, error) {
+	expandedHeaders := make([]string, len(headers))
+	for idx, header := range headers {
+		if header == "" {
+			continue
+		}
+		results := envVarRegex.FindAllString(header, -1)
+		if len(results) == 0 {
+			expandedHeaders[idx] = headers[idx]
+			continue
+		}
+		expandedHeader := header
+		for _, result := range results {
+			envVarName := result[2 : len(result)-1] // strip leading `${` and trailing `}`
+			envVarValue, ok := os.LookupEnv(envVarName)
+			if !ok {
+				return nil, fmt.Errorf("header %q refers to missing environment variable %q", header, envVarName)
+			}
+			expandedHeader = strings.Replace(expandedHeader, result, envVarValue, -1)
+		}
+		expandedHeaders[idx] = expandedHeader
+	}
+	return expandedHeaders, nil
+}
+
 var base64Codecs = []*base64.Encoding{base64.StdEncoding, base64.URLEncoding, base64.RawStdEncoding, base64.RawURLEncoding}
 
 func decode(val string) (string, error) {

grpcurl_test.go

@@ -300,6 +300,33 @@ 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}", "key8: ${EMPTY}"}
+	os.Setenv("value", "value")
+	os.Setenv("TEST", "value5")
+	os.Setenv("TEST_VAR", "value6")
+	os.Setenv("EMPTY", "")
+	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, "key8: ": true}
+
+	outHeaders, err := ExpandHeaders(inHeaders)
+	if err != nil {
+		t.Errorf("The ExpandHeaders function generated an unexpected error %s", err)
+	}
+	for _, expandedHeader := range outHeaders {
+		if _, ok := expectedHeaders[expandedHeader]; !ok {
+			t.Errorf("The ExpandHeaders function has returned an unexpected header. Received unexpected header %s", expandedHeader)
+		}
+	}
+
+	badHeaders := []string{"key: ${DNE}"}
+	_, err = ExpandHeaders(badHeaders)
+	if err == nil {
+		t.Errorf("The ExpandHeaders function should return an error for missing environment variables %q", badHeaders)
+	}
+}
+
 func fileNames(files []*desc.FileDescriptor) []string {
 	names := make([]string, len(files))
 	for i, f := range files {