Add Expand Headers Feature

See discussion of the feature here:
https://github.com/fullstorydev/grpcurl/issues/116
This commit is contained in:
jammerful 2019-09-23 22:59:04 -04:00
parent 4054d1d115
commit 21ea8b8e04
3 changed files with 72 additions and 9 deletions

View File

@ -52,13 +52,19 @@ var (
key = flags.String("key", "", prettify(` key = flags.String("key", "", prettify(`
File containing client private key, to present to the server. Not valid File containing client private key, to present to the server. Not valid
with -plaintext option. Must also provide -cert option.`)) with -plaintext option. Must also provide -cert option.`))
protoset multiString protoset multiString
protoFiles multiString protoFiles multiString
importPaths multiString importPaths multiString
addlHeaders multiString addlHeaders multiString
rpcHeaders multiString rpcHeaders multiString
reflHeaders multiString reflHeaders multiString
authority = flags.String("authority", "", prettify(` 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 Value of :authority pseudo-header to be use with underlying HTTP/2
requests. It defaults to the given address.`)) requests. It defaults to the given address.`))
data = flags.String("d", "", prettify(` data = flags.String("d", "", prettify(`
@ -313,6 +319,13 @@ func main() {
return cc return cc
} }
var headers []string
if *expandHeaders {
headers = grpcurl.ExpandHeaders(addlHeaders)
} else {
headers = addlHeaders
}
var cc *grpc.ClientConn var cc *grpc.ClientConn
var descSource grpcurl.DescriptorSource var descSource grpcurl.DescriptorSource
var refClient *grpcreflect.Client var refClient *grpcreflect.Client
@ -329,7 +342,7 @@ func main() {
fail(err, "Failed to process proto source files.") fail(err, "Failed to process proto source files.")
} }
} else { } else {
md := grpcurl.MetadataFromHeaders(append(addlHeaders, reflHeaders...)) md := grpcurl.MetadataFromHeaders(append(headers, reflHeaders...))
refCtx := metadata.NewOutgoingContext(ctx, md) refCtx := metadata.NewOutgoingContext(ctx, md)
cc = dial() cc = dial()
refClient = grpcreflect.NewClient(refCtx, reflectpb.NewServerReflectionClient(cc)) refClient = grpcreflect.NewClient(refCtx, reflectpb.NewServerReflectionClient(cc))
@ -502,7 +515,7 @@ func main() {
} }
h := grpcurl.NewDefaultEventHandler(os.Stdout, descSource, formatter, *verbose) 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 { if err != nil {
fail(err, "Error invoking method %q", symbol) fail(err, "Error invoking method %q", symbol)
} }

View File

@ -15,6 +15,8 @@ import (
"fmt" "fmt"
"io/ioutil" "io/ioutil"
"net" "net"
"os"
"regexp"
"sort" "sort"
"strings" "strings"
@ -161,6 +163,37 @@ func MetadataFromHeaders(headers []string) metadata.MD {
return 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} var base64Codecs = []*base64.Encoding{base64.StdEncoding, base64.URLEncoding, base64.RawStdEncoding, base64.RawURLEncoding}
func decode(val string) (string, error) { func decode(val string) (string, error) {

View File

@ -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 { func fileNames(files []*desc.FileDescriptor) []string {
names := make([]string, len(files)) names := make([]string, len(files))
for i, f := range files { for i, f := range files {