Detecting content type automatically to address issue #16

awslabs · Jul 13, 2018 · 3a4d350 · 3a4d350
1 parent 6986cdf
commit 3a4d350
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 2 deletions.
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())