diff --git a/metadata/metadata.go b/metadata/metadata.go index 8e0f6abe89d..98d62e0675f 100644 --- a/metadata/metadata.go +++ b/metadata/metadata.go @@ -50,7 +50,7 @@ type MD map[string][]string // Keys beginning with "grpc-" are reserved for grpc-internal use only and may // result in errors if set in metadata. func New(m map[string]string) MD { - md := MD{} + md := make(MD, len(m)) for k, val := range m { key := strings.ToLower(k) md[key] = append(md[key], val) @@ -74,7 +74,7 @@ func Pairs(kv ...string) MD { if len(kv)%2 == 1 { panic(fmt.Sprintf("metadata: Pairs got the odd number of input pairs for metadata: %d", len(kv))) } - md := MD{} + md := make(MD, len(kv)/2) for i := 0; i < len(kv); i += 2 { key := strings.ToLower(kv[i]) md[key] = append(md[key], kv[i+1]) @@ -182,19 +182,51 @@ func FromIncomingContext(ctx context.Context) (MD, bool) { if !ok { return nil, false } - out := MD{} + out := make(MD, len(md)) for k, v := range md { // We need to manually convert all keys to lower case, because MD is a // map, and there's no guarantee that the MD attached to the context is // created using our helper functions. key := strings.ToLower(k) - s := make([]string, len(v)) - copy(s, v) - out[key] = s + out[key] = copyOf(v) } return out, true } +// ValueFromIncomingContext returns the metadata value corresponding to the metadata +// key from the incoming metadata if it exists. Key must be lower-case. +// +// Experimental +// +// Notice: This API is EXPERIMENTAL and may be changed or removed in a +// later release. +func ValueFromIncomingContext(ctx context.Context, key string) []string { + md, ok := ctx.Value(mdIncomingKey{}).(MD) + if !ok { + return nil + } + + if v, ok := md[key]; ok { + return copyOf(v) + } + for k, v := range md { + // We need to manually convert all keys to lower case, because MD is a + // map, and there's no guarantee that the MD attached to the context is + // created using our helper functions. + if strings.ToLower(k) == key { + return copyOf(v) + } + } + return nil +} + +// the returned slice must not be modified in place +func copyOf(v []string) []string { + vals := make([]string, len(v)) + copy(vals, v) + return vals +} + // FromOutgoingContextRaw returns the un-merged, intermediary contents of rawMD. // // Remember to perform strings.ToLower on the keys, for both the returned MD (MD @@ -222,15 +254,18 @@ func FromOutgoingContext(ctx context.Context) (MD, bool) { return nil, false } - out := MD{} + mdSize := len(raw.md) + for i := range raw.added { + mdSize += len(raw.added[i]) / 2 + } + + out := make(MD, mdSize) for k, v := range raw.md { // We need to manually convert all keys to lower case, because MD is a // map, and there's no guarantee that the MD attached to the context is // created using our helper functions. key := strings.ToLower(k) - s := make([]string, len(v)) - copy(s, v) - out[key] = s + out[key] = copyOf(v) } for _, added := range raw.added { if len(added)%2 == 1 { diff --git a/metadata/metadata_test.go b/metadata/metadata_test.go index 89be06eaada..57763cd6a97 100644 --- a/metadata/metadata_test.go +++ b/metadata/metadata_test.go @@ -198,6 +198,43 @@ func (s) TestDelete(t *testing.T) { } } +func (s) TestValueFromIncomingContext(t *testing.T) { + md := Pairs( + "X-My-Header-1", "42", + "X-My-Header-2", "43-1", + "X-My-Header-2", "43-2", + "x-my-header-3", "44", + ) + ctx := NewIncomingContext(context.Background(), md) + + for _, test := range []struct { + key string + want []string + }{ + { + key: "x-my-header-1", + want: []string{"42"}, + }, + { + key: "x-my-header-2", + want: []string{"43-1", "43-2"}, + }, + { + key: "x-my-header-3", + want: []string{"44"}, + }, + { + key: "x-unknown", + want: nil, + }, + } { + v := ValueFromIncomingContext(ctx, test.key) + if !reflect.DeepEqual(v, test.want) { + t.Errorf("value of metadata is %v, want %v", v, test.want) + } + } +} + func (s) TestAppendToOutgoingContext(t *testing.T) { // Pre-existing metadata tCtx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)