Skip to content

Commit a660c21

Browse files
bmoffattbmoffatt
authored
Add response type for Lambda Function URL Streaming Responses (#494)
1 parent 47e703d commit a660c21

File tree

2 files changed

+166
-0
lines changed

2 files changed

+166
-0
lines changed

‎events/lambda_function_urls.go

+74
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,14 @@
22

33
package events
44

5+
import (
6+
"bytes"
7+
"encoding/json"
8+
"errors"
9+
"io"
10+
"net/http"
11+
)
12+
513
// LambdaFunctionURLRequest contains data coming from the HTTP request to a Lambda Function URL.
614
type LambdaFunctionURLRequest struct {
715
Version string `json:"version"` // Version is expected to be `"2.0"`
@@ -59,3 +67,69 @@ type LambdaFunctionURLResponse struct {
5967
IsBase64Encoded bool `json:"isBase64Encoded"`
6068
Cookies []string `json:"cookies"`
6169
}
70+
71+
// LambdaFunctionURLStreamingResponse models the response to a Lambda Function URL when InvokeMode is RESPONSE_STREAM.
72+
// If the InvokeMode of the Function URL is BUFFERED (default), use LambdaFunctionURLResponse instead.
73+
//
74+
// Example:
75+
//
76+
// lambda.Start(func() (*events.LambdaFunctionURLStreamingResponse, error) {
77+
// return &events.LambdaFunctionURLStreamingResponse{
78+
// StatusCode: 200,
79+
// Headers: map[string]string{
80+
// "Content-Type": "text/html",
81+
// },
82+
// Body: strings.NewReader("<html><body>Hello World!</body></html>"),
83+
// }, nil
84+
// })
85+
type LambdaFunctionURLStreamingResponse struct {
86+
prelude *bytes.Buffer
87+
88+
StatusCode int
89+
Headers map[string]string
90+
Body io.Reader
91+
Cookies []string
92+
}
93+
94+
func (r *LambdaFunctionURLStreamingResponse) Read(p []byte) (n int, err error) {
95+
if r.prelude == nil {
96+
if r.StatusCode == 0 {
97+
r.StatusCode = http.StatusOK
98+
}
99+
b, err := json.Marshal(struct {
100+
StatusCode int `json:"statusCode"`
101+
Headers map[string]string `json:"headers,omitempty"`
102+
Cookies []string `json:"cookies,omitempty"`
103+
}{
104+
StatusCode: r.StatusCode,
105+
Headers: r.Headers,
106+
Cookies: r.Cookies,
107+
})
108+
if err != nil {
109+
return 0, err
110+
}
111+
r.prelude = bytes.NewBuffer(append(b, 0, 0, 0, 0, 0, 0, 0, 0))
112+
}
113+
if r.prelude.Len() > 0 {
114+
return r.prelude.Read(p)
115+
}
116+
if r.Body == nil {
117+
return 0, io.EOF
118+
}
119+
return r.Body.Read(p)
120+
}
121+
122+
func (r *LambdaFunctionURLStreamingResponse) Close() error {
123+
if closer, ok := r.Body.(io.ReadCloser); ok {
124+
return closer.Close()
125+
}
126+
return nil
127+
}
128+
129+
func (r *LambdaFunctionURLStreamingResponse) MarshalJSON() ([]byte, error) {
130+
return nil, errors.New("not json")
131+
}
132+
133+
func (r *LambdaFunctionURLStreamingResponse) ContentType() string {
134+
return "application/vnd.awslambda.http-integration-response"
135+
}

‎events/lambda_function_urls_test.go

+92
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,14 @@ package events
44

55
import (
66
"encoding/json"
7+
"errors"
78
"io/ioutil" //nolint: staticcheck
9+
"net/http"
10+
"strings"
811
"testing"
912

1013
"github.com/stretchr/testify/assert"
14+
"github.com/stretchr/testify/require"
1115
)
1216

1317
func TestLambdaFunctionURLResponseMarshaling(t *testing.T) {
@@ -55,3 +59,91 @@ func TestLambdaFunctionURLRequestMarshaling(t *testing.T) {
5559

5660
assert.JSONEq(t, string(inputJSON), string(outputJSON))
5761
}
62+
63+
func TestLambdaFunctionURLStreamingResponseMarshaling(t *testing.T) {
64+
for _, test := range []struct {
65+
name string
66+
response *LambdaFunctionURLStreamingResponse
67+
expectedHead string
68+
expectedBody string
69+
}{
70+
{
71+
"empty",
72+
&LambdaFunctionURLStreamingResponse{},
73+
`{"statusCode":200}`,
74+
"",
75+
},
76+
{
77+
"just the status code",
78+
&LambdaFunctionURLStreamingResponse{
79+
StatusCode: http.StatusTeapot,
80+
},
81+
`{"statusCode":418}`,
82+
"",
83+
},
84+
{
85+
"status and headers and cookies and body",
86+
&LambdaFunctionURLStreamingResponse{
87+
StatusCode: http.StatusTeapot,
88+
Headers: map[string]string{"hello": "world"},
89+
Cookies: []string{"cookies", "are", "yummy"},
90+
Body: strings.NewReader(`<html>Hello Hello</html>`),
91+
},
92+
`{"statusCode":418, "headers":{"hello":"world"}, "cookies":["cookies","are","yummy"]}`,
93+
`<html>Hello Hello</html>`,
94+
},
95+
} {
96+
t.Run(test.name, func(t *testing.T) {
97+
response, err := ioutil.ReadAll(test.response)
98+
require.NoError(t, err)
99+
sep := "\x00\x00\x00\x00\x00\x00\x00\x00"
100+
responseParts := strings.Split(string(response), sep)
101+
require.Len(t, responseParts, 2)
102+
head := string(responseParts[0])
103+
body := string(responseParts[1])
104+
assert.JSONEq(t, test.expectedHead, head)
105+
assert.Equal(t, test.expectedBody, body)
106+
assert.NoError(t, test.response.Close())
107+
})
108+
}
109+
}
110+
111+
type readCloser struct {
112+
closed bool
113+
err error
114+
reader *strings.Reader
115+
}
116+
117+
func (r *readCloser) Read(p []byte) (int, error) {
118+
return r.reader.Read(p)
119+
}
120+
121+
func (r *readCloser) Close() error {
122+
r.closed = true
123+
return r.err
124+
}
125+
126+
func TestLambdaFunctionURLStreamingResponsePropogatesInnerClose(t *testing.T) {
127+
for _, test := range []struct {
128+
name string
129+
closer *readCloser
130+
err error
131+
}{
132+
{
133+
"closer no err",
134+
&readCloser{},
135+
nil,
136+
},
137+
{
138+
"closer with err",
139+
&readCloser{err: errors.New("yolo")},
140+
errors.New("yolo"),
141+
},
142+
} {
143+
t.Run(test.name, func(t *testing.T) {
144+
response := &LambdaFunctionURLStreamingResponse{Body: test.closer}
145+
assert.Equal(t, test.err, response.Close())
146+
assert.True(t, test.closer.closed)
147+
})
148+
}
149+
}

0 commit comments

Comments
 (0)