From 0057c9d7af4bcb829dc107a28471ad6b2a6db3ba Mon Sep 17 00:00:00 2001
From: Janne Snabb <snabb@epipe.com>
Date: Thu, 7 Feb 2019 20:58:16 +0200
Subject: [PATCH] Pass context.Context from Lambda runtime to http.Request.

Fixes https://github.com/awslabs/aws-lambda-go-api-proxy/issues/27
---
 chi/adapter.go              | 5 +++--
 chi/chilambda_test.go       | 3 ++-
 gin/adapter.go              | 5 +++--
 gin/ginlambda_test.go       | 3 ++-
 gorillamux/adapter.go       | 5 +++--
 gorillamux/adapter_test.go  | 5 +++--
 handlerfunc/adapter.go      | 5 +++--
 handlerfunc/adapter_test.go | 3 ++-
 httpadapter/adapter.go      | 5 +++--
 httpadapter/adapter_test.go | 3 ++-
 negroni/adapter.go          | 5 +++--
 negroni/adapter_test.go     | 5 +++--
 sample/main.go              | 5 +++--
 13 files changed, 35 insertions(+), 22 deletions(-)

diff --git a/chi/adapter.go b/chi/adapter.go
index e886960..44afbff 100644
--- a/chi/adapter.go
+++ b/chi/adapter.go
@@ -4,6 +4,7 @@
 package chiadapter
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -30,7 +31,7 @@ func New(chi *chi.Mux) *ChiLambda {
 // Proxy receives an API Gateway proxy event, transforms it into an http.Request
 // object, and sends it to the chi.Mux for routing.
 // It returns a proxy response object gneerated from the http.ResponseWriter.
-func (g *ChiLambda) Proxy(req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func (g *ChiLambda) Proxy(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	chiRequest, err := g.ProxyEventToHTTPRequest(req)
 
 	if err != nil {
@@ -38,7 +39,7 @@ func (g *ChiLambda) Proxy(req events.APIGatewayProxyRequest) (events.APIGatewayP
 	}
 
 	respWriter := core.NewProxyResponseWriter()
-	g.chiMux.ServeHTTP(http.ResponseWriter(respWriter), chiRequest)
+	g.chiMux.ServeHTTP(http.ResponseWriter(respWriter), chiRequest.WithContext(ctx))
 
 	proxyResponse, err := respWriter.GetProxyResponse()
 	if err != nil {
diff --git a/chi/chilambda_test.go b/chi/chilambda_test.go
index 37b6105..7f7080a 100644
--- a/chi/chilambda_test.go
+++ b/chi/chilambda_test.go
@@ -1,6 +1,7 @@
 package chiadapter_test
 
 import (
+	"context"
 	"log"
 	"net/http"
 
@@ -29,7 +30,7 @@ var _ = Describe("ChiLambda tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			resp, err := adapter.Proxy(req)
+			resp, err := adapter.Proxy(context.Background(), req)
 
 			Expect(err).To(BeNil())
 			Expect(resp.StatusCode).To(Equal(200))
diff --git a/gin/adapter.go b/gin/adapter.go
index 9372449..7ee2d20 100644
--- a/gin/adapter.go
+++ b/gin/adapter.go
@@ -4,6 +4,7 @@
 package ginadapter
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -30,7 +31,7 @@ func New(gin *gin.Engine) *GinLambda {
 // Proxy receives an API Gateway proxy event, transforms it into an http.Request
 // object, and sends it to the gin.Engine for routing.
 // It returns a proxy response object gneerated from the http.ResponseWriter.
-func (g *GinLambda) Proxy(req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func (g *GinLambda) Proxy(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	ginRequest, err := g.ProxyEventToHTTPRequest(req)
 
 	if err != nil {
@@ -38,7 +39,7 @@ func (g *GinLambda) Proxy(req events.APIGatewayProxyRequest) (events.APIGatewayP
 	}
 
 	respWriter := core.NewProxyResponseWriter()
-	g.ginEngine.ServeHTTP(http.ResponseWriter(respWriter), ginRequest)
+	g.ginEngine.ServeHTTP(http.ResponseWriter(respWriter), ginRequest.WithContext(ctx))
 
 	proxyResponse, err := respWriter.GetProxyResponse()
 	if err != nil {
diff --git a/gin/ginlambda_test.go b/gin/ginlambda_test.go
index 34f7064..4091be9 100644
--- a/gin/ginlambda_test.go
+++ b/gin/ginlambda_test.go
@@ -1,6 +1,7 @@
 package ginadapter_test
 
 import (
+	"context"
 	"log"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -30,7 +31,7 @@ var _ = Describe("GinLambda tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			resp, err := adapter.Proxy(req)
+			resp, err := adapter.Proxy(context.Background(), req)
 
 			Expect(err).To(BeNil())
 			Expect(resp.StatusCode).To(Equal(200))
diff --git a/gorillamux/adapter.go b/gorillamux/adapter.go
index d5623c6..ebf2cfe 100644
--- a/gorillamux/adapter.go
+++ b/gorillamux/adapter.go
@@ -1,6 +1,7 @@
 package gorillamux
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -19,14 +20,14 @@ func New(router *mux.Router) *GorillaMuxAdapter {
 	}
 }
 
-func (h *GorillaMuxAdapter) Proxy(event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func (h *GorillaMuxAdapter) Proxy(ctx context.Context, event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	req, err := h.ProxyEventToHTTPRequest(event)
 	if err != nil {
 		return core.GatewayTimeout(), core.NewLoggedError("Could not convert proxy event to request: %v", err)
 	}
 
 	w := core.NewProxyResponseWriter()
-	h.router.ServeHTTP(http.ResponseWriter(w), req)
+	h.router.ServeHTTP(http.ResponseWriter(w), req.WithContext(ctx))
 
 	resp, err := w.GetProxyResponse()
 	if err != nil {
diff --git a/gorillamux/adapter_test.go b/gorillamux/adapter_test.go
index 73a5441..54d9979 100644
--- a/gorillamux/adapter_test.go
+++ b/gorillamux/adapter_test.go
@@ -1,6 +1,7 @@
 package gorillamux_test
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 
@@ -36,7 +37,7 @@ var _ = Describe("GorillaMuxAdapter tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			homePageResp, homePageReqErr := adapter.Proxy(homePageReq)
+			homePageResp, homePageReqErr := adapter.Proxy(context.Background(), homePageReq)
 
 			Expect(homePageReqErr).To(BeNil())
 			Expect(homePageResp.StatusCode).To(Equal(200))
@@ -47,7 +48,7 @@ var _ = Describe("GorillaMuxAdapter tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			productsPageResp, productsPageReqErr := adapter.Proxy(productsPageReq)
+			productsPageResp, productsPageReqErr := adapter.Proxy(context.Background(), productsPageReq)
 
 			Expect(productsPageReqErr).To(BeNil())
 			Expect(productsPageResp.StatusCode).To(Equal(200))
diff --git a/handlerfunc/adapter.go b/handlerfunc/adapter.go
index f39f598..1770847 100644
--- a/handlerfunc/adapter.go
+++ b/handlerfunc/adapter.go
@@ -1,6 +1,7 @@
 package handlerfunc
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -18,14 +19,14 @@ func New(handlerFunc http.HandlerFunc) *HandlerFuncAdapter {
 	}
 }
 
-func (h *HandlerFuncAdapter) Proxy(event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func (h *HandlerFuncAdapter) Proxy(ctx context.Context, event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	req, err := h.ProxyEventToHTTPRequest(event)
 	if err != nil {
 		return core.GatewayTimeout(), core.NewLoggedError("Could not convert proxy event to request: %v", err)
 	}
 
 	w := core.NewProxyResponseWriter()
-	h.handlerFunc.ServeHTTP(http.ResponseWriter(w), req)
+	h.handlerFunc.ServeHTTP(http.ResponseWriter(w), req.WithContext(ctx))
 
 	resp, err := w.GetProxyResponse()
 	if err != nil {
diff --git a/handlerfunc/adapter_test.go b/handlerfunc/adapter_test.go
index 29850d5..cadd53b 100644
--- a/handlerfunc/adapter_test.go
+++ b/handlerfunc/adapter_test.go
@@ -1,6 +1,7 @@
 package handlerfunc_test
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
@@ -29,7 +30,7 @@ var _ = Describe("HandlerFuncAdapter tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			resp, err := adapter.Proxy(req)
+			resp, err := adapter.Proxy(context.Background(), req)
 
 			Expect(err).To(BeNil())
 			Expect(resp.StatusCode).To(Equal(200))
diff --git a/httpadapter/adapter.go b/httpadapter/adapter.go
index 0f3032c..2c28912 100644
--- a/httpadapter/adapter.go
+++ b/httpadapter/adapter.go
@@ -1,6 +1,7 @@
 package httpadapter
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -18,14 +19,14 @@ func New(handler http.Handler) *HandlerAdapter {
 	}
 }
 
-func (h *HandlerAdapter) Proxy(event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func (h *HandlerAdapter) Proxy(ctx context.Context, event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	req, err := h.ProxyEventToHTTPRequest(event)
 	if err != nil {
 		return core.GatewayTimeout(), core.NewLoggedError("Could not convert proxy event to request: %v", err)
 	}
 
 	w := core.NewProxyResponseWriter()
-	h.handler.ServeHTTP(http.ResponseWriter(w), req)
+	h.handler.ServeHTTP(http.ResponseWriter(w), req.WithContext(ctx))
 
 	resp, err := w.GetProxyResponse()
 	if err != nil {
diff --git a/httpadapter/adapter_test.go b/httpadapter/adapter_test.go
index 1e31bfd..a03b52f 100644
--- a/httpadapter/adapter_test.go
+++ b/httpadapter/adapter_test.go
@@ -1,6 +1,7 @@
 package httpadapter_test
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
@@ -36,7 +37,7 @@ var _ = Describe("HTTPAdapter tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			resp, err := adapter.Proxy(req)
+			resp, err := adapter.Proxy(context.Background(), req)
 
 			Expect(err).To(BeNil())
 			Expect(resp.StatusCode).To(Equal(200))
diff --git a/negroni/adapter.go b/negroni/adapter.go
index 74e1ce9..6ca06e3 100644
--- a/negroni/adapter.go
+++ b/negroni/adapter.go
@@ -1,6 +1,7 @@
 package negroniadapter
 
 import (
+	"context"
 	"net/http"
 
 	"github.com/aws/aws-lambda-go/events"
@@ -19,14 +20,14 @@ func New(n *negroni.Negroni) *NegroniAdapter {
 	}
 }
 
-func (h *NegroniAdapter) Proxy(event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func (h *NegroniAdapter) Proxy(ctx context.Context, event events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	req, err := h.ProxyEventToHTTPRequest(event)
 	if err != nil {
 		return core.GatewayTimeout(), core.NewLoggedError("Could not convert proxy event to request: %v", err)
 	}
 
 	w := core.NewProxyResponseWriter()
-	h.n.ServeHTTP(http.ResponseWriter(w), req)
+	h.n.ServeHTTP(http.ResponseWriter(w), req.WithContext(ctx))
 
 	resp, err := w.GetProxyResponse()
 	if err != nil {
diff --git a/negroni/adapter_test.go b/negroni/adapter_test.go
index 22480b6..d689427 100644
--- a/negroni/adapter_test.go
+++ b/negroni/adapter_test.go
@@ -1,6 +1,7 @@
 package negroniadapter_test
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
@@ -42,7 +43,7 @@ var _ = Describe("NegroniAdapter tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			homePageResp, homePageReqErr := adapter.Proxy(homePageReq)
+			homePageResp, homePageReqErr := adapter.Proxy(context.Background(), homePageReq)
 
 			Expect(homePageReqErr).To(BeNil())
 			Expect(homePageResp.StatusCode).To(Equal(200))
@@ -53,7 +54,7 @@ var _ = Describe("NegroniAdapter tests", func() {
 				HTTPMethod: "GET",
 			}
 
-			productsPageResp, productsPageReqErr := adapter.Proxy(productsPageReq)
+			productsPageResp, productsPageReqErr := adapter.Proxy(context.Background(), productsPageReq)
 
 			Expect(productsPageReqErr).To(BeNil())
 			Expect(productsPageResp.StatusCode).To(Equal(200))
diff --git a/sample/main.go b/sample/main.go
index 747dde0..4b20f98 100644
--- a/sample/main.go
+++ b/sample/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"log"
 	"net/http"
 	"strconv"
@@ -15,7 +16,7 @@ var ginLambda *ginadapter.GinLambda
 
 // Handler is the main entry point for Lambda. Receives a proxy request and
 // returns a proxy response
-func Handler(req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+func Handler(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	if ginLambda == nil {
 		// stdout and stderr are sent to AWS CloudWatch Logs
 		log.Printf("Gin cold start")
@@ -27,7 +28,7 @@ func Handler(req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse,
 		ginLambda = ginadapter.New(r)
 	}
 
-	return ginLambda.Proxy(req)
+	return ginLambda.Proxy(ctx, req)
 }
 
 func main() {