Merge 5b62430 into 398a6d1

uber · Jul 28, 2020 · 2bebfbd · 2bebfbd
2 parents 398a6d1 + 5b62430
commit 2bebfbd
Show file tree

Hide file tree

Showing 2 changed files with 145 additions and 1 deletion.
diff --git a/runtime/server_header.go b/runtime/server_header.go
@@ -33,7 +33,14 @@ import (
 
 // Header defines methods on ServerHeaders
 type Header interface {
+	// Get mirrors the implementation of http.header and returns a single header value.
+	// When a key contains multiple values, -only- the first one is returned.
+	// Ref: https://golang.org/pkg/net/http/#Header.Get
 	Get(key string) (string, bool)
+	// Values mirrors the implementation of http.header and returns a slice of header values.
+	// When a key contains multiple values, the entire collection is returned.
+	// Ref: https://golang.org/pkg/net/http/#Header.Values
+	//Values(key string) ([]string, bool)
 	Add(key string, value string)
 	Set(key string, value string)
 	Keys() []string
@@ -53,7 +60,7 @@ func NewServerHTTPHeader(h http.Header) ServerHTTPHeader {
 
 // Get retrieves the first string stored on a given header. Bool
 // return value is used to distinguish between the presence of a
-// header with golang's zerovalue string and the absence of the string.
+// header with golang's zerovalue string and the absence of the header.
 func (zh ServerHTTPHeader) Get(key string) (string, bool) {
 	httpKey := textproto.CanonicalMIMEHeaderKey(key)
 	h := zh[httpKey]
@@ -63,6 +70,17 @@ func (zh ServerHTTPHeader) Get(key string) (string, bool) {
 	return "", false
 }
 
+// Values retrieves the entire collection of values stored on a given header.
+// Bool return value is used to distinguish between the presence of a
+// header with golang's zerovalue slice and the absence of the header.
+func (zh ServerHTTPHeader) Values(key string) ([]string, bool) {
+	httpKey := textproto.CanonicalMIMEHeaderKey(key)
+	if _, ok := zh[httpKey]; ok {
+		return textproto.MIMEHeader(zh).Values(key), true
+	}
+	return []string{}, false
+}
+
 // GetOrEmptyStr retrieves the first string stored on a given header or
 // the empty string (golang's zero vlaue for string types)
 func (zh ServerHTTPHeader) GetOrEmptyStr(key string) string {
@@ -151,6 +169,19 @@ func (th ServerTChannelHeader) Get(key string) (string, bool) {
 	return value, ok
 }
 
+// Values retrieves the entire collection of values stored on a given header.
+// Bool return value is used to distinguish between the presence of a
+// header with golang's zerovalue slice and the absence of the header.
+func (th ServerTChannelHeader) Values(key string) ([]string, bool) {
+	if value, ok := th.Get(key); ok {
+		// In the case of TChannel, ServerTChannelHeader does not support
+		//  multiple, disparate values so we defer to Get and package it
+		// in a slice to meet the interface's requirement.
+		return []string{value}, ok
+	}
+	return []string{}, false
+}
+
 // Add is an alias to Set.
 func (th ServerTChannelHeader) Add(key string, value string) {
 	th.Set(key, value)

diff --git a/runtime/server_header_test.go b/runtime/server_header_test.go
@@ -73,6 +73,67 @@ func TestGetMultivalueKey(t *testing.T) {
 	assert.Equal(t, "headOne", v)
 }
 
+func TestValues(t *testing.T) {
+	key := "Canonicalized-Key"
+	testCases := []struct {
+		title          string
+		header         func() zanzibar.ServerHTTPHeader
+		key            string
+		expectedValues []string
+		expectedBool   bool
+	}{
+		{
+			title: "Multiple values for a valid key",
+			header: func() zanzibar.ServerHTTPHeader {
+				zh := zanzibar.NewServerHTTPHeader(http.Header{})
+				zh.Set(key, "headerOne")
+				zh.Add(key, "headerTwo")
+				return zh
+			},
+			expectedValues: []string{"headerOne", "headerTwo"},
+			expectedBool:   true,
+		},
+		{
+			title: "Single value for a valid key",
+			header: func() zanzibar.ServerHTTPHeader {
+				zh := zanzibar.NewServerHTTPHeader(http.Header{})
+				zh.Set(key, "headerOne")
+				return zh
+			},
+			expectedValues: []string{"headerOne"},
+			expectedBool:   true,
+		},
+		{
+			title: "Zero values for a valid key",
+			header: func() zanzibar.ServerHTTPHeader {
+				zh := zanzibar.NewServerHTTPHeader(http.Header{
+					key: []string{},
+				})
+				return zh
+			},
+			expectedValues: []string{},
+			expectedBool:   true,
+		},
+		{
+			title: "Missing header key",
+			header: func() zanzibar.ServerHTTPHeader {
+				zh := zanzibar.NewServerHTTPHeader(http.Header{})
+				return zh
+			},
+			expectedValues: []string{},
+			expectedBool:   false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.title, func(t *testing.T) {
+			actualValues, actualBool := tc.header().Values(key)
+			assert.Equal(t, tc.expectedValues, actualValues, tc.title)
+			assert.Equal(t, tc.expectedBool, actualBool, tc.title)
+		})
+	}
+}
+
 func TestAdd(t *testing.T) {
 	zh := zanzibar.NewServerHTTPHeader(http.Header{})
 	zh.Set("bar", "otherHeader")
@@ -164,6 +225,58 @@ func TestSTHGetMissingKey(t *testing.T) {
 	assert.Equal(t, "", v)
 }
 
+func TestSTHValues(t *testing.T) {
+	key := "Canonicalized-Key"
+	testCases := []struct {
+		title          string
+		header         func() zanzibar.ServerTChannelHeader
+		key            string
+		expectedValues []string
+		expectedBool   bool
+	}{
+		{
+			title: "Multiple values set for a valid key",
+			header: func() zanzibar.ServerTChannelHeader {
+				zh := zanzibar.ServerTChannelHeader{}
+				zh.Set(key, "headerOne")
+				// For ServerTChannelHeader, Add is an alias to Set so
+				// this will overwrite the existing key with the new value.
+				zh.Add(key, "headerTwo")
+				return zh
+			},
+			expectedValues: []string{"headerTwo"},
+			expectedBool:   true,
+		},
+		{
+			title: "Single value for a valid key",
+			header: func() zanzibar.ServerTChannelHeader {
+				zh := zanzibar.ServerTChannelHeader{}
+				zh.Set(key, "headerOne")
+				return zh
+			},
+			expectedValues: []string{"headerOne"},
+			expectedBool:   true,
+		},
+		{
+			title: "Missing header key",
+			header: func() zanzibar.ServerTChannelHeader {
+				zh := zanzibar.ServerTChannelHeader{}
+				return zh
+			},
+			expectedValues: []string{},
+			expectedBool:   false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.title, func(t *testing.T) {
+			actualValues, actualBool := tc.header().Values(key)
+			assert.Equal(t, tc.expectedValues, actualValues, tc.title)
+			assert.Equal(t, tc.expectedBool, actualBool, tc.title)
+		})
+	}
+}
+
 func TestSTHAdd(t *testing.T) {
 	zh := zanzibar.ServerTChannelHeader{}