Skip to content

Commit

Permalink
feat: add support for LambdaFuctionURLRequest/Response (awslabs#172)
Browse files Browse the repository at this point in the history
Signed-off-by: thomasgouveia <[email protected]>
  • Loading branch information
thomasgouveia authored and dza89 committed Sep 12, 2023
1 parent 8c74d92 commit 5804758
Show file tree
Hide file tree
Showing 7 changed files with 459 additions and 0 deletions.
169 changes: 169 additions & 0 deletions core/requestFnURL.go
Original file line number Diff line number Diff line change
@@ -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
}
117 changes: 117 additions & 0 deletions core/responseFnURL.go
Original file line number Diff line number Diff line change
@@ -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
}
12 changes: 12 additions & 0 deletions core/typesFnURL.go
Original file line number Diff line number Diff line change
@@ -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}
}
13 changes: 13 additions & 0 deletions handlerfunc/adapterFnURL.go
Original file line number Diff line number Diff line change
@@ -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)
}
48 changes: 48 additions & 0 deletions handlerfunc/adapterFnURL_test.go
Original file line number Diff line number Diff line change
@@ -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))
})
})
})
Loading

0 comments on commit 5804758

Please sign in to comment.