@@ -18,24 +18,76 @@ import (
1818 "github.com/aws/aws-lambda-go/lambda"
1919)
2020
21+ type detectContentTypeContextKey struct {}
22+
23+ // WithDetectContentType sets the behavior of content type detection when the Content-Type header is not already provided.
24+ // When true, the first Write call will pass the intial bytes to http.DetectContentType.
25+ // When false, and if no Content-Type is provided, no Content-Type will be sent back to Lambda,
26+ // and the Lambda Function URL will fallback to it's default.
27+ //
28+ // Note: The http.ResponseWriter passed to the handler is unbuffered.
29+ // This may result in different Content-Type headers in the Function URL response when compared to http.ListenAndServe.
30+ //
31+ // Usage:
32+ //
33+ // lambdaurl.Start(
34+ // http.HandlerFunc(func (w http.ResponseWriter, r *http.Request) {
35+ // w.Write("<!DOCTYPE html><html></html>")
36+ // }),
37+ // lambdaurl.WithDetectContentType(true)
38+ // )
39+ func WithDetectContentType (detectContentType bool ) lambda.Option {
40+ return lambda .WithContextValue (detectContentTypeContextKey {}, detectContentType )
41+ }
42+
2143type httpResponseWriter struct {
44+ detectContentType bool
45+ header http.Header
46+ writer io.Writer
47+ once sync.Once
48+ ready chan <- header
49+ }
50+
51+ type header struct {
52+ code int
2253 header http.Header
23- writer io.Writer
24- once sync.Once
25- status chan <- int
2654}
2755
2856func (w * httpResponseWriter ) Header () http.Header {
57+ if w .header == nil {
58+ w .header = http.Header {}
59+ }
2960 return w .header
3061}
3162
3263func (w * httpResponseWriter ) Write (p []byte ) (int , error ) {
33- w .once . Do ( func () { w . status <- http .StatusOK } )
64+ w .writeHeader ( http .StatusOK , p )
3465 return w .writer .Write (p )
3566}
3667
3768func (w * httpResponseWriter ) WriteHeader (statusCode int ) {
38- w .once .Do (func () { w .status <- statusCode })
69+ w .writeHeader (statusCode , nil )
70+ }
71+
72+ func (w * httpResponseWriter ) writeHeader (statusCode int , initialPayload []byte ) {
73+ w .once .Do (func () {
74+ if w .detectContentType {
75+ if w .Header ().Get ("Content-Type" ) == "" {
76+ w .Header ().Set ("Content-Type" , detectContentType (initialPayload ))
77+ }
78+ }
79+ w .ready <- header {code : statusCode , header : w .header }
80+ })
81+ }
82+
83+ func detectContentType (p []byte ) string {
84+ // http.DetectContentType returns "text/plain; charset=utf-8" for nil and zero-length byte slices.
85+ // This is a weird behavior, since otherwise it defaults to "application/octet-stream"! So we'll do that.
86+ // This differs from http.ListenAndServe, which set no Content-Type when the initial Flush body is empty.
87+ if len (p ) == 0 {
88+ return "application/octet-stream"
89+ }
90+ return http .DetectContentType (p )
3991}
4092
4193type requestContextKey struct {}
@@ -46,11 +98,13 @@ func RequestFromContext(ctx context.Context) (*events.LambdaFunctionURLRequest,
4698 return req , ok
4799}
48100
49- // Wrap converts an http.Handler into a lambda request handler.
101+ // Wrap converts an http.Handler into a Lambda request handler.
102+ //
50103// Only Lambda Function URLs configured with `InvokeMode: RESPONSE_STREAM` are supported with the returned handler.
51- // The response body of the handler will conform to the content-type `application/vnd.awslambda.http-integration-response`
104+ // The response body of the handler will conform to the content-type `application/vnd.awslambda.http-integration-response`.
52105func Wrap (handler http.Handler ) func (context.Context , * events.LambdaFunctionURLRequest ) (* events.LambdaFunctionURLStreamingResponse , error ) {
53106 return func (ctx context.Context , request * events.LambdaFunctionURLRequest ) (* events.LambdaFunctionURLStreamingResponse , error ) {
107+
54108 var body io.Reader = strings .NewReader (request .Body )
55109 if request .IsBase64Encoded {
56110 body = base64 .NewDecoder (base64 .StdEncoding , body )
@@ -67,21 +121,28 @@ func Wrap(handler http.Handler) func(context.Context, *events.LambdaFunctionURLR
67121 for k , v := range request .Headers {
68122 httpRequest .Header .Add (k , v )
69123 }
70- status := make ( chan int ) // Signals when it's OK to start returning the response body to Lambda
71- header := http. Header {}
124+
125+ ready := make ( chan header ) // Signals when it's OK to start returning the response body to Lambda
72126 r , w := io .Pipe ()
127+ responseWriter := & httpResponseWriter {writer : w , ready : ready }
128+ if detectContentType , ok := ctx .Value (detectContentTypeContextKey {}).(bool ); ok {
129+ responseWriter .detectContentType = detectContentType
130+ }
73131 go func () {
74- defer close (status )
132+ defer close (ready )
75133 defer w .Close () // TODO: recover and CloseWithError the any panic value once the runtime API client supports plumbing fatal errors through the reader
76- handler .ServeHTTP (& httpResponseWriter {writer : w , header : header , status : status }, httpRequest )
134+ //nolint:errcheck
135+ defer responseWriter .Write (nil ) // force default status, headers, content type detection, if none occured during the execution of the handler
136+ handler .ServeHTTP (responseWriter , httpRequest )
77137 }()
138+ header := <- ready
78139 response := & events.LambdaFunctionURLStreamingResponse {
79140 Body : r ,
80- StatusCode : <- status ,
141+ StatusCode : header . code ,
81142 }
82- if len (header ) > 0 {
83- response .Headers = make (map [string ]string , len (header ))
84- for k , v := range header {
143+ if len (header . header ) > 0 {
144+ response .Headers = make (map [string ]string , len (header . header ))
145+ for k , v := range header . header {
85146 if k == "Set-Cookie" {
86147 response .Cookies = v
87148 } else {
0 commit comments