Make Expand Headers Generate Expected Errors

Expand Headers should throw an error for unset environmental
variables. Also address additional PR feedback.

cmd/grpcurl/grpcurl.go

@@ -59,11 +59,12 @@ var (
 	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
+		If set any environmental variables contained contained in any additional, rpc,
+		or reflection header string will be substituted for by the value of the 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.`))
+		will be expanded to 'key: foo'. Any undefined environmental variables
+		will throw an error. Note that verbose mode shows that the headers are
+		being expanded to. Escaping '${' is not supported.`))
 	authority = flags.String("authority", "", prettify(`
 		Value of :authority pseudo-header to be use with underlying HTTP/2
 		requests. It defaults to the given address.`))
@@ -319,11 +320,20 @@ func main() {
 		return cc
 	}
 
-	var headers []string
 	if *expandHeaders {
-		headers = grpcurl.ExpandHeaders(addlHeaders)
-	} else {
-		headers = addlHeaders
+		var err error
+		addlHeaders, err = grpcurl.ExpandHeaders(addlHeaders)
+		if err != nil {
+			fail(err, "Failed to expand additional headers, missing environmental variable")
+		}
+		rpcHeaders, err = grpcurl.ExpandHeaders(rpcHeaders)
+		if err != nil {
+			fail(err, "Failed to expand rpc headers, missing environmental variable")
+		}
+		reflHeaders, err = grpcurl.ExpandHeaders(reflHeaders)
+		if err != nil {
+			fail(err, "Failed to expand reflection headers, missing environmental variable")
+		}
 	}
 
 	var cc *grpc.ClientConn
@@ -342,7 +352,7 @@ func main() {
 			fail(err, "Failed to process proto source files.")
 		}
 	} else {
-		md := grpcurl.MetadataFromHeaders(append(headers, reflHeaders...))
+		md := grpcurl.MetadataFromHeaders(append(addlHeaders, reflHeaders...))
 		refCtx := metadata.NewOutgoingContext(ctx, md)
 		cc = dial()
 		refClient = grpcreflect.NewClient(refCtx, reflectpb.NewServerReflectionClient(cc))
@@ -515,7 +525,7 @@ func main() {
 		}
 		h := grpcurl.NewDefaultEventHandler(os.Stdout, descSource, formatter, *verbose)
 
-		err = grpcurl.InvokeRPC(ctx, descSource, cc, symbol, append(headers, rpcHeaders...), h, rf.Next)
+		err = grpcurl.InvokeRPC(ctx, descSource, cc, symbol, append(addlHeaders, rpcHeaders...), h, rf.Next)
 		if err != nil {
 			fail(err, "Error invoking method %q", symbol)
 		}

grpcurl.go

@@ -163,25 +163,24 @@ 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 {
+var envVarRegex = regexp.MustCompile(`\${\w+}`)
+
+// ExpandHeaders expands environmental variables contained in the header string.
+// If no corresponding environmental variable is found an error is thrown.
+// TODO: Add escaping for `${`
+func ExpandHeaders(headers []string) ([]string, error) {
 	expandedHeaders := make([]string, len(headers))
 	for idx, header := range headers {
 		if header != "" {
-			regex := regexp.MustCompile(`\${\w+}`)
-			results := regex.FindAllString(header, -1)
+			results := envVarRegex.FindAllString(header, -1)
 			if results != nil {
 				expandedHeader := header
 				for _, result := range results {
-					envVarValue := os.Getenv(result[2 : len(result)-1])
+					envVarName := result[2 : len(result)-1]
+					envVarValue := os.Getenv(envVarName)
 					replacementValue := envVarValue
-					// If no corresponding env var is found, leave the header as is.
 					if len(envVarValue) == 0 {
-						replacementValue = result
+						return nil, fmt.Errorf("environmental variable '%s' is not set", envVarName)
 					}
 					expandedHeader = strings.Replace(expandedHeader, result, replacementValue, -1)
 				}
@@ -191,7 +190,7 @@ func ExpandHeaders(headers []string) []string {
 			}
 		}
 	}
-	return expandedHeaders
+	return expandedHeaders, nil
 }
 
 var base64Codecs = []*base64.Encoding{base64.StdEncoding, base64.URLEncoding, base64.RawStdEncoding, base64.RawURLEncoding}

grpcurl_test.go

@@ -301,20 +301,29 @@ func TestGetAllFiles(t *testing.T) {
 }
 
 func TestExpandHeaders(t *testing.T) {
-	inHeaders := []string{"key1: ${value}", "key2: bar", "key3: ${woo", "key4: woo}", "key5: ${TEST}",
+	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,
+	expectedHeaders := map[string]bool{"": true, "key1: value": true, "key2: bar": true, "key3: ${woo": true, "key4: woo}": true,
 		"key5: value5": true, "key6: value6": true, "value5: value6": true}
 
-	outHeaders := ExpandHeaders(inHeaders)
+	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 generated an unexpected header. Recieved unexpected header %s", expandedHeader)
 		}
 	}
+
+	badHeaders := []string{"key: ${DNE}"}
+	_, err = ExpandHeaders(badHeaders)
+	if err == nil {
+		t.Errorf("The ExpandHeaders function should generate an error for missing environmental variables %q", badHeaders)
+	}
 }
 
 func fileNames(files []*desc.FileDescriptor) []string {