Merge pull request aws#339 from SAPessi/develop

PaulMaddox · web-flow · commit 290b208b73c4 · 2018-04-05T13:33:06.000+01:00
Added Forward headers similar to API Gateway's
diff --git a/router/event.go b/router/event.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"io/ioutil"
 	"net/http"
+	"strings"
 
 	"github.com/gorilla/mux"
 )
@@ -69,6 +70,17 @@ func NewEvent(req *http.Request, isBase64Encoded bool) (*Event, error) {
 		}
 	}
 
+	// add the forwarded headers we expect to see from an API Gateway request
+	if _, ok := headers["Host"]; !ok {
+		host := req.Host
+		if strings.Contains(host, ":") {
+			host = host[:strings.Index(host, ":")]
+		}
+		headers["Host"] = host
+	}
+	headers["X-Forwarded-Proto"] = req.URL.Scheme
+	headers["X-Forwarded-Port"] = req.URL.Port()
+
 	query := map[string]string{}
 	for name, values := range req.URL.Query() {
 		for _, value := range values {
diff --git a/router/event_test.go b/router/event_test.go
@@ -81,6 +81,40 @@ var _ = Describe("Event", func() {
 					r.Router().ServeHTTP(rec, req)
 				})
 			})
+
+			Context("Includes forwarded headers", func() {
+				req, _ := http.NewRequest("GET", "http://localhost:3000/get", new(bytes.Buffer))
+
+				It("Includes Host header", func() {
+					r.AddFunction(function, func(w http.ResponseWriter, e *Event) {
+						hostHeader, ok := e.Headers["Host"]
+						Expect(ok).To(BeTrue())
+						Expect(hostHeader).To(BeIdenticalTo("localhost"))
+					})
+					rec := httptest.NewRecorder()
+					r.Router().ServeHTTP(rec, req)
+				})
+
+				It("Includes X-Forwarded-Proto header", func() {
+					r.AddFunction(function, func(w http.ResponseWriter, e *Event) {
+						hostHeader, ok := e.Headers["X-Forwarded-Proto"]
+						Expect(ok).To(BeTrue())
+						Expect(hostHeader).To(BeIdenticalTo("http"))
+					})
+					rec := httptest.NewRecorder()
+					r.Router().ServeHTTP(rec, req)
+				})
+
+				It("Includes X-Forwarded-Port header", func() {
+					r.AddFunction(function, func(w http.ResponseWriter, e *Event) {
+						hostHeader, ok := e.Headers["X-Forwarded-Port"]
+						Expect(ok).To(BeTrue())
+						Expect(hostHeader).To(BeIdenticalTo("3000"))
+					})
+					rec := httptest.NewRecorder()
+					r.Router().ServeHTTP(rec, req)
+				})
+			})
 		})
 
 		Context("with no parameters on the route", func() {
