From 45b5ae08a961a7fb979a37ccafb0ddd2cbc745b6 Mon Sep 17 00:00:00 2001 From: jammerful Date: Wed, 25 Sep 2019 20:52:21 -0400 Subject: [PATCH] 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 | 30 ++++++++++++++++++++---------- grpcurl.go | 23 +++++++++++------------ grpcurl_test.go | 15 ++++++++++++--- 3 files changed, 43 insertions(+), 25 deletions(-) diff --git a/cmd/grpcurl/grpcurl.go b/cmd/grpcurl/grpcurl.go index f40b1b5..20ae39a 100644 --- a/cmd/grpcurl/grpcurl.go +++ b/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) } diff --git a/grpcurl.go b/grpcurl.go index d3e74eb..d86d108 100644 --- a/grpcurl.go +++ b/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} diff --git a/grpcurl_test.go b/grpcurl_test.go index 5c65478..a0babf7 100644 --- a/grpcurl_test.go +++ b/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 {