feat: add support for LambdaFuctionURLRequest/Response (awslabs#172)

thomasgouveia · dza89 · commit 5804758a8a8c · 2023-09-12T10:46:14.000+02:00
Signed-off-by: thomasgouveia &lt;gouveia.thomas@outlook.fr&gt;
diff --git a/core/requestFnURL.go b/core/requestFnURL.go
@@ -0,0 +1,169 @@
+// Package core provides utility methods that help convert ALB events
+// into an http.Request and http.ResponseWriter
+package core
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log"
+	"net/http"
+	"os"
+	"strings"
+
+	"github.com/aws/aws-lambda-go/events"
+	"github.com/aws/aws-lambda-go/lambdacontext"
+)
+
+const (
+	// FnURLContextHeader is the custom header key used to store the
+	// Function URL context. To access the Context properties use the
+	// GetContext method of the RequestAccessorFnURL object.
+	FnURLContextHeader = "X-GoLambdaProxy-FnURL-Context"
+)
+
+// RequestAccessorFnURL objects give access to custom Function URL properties
+// in the request.
+type RequestAccessorFnURL struct {
+	stripBasePath string
+}
+
+// GetALBContext extracts the ALB context object from a request's custom header.
+// Returns a populated events.ALBTargetGroupRequestContext object from the request.
+func (r *RequestAccessorFnURL) GetContext(req *http.Request) (events.LambdaFunctionURLRequestContext, error) {
+	if req.Header.Get(FnURLContextHeader) == "" {
+		return events.LambdaFunctionURLRequestContext{}, errors.New("no context header in request")
+	}
+	context := events.LambdaFunctionURLRequestContext{}
+	err := json.Unmarshal([]byte(req.Header.Get(FnURLContextHeader)), &context)
+	if err != nil {
+		log.Println("Error while unmarshalling context")
+		log.Println(err)
+		return events.LambdaFunctionURLRequestContext{}, err
+	}
+	return context, nil
+}
+
+// StripBasePath instructs the RequestAccessor object that the given base
+// path should be removed from the request path before sending it to the
+// framework for routing. This is used when API Gateway is configured with
+// base path mappings in custom domain names.
+func (r *RequestAccessorFnURL) StripBasePath(basePath string) string {
+	if strings.Trim(basePath, " ") == "" {
+		r.stripBasePath = ""
+		return ""
+	}
+
+	newBasePath := basePath
+	if !strings.HasPrefix(newBasePath, "/") {
+		newBasePath = "/" + newBasePath
+	}
+
+	if strings.HasSuffix(newBasePath, "/") {
+		newBasePath = newBasePath[:len(newBasePath)-1]
+	}
+
+	r.stripBasePath = newBasePath
+
+	return newBasePath
+}
+
+// FunctionURLEventToHTTPRequest converts an a Function URL event into a http.Request object.
+// Returns the populated http request with additional custom header for the Function URL context.
+// To access these properties use the GetContext method of the RequestAccessorFnURL object.
+func (r *RequestAccessorFnURL) FunctionURLEventToHTTPRequest(req events.LambdaFunctionURLRequest) (*http.Request, error) {
+	httpRequest, err := r.EventToRequest(req)
+	if err != nil {
+		log.Println(err)
+		return nil, err
+	}
+	return addToHeaderFnURL(httpRequest, req)
+}
+
+// FunctionURLEventToHTTPRequestWithContext converts a Function URL event and context into an http.Request object.
+// Returns the populated http request with lambda context, Function URL RequestContext as part of its context.
+func (r *RequestAccessorFnURL) FunctionURLEventToHTTPRequestWithContext(ctx context.Context, req events.LambdaFunctionURLRequest) (*http.Request, error) {
+	httpRequest, err := r.EventToRequest(req)
+	if err != nil {
+		log.Println(err)
+		return nil, err
+	}
+	return addToContextFnURL(ctx, httpRequest, req), nil
+}
+
+// EventToRequest converts a Function URL event into an http.Request object.
+// Returns the populated request maintaining headers
+func (r *RequestAccessorFnURL) EventToRequest(req events.LambdaFunctionURLRequest) (*http.Request, error) {
+	decodedBody := []byte(req.Body)
+	if req.IsBase64Encoded {
+		base64Body, err := base64.StdEncoding.DecodeString(req.Body)
+		if err != nil {
+			return nil, err
+		}
+		decodedBody = base64Body
+	}
+
+	path := req.RawPath
+	if r.stripBasePath != "" && len(r.stripBasePath) > 1 {
+		if strings.HasPrefix(path, r.stripBasePath) {
+			path = strings.Replace(path, r.stripBasePath, "", 1)
+		}
+	}
+	if !strings.HasPrefix(path, "/") {
+		path = "/" + path
+	}
+
+	serverAddress := "https://" + req.RequestContext.DomainName
+	if customAddress, ok := os.LookupEnv(CustomHostVariable); ok {
+		serverAddress = customAddress
+	}
+
+	path = serverAddress + path + "?" + req.RawQueryString
+
+	httpRequest, err := http.NewRequest(
+		strings.ToUpper(req.RequestContext.HTTP.Method),
+		path,
+		bytes.NewReader(decodedBody),
+	)
+
+	if err != nil {
+		fmt.Printf("Could not convert request %s:%s to http.Request\n", req.RequestContext.HTTP.Method, req.RawPath)
+		log.Println(err)
+		return nil, err
+	}
+
+	for header, val := range req.Headers {
+		httpRequest.Header.Add(header, val)
+	}
+
+	httpRequest.RemoteAddr = req.RequestContext.HTTP.SourceIP
+	httpRequest.RequestURI = httpRequest.URL.RequestURI()
+
+	return httpRequest, nil
+}
+
+func addToHeaderFnURL(req *http.Request, fnUrlRequest events.LambdaFunctionURLRequest) (*http.Request, error) {
+	ctx, err := json.Marshal(fnUrlRequest.RequestContext)
+	if err != nil {
+		log.Println("Could not Marshal Function URL context for custom header")
+		return req, err
+	}
+	req.Header.Set(FnURLContextHeader, string(ctx))
+	return req, nil
+}
+
+// adds context data to http request so we can pass
+func addToContextFnURL(ctx context.Context, req *http.Request, fnUrlRequest events.LambdaFunctionURLRequest) *http.Request {
+	lc, _ := lambdacontext.FromContext(ctx)
+	rc := requestContextFnURL{lambdaContext: lc, fnUrlContext: fnUrlRequest.RequestContext}
+	ctx = context.WithValue(ctx, ctxKey{}, rc)
+	return req.WithContext(ctx)
+}
+
+type requestContextFnURL struct {
+	lambdaContext *lambdacontext.LambdaContext
+	fnUrlContext  events.LambdaFunctionURLRequestContext
+}
diff --git a/core/responseFnURL.go b/core/responseFnURL.go
@@ -0,0 +1,117 @@
+// Package core provides utility methods that help convert proxy events
+// into an http.Request and http.ResponseWriter
+package core
+
+import (
+	"bytes"
+	"encoding/base64"
+	"errors"
+	"net/http"
+	"unicode/utf8"
+
+	"github.com/aws/aws-lambda-go/events"
+)
+
+// ProxyResponseWriterFunctionURL implements http.ResponseWriter and adds the method
+// necessary to return an events.LambdaFunctionURLResponse object
+type ProxyResponseWriterFunctionURL struct {
+	status    int
+	headers   http.Header
+	body      bytes.Buffer
+	observers []chan<- bool
+}
+
+// Ensure implementation satisfies http.ResponseWriter interface
+var (
+	_ http.ResponseWriter = &ProxyResponseWriterFunctionURL{}
+)
+
+// NewProxyResponseWriterFnURL returns a new ProxyResponseWriterFunctionURL object.
+// The object is initialized with an empty map of headers and a status code of -1
+func NewProxyResponseWriterFnURL() *ProxyResponseWriterFunctionURL {
+	return &ProxyResponseWriterFunctionURL{
+		headers:   make(http.Header),
+		status:    defaultStatusCode,
+		observers: make([]chan<- bool, 0),
+	}
+}
+
+func (r *ProxyResponseWriterFunctionURL) CloseNotify() <-chan bool {
+	ch := make(chan bool, 1)
+
+	r.observers = append(r.observers, ch)
+
+	return ch
+}
+
+func (r *ProxyResponseWriterFunctionURL) notifyClosed() {
+	for _, v := range r.observers {
+		v <- true
+	}
+}
+
+// Header implementation from the http.ResponseWriter interface.
+func (r *ProxyResponseWriterFunctionURL) Header() http.Header {
+	return r.headers
+}
+
+// Write sets the response body in the object. If no status code
+// was set before with the WriteHeader method it sets the status
+// for the response to 200 OK.
+func (r *ProxyResponseWriterFunctionURL) Write(body []byte) (int, error) {
+	if r.status == defaultStatusCode {
+		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)
+}
+
+// WriteHeader sets a status code for the response. This method is used
+// for error responses.
+func (r *ProxyResponseWriterFunctionURL) WriteHeader(status int) {
+	r.status = status
+}
+
+// GetProxyResponse converts the data passed to the response writer into
+// an events.LambdaFunctionURLResponse object.
+// Returns a populated proxy response object. If the response is invalid, for example
+// has no headers or an invalid status code returns an error.
+func (r *ProxyResponseWriterFunctionURL) GetProxyResponse() (events.LambdaFunctionURLResponse, error) {
+	r.notifyClosed()
+
+	if r.status == defaultStatusCode {
+		return events.LambdaFunctionURLResponse{}, errors.New("status code not set on response")
+	}
+
+	var output string
+	isBase64 := false
+
+	bb := (&r.body).Bytes()
+
+	if utf8.Valid(bb) {
+		output = string(bb)
+	} else {
+		output = base64.StdEncoding.EncodeToString(bb)
+		isBase64 = true
+	}
+
+	headers := make(map[string]string)
+	for h, v := range r.Header() {
+		headers[h] = v[0]
+	}
+
+	return events.LambdaFunctionURLResponse{
+		StatusCode:      r.status,
+		Headers:         headers,
+		Body:            output,
+		IsBase64Encoded: isBase64,
+	}, nil
+}
diff --git a/core/typesFnURL.go b/core/typesFnURL.go
@@ -0,0 +1,12 @@
+package core
+
+import (
+	"net/http"
+
+	"github.com/aws/aws-lambda-go/events"
+)
+
+// GatewayTimeoutFnURL returns a dafault Gateway Timeout (504) response
+func GatewayTimeoutFnURL() events.LambdaFunctionURLResponse {
+	return events.LambdaFunctionURLResponse{StatusCode: http.StatusGatewayTimeout}
+}
diff --git a/handlerfunc/adapterFnURL.go b/handlerfunc/adapterFnURL.go
@@ -0,0 +1,13 @@
+package handlerfunc
+
+import (
+	"net/http"
+
+	"github.com/awslabs/aws-lambda-go-api-proxy/httpadapter"
+)
+
+type HandlerFuncAdapterFnURL = httpadapter.HandlerAdapterFnURL
+
+func NewFunctionURL(handlerFunc http.HandlerFunc) *HandlerFuncAdapterFnURL {
+	return httpadapter.NewFunctionURL(handlerFunc)
+}
diff --git a/handlerfunc/adapterFnURL_test.go b/handlerfunc/adapterFnURL_test.go
@@ -0,0 +1,48 @@
+package handlerfunc_test
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"net/http"
+
+	"github.com/aws/aws-lambda-go/events"
+	"github.com/awslabs/aws-lambda-go-api-proxy/httpadapter"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("HandlerFuncAdapter tests", func() {
+	Context("Simple ping request", func() {
+		It("Proxies the event correctly", func() {
+			log.Println("Starting test")
+
+			var handler http.Handler = http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+				w.Header().Add("unfortunately-required-header", "")
+				fmt.Fprintf(w, "Go Lambda!!")
+			})
+
+			adapter := httpadapter.NewFunctionURL(handler)
+
+			req := events.LambdaFunctionURLRequest{
+				RequestContext: events.LambdaFunctionURLRequestContext{
+					HTTP: events.LambdaFunctionURLRequestContextHTTPDescription{
+						Method: http.MethodGet,
+						Path:   "/ping",
+					},
+				},
+			}
+
+			resp, err := adapter.ProxyWithContext(context.Background(), req)
+
+			Expect(err).To(BeNil())
+			Expect(resp.StatusCode).To(Equal(200))
+
+			resp, err = adapter.Proxy(req)
+
+			Expect(err).To(BeNil())
+			Expect(resp.StatusCode).To(Equal(200))
+		})
+	})
+})
diff --git a/httpadapter/adapterFnURL.go b/httpadapter/adapterFnURL.go
diff --git a/httpadapter/adapterFnURL_test.go b/httpadapter/adapterFnURL_test.go
