NewAPIGatewayProxyResponse handles Headers with multiple values, and …

…ensuring canonicalized header keys
iamatypeofwalrus · Jan 28, 2018 · 7c052c7 · 7c052c7
1 parent 7055c9e
commit 7c052c7
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 6 deletions.
diff --git a/api_gateway_proxy_response.go b/api_gateway_proxy_response.go
@@ -2,36 +2,46 @@ package shim
 
 import (
 	"net/http"
+	"strings"
 
 	"github.com/aws/aws-lambda-go/events"
 )
 
+const (
+	multipleValueSeperator = ","
+)
+
 // NewAPIGatewayProxyResponse converts a shim.ResponseWriter into an events.APIGatewayProxyResponse
 func NewAPIGatewayProxyResponse(rw *ResponseWriter) events.APIGatewayProxyResponse {
 	headers := formatHeaders(rw.Headers)
-
 	setDefaultContentType(headers, rw.Body.Bytes())
 
+	// TODO: if body type is not mime type convert body to base64
+
 	return events.APIGatewayProxyResponse{
 		StatusCode: rw.Code,
 		Body:       rw.Body.String(),
 		Headers:    headers,
 	}
 }
 
+// formatHeaders converts an http.Headers map into a map[string]string which is expected by Lambda.
 func formatHeaders(h http.Header) map[string]string {
 	headers := make(map[string]string)
 
 	for k, v := range h {
-		// No great options here. Rather than cat-ing the string array together
-		// and on average getting it right, I'd rather just take the first item
-		// in the array.
+		// e.g. convert accept-encoding into Accept-Encoding
+		canonicalKey := http.CanonicalHeaderKey(k)
 		var str string
-		if len(v) > 0 {
+		if len(v) == 1 {
 			str = v[0]
+		} else if len(v) > 1 {
+			// Per RFC 2616 combine headers with multiple values with ","
+			// Source: https://www.w3.org/Protocols/rfc2616/rfc2616-sec4.html#sec4.2
+			str = strings.Join(v, multipleValueSeperator)
 		}
 
-		headers[k] = str
+		headers[canonicalKey] = str
 	}
 
 	return headers

diff --git a/api_gateway_proxy_response_test.go b/api_gateway_proxy_response_test.go
@@ -0,0 +1,52 @@
+package shim
+
+import (
+	"net/http"
+	"sort"
+	"strings"
+	"testing"
+)
+
+func TestFormatHeadersCanonicalizesKeyNames(t *testing.T) {
+	h := make(http.Header)
+	encoding := "gzip"
+	h["accept-encoding"] = []string{"gzip"}
+	f := formatHeaders(h)
+
+	v, ok := f["Accept-Encoding"]
+	if !ok {
+		t.Error("expected to get key Accept-Encoding to have a value")
+	}
+
+	if v != encoding {
+		t.Errorf("expected encoding to be %v but was %v", encoding, v)
+	}
+}
+
+func TestFormatHeadersHandlesMultipleValusForAKey(t *testing.T) {
+	h := make(http.Header)
+	gzip := "gzip"
+	deflate := "deflate"
+	encodings := []string{gzip, deflate}
+	h["Accept-Encoding"] = encodings
+	f := formatHeaders(h)
+
+	v, ok := f["Accept-Encoding"]
+	if !ok {
+		t.Fatal("expected value for Accept-Encoding to be present")
+	}
+
+	encodingVals := strings.Split(v, multipleValueSeperator)
+	if len(encodingVals) != len(encodings) {
+		t.Fatalf("expected number of encoded values to be %v but was %v\n", len(encodings), len(encodingVals))
+	}
+
+	sort.Strings(encodings)
+	sort.Strings(encodingVals)
+
+	for i, encoding := range encodings {
+		if encodingVals[i] != encoding {
+			t.Errorf("expected encoding at %v to be %v but was %v\n", i, encoding, encodingVals[i])
+		}
+	}
+}