Correctly propagate statusCode

Author: stesie

main.go

@@ -38,6 +38,7 @@ type message struct {
 
 type responseWrapper struct {
 	http.ResponseWriter
+	StatusCode int
 	Body []byte
 }
 
@@ -71,6 +72,10 @@ func (rw *responseWrapper) Write(data []byte) (int, error) {
 	return len(data), nil
 }
 
+func (rw *responseWrapper) WriteHeader(statusCode int) {
+	rw.StatusCode = statusCode;
+}
+
 func setupEnv() error {
 	environment["_HANDLER"], _ = os.LookupEnv("_HANDLER")
 	environment["LAMBDA_TASK_ROOT"], _ = os.LookupEnv("LAMBDA_TASK_ROOT")
@@ -206,12 +211,12 @@ func responseHandler(w http.ResponseWriter, r *http.Request) {
 func mapEvent(h http.Handler) http.Handler {
 	eventType, _ := os.LookupEnv("EVENT")
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		rw := responseWrapper{w, []byte{}}
+		rw := responseWrapper{w, 200, []byte{}}
 		switch eventType {
 		case "API_GATEWAY":
 			apiGateway.Request(r)
 			h.ServeHTTP(&rw, r)
-			apiGateway.Response(w, rw.Body)
+			apiGateway.Response(w, rw.StatusCode, rw.Body)
 		default:
 			h.ServeHTTP(w, r)
 		}

pkg/events/apiGateway/mapper.go

@@ -40,14 +40,19 @@ func Request(r *http.Request) {
 	r.Body = ioutil.NopCloser(bytes.NewBuffer(js))
 }
 
-func Response(w http.ResponseWriter, data []byte) (int, error) {
+func Response(w http.ResponseWriter, statusCode int, data []byte) (int, error) {
 	var js events.APIGatewayProxyResponse
 	if err := json.Unmarshal(data, &js); err != nil {
 		return 0, err
 	}
 	for k, v := range js.Headers {
 		w.Header().Set(k, v)
 	}
-	w.WriteHeader(js.StatusCode)
+
+	if js.StatusCode >= 200 {
+		statusCode = js.StatusCode
+	}
+	w.WriteHeader(statusCode)
+
 	return w.Write([]byte(js.Body))
 }