Detecting content type automatically to address issue awslabs#16

sapessi · sapessi · commit 3a4d35011e4a · 2018-07-13T08:33:04.000-07:00
diff --git a/core/response.go b/core/response.go
@@ -13,6 +13,7 @@ import (
 )
 
 const defaultStatusCode = -1
+const contentTypeHeaderKey = "Content-Type"
 
 // ProxyResponseWriter implements http.ResponseWriter and adds the method
 // necessary to return an events.APIGatewayProxyResponse object
@@ -46,6 +47,14 @@ func (r *ProxyResponseWriter) Write(body []byte) (int, error) {
 		r.status = http.StatusOK
 	}
 
+	// if the content type header is not set when we write the body we try to
+	// detect one and set it by default. If the content type cannot be detected
+	// it is automatically set to "application/octet-stream" by the
+	// DetectContentType method
+	if r.Header().Get(contentTypeHeaderKey) == "" {
+		r.Header().Add(contentTypeHeaderKey, http.DetectContentType(body))
+	}
+
 	return (&r.body).Write(body)
 }
 
diff --git a/core/response_test.go b/core/response_test.go
@@ -4,6 +4,7 @@ import (
 	"encoding/base64"
 	"math/rand"
 	"net/http"
+	"strings"
 
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
@@ -49,6 +50,50 @@ var _ = Describe("ResponseWriter tests", func() {
 		})
 	})
 
+	Context("Automatically set response content type", func() {
+		xmlBodyContent := "<?xml version=\"1.0\" encoding=\"UTF-8\"?><note><to>Tove</to><from>Jani</from><heading>Reminder</heading><body>Don't forget me this weekend!</body></note>"
+		htmlBodyContent := " <!DOCTYPE html><html><head><meta charset=\"UTF-8\"><title>Title of the document</title></head><body>Content of the document......</body></html>"
+		It("Does not set the content type if it's already set", func() {
+			resp := NewProxyResponseWriter()
+			resp.Header().Add("Content-Type", "application/json")
+
+			resp.Write([]byte(xmlBodyContent))
+
+			Expect("application/json").To(Equal(resp.Header().Get("Content-Type")))
+			proxyResp, err := resp.GetProxyResponse()
+			Expect(err).To(BeNil())
+			Expect(1).To(Equal(len(proxyResp.Headers)))
+			Expect("application/json").To(Equal(proxyResp.Headers["Content-Type"]))
+			Expect(xmlBodyContent).To(Equal(proxyResp.Body))
+		})
+
+		It("Sets the conte type to text/xml given the body", func() {
+			resp := NewProxyResponseWriter()
+			resp.Write([]byte(xmlBodyContent))
+
+			Expect("").ToNot(Equal(resp.Header().Get("Content-Type")))
+			Expect(true).To(Equal(strings.HasPrefix(resp.Header().Get("Content-Type"), "text/xml;")))
+			proxyResp, err := resp.GetProxyResponse()
+			Expect(err).To(BeNil())
+			Expect(1).To(Equal(len(proxyResp.Headers)))
+			Expect(true).To(Equal(strings.HasPrefix(proxyResp.Headers["Content-Type"], "text/xml;")))
+			Expect(xmlBodyContent).To(Equal(proxyResp.Body))
+		})
+
+		It("Sets the conte type to text/html given the body", func() {
+			resp := NewProxyResponseWriter()
+			resp.Write([]byte(htmlBodyContent))
+
+			Expect("").ToNot(Equal(resp.Header().Get("Content-Type")))
+			Expect(true).To(Equal(strings.HasPrefix(resp.Header().Get("Content-Type"), "text/html;")))
+			proxyResp, err := resp.GetProxyResponse()
+			Expect(err).To(BeNil())
+			Expect(1).To(Equal(len(proxyResp.Headers)))
+			Expect(true).To(Equal(strings.HasPrefix(proxyResp.Headers["Content-Type"], "text/html;")))
+			Expect(htmlBodyContent).To(Equal(proxyResp.Body))
+		})
+	})
+
 	Context("Export API Gateway proxy response", func() {
 		emtpyResponse := NewProxyResponseWriter()
 		emtpyResponse.Header().Add("Content-Type", "application/json")
@@ -70,7 +115,7 @@ var _ = Describe("ResponseWriter tests", func() {
 			Expect("hello").To(Equal(proxyResponse.Body))
 			Expect(http.StatusOK).To(Equal(proxyResponse.StatusCode))
 			Expect(1).To(Equal(len(proxyResponse.Headers)))
-			Expect("text/plain").To(Equal(proxyResponse.Headers["Content-Type"]))
+			Expect(true).To(Equal(strings.HasPrefix(proxyResponse.Headers["Content-Type"], "text/plain")))
 			Expect(proxyResponse.IsBase64Encoded).To(BeFalse())
 		})
 
diff --git a/core/types.go b/core/types.go
@@ -7,11 +7,12 @@ import (
 	"github.com/aws/aws-lambda-go/events"
 )
 
-// Returns a dafault Gateway Timeout (504) response
+// GatewayTimeout returns a dafault Gateway Timeout (504) response
 func GatewayTimeout() events.APIGatewayProxyResponse {
 	return events.APIGatewayProxyResponse{StatusCode: http.StatusGatewayTimeout}
 }
 
+// New Logged Error
 func NewLoggedError(format string, a ...interface{}) error {
 	err := fmt.Errorf(format, a...)
 	fmt.Println(err.Error())

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,7 @@ import (`
`13`	`13`	`)`
`14`	`14`
`15`	`15`	`const defaultStatusCode = -1`
	`16`	`+const contentTypeHeaderKey = "Content-Type"`
`16`	`17`
`17`	`18`	`// ProxyResponseWriter implements http.ResponseWriter and adds the method`
`18`	`19`	`// necessary to return an events.APIGatewayProxyResponse object`
`@@ -46,6 +47,14 @@ func (r *ProxyResponseWriter) Write(body []byte) (int, error) {`
`46`	`47`	`r.status = http.StatusOK`
`47`	`48`	`}`
`48`	`49`
	`50`	`+ // if the content type header is not set when we write the body we try to`
	`51`	`+ // detect one and set it by default. If the content type cannot be detected`
	`52`	`+ // it is automatically set to "application/octet-stream" by the`
	`53`	`+ // DetectContentType method`
	`54`	`+ if r.Header().Get(contentTypeHeaderKey) == "" {`
	`55`	`+ r.Header().Add(contentTypeHeaderKey, http.DetectContentType(body))`
	`56`	`+ }`
	`57`	`+`
`49`	`58`	`return (&r.body).Write(body)`
`50`	`59`	`}`
`51`	`60`
Original file line number	Diff line number	Diff line change
`@@ -7,11 +7,12 @@ import (`
`7`	`7`	`"github.com/aws/aws-lambda-go/events"`
`8`	`8`	`)`
`9`	`9`
`10`		`-// Returns a dafault Gateway Timeout (504) response`
	`10`	`+// GatewayTimeout returns a dafault Gateway Timeout (504) response`
`11`	`11`	`func GatewayTimeout() events.APIGatewayProxyResponse {`
`12`	`12`	`return events.APIGatewayProxyResponse{StatusCode: http.StatusGatewayTimeout}`
`13`	`13`	`}`
`14`	`14`
	`15`	`+// New Logged Error`
`15`	`16`	`func NewLoggedError(format string, a ...interface{}) error {`
`16`	`17`	`err := fmt.Errorf(format, a...)`
`17`	`18`	`fmt.Println(err.Error())`