Skip to content

Commit 503bab0

Browse files
committed
fix: streaming mode response writer ownership
1 parent 65b7b73 commit 503bab0

File tree

1 file changed

+19
-10
lines changed

1 file changed

+19
-10
lines changed

handler/functionurl.go

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -180,10 +180,17 @@ func (w *functionURLStreamingResponseWriter) WriteHeader(statusCode int) {
180180
}
181181

182182
w.resCh <- w.res
183-
close(w.resCh)
184183
}
185184
}
186185

186+
func (w *functionURLStreamingResponseWriter) Close() error {
187+
if w.body == nil {
188+
return nil
189+
}
190+
191+
return w.body.Close()
192+
}
193+
187194
func handleFunctionURLStreaming(ctx context.Context, event events.LambdaFunctionURLRequest, adapter AdapterFunc) (*events.LambdaFunctionURLStreamingResponse, error) {
188195
req, err := convertFunctionURLRequest(ctx, event)
189196
if err != nil {
@@ -192,20 +199,22 @@ func handleFunctionURLStreaming(ctx context.Context, event events.LambdaFunction
192199

193200
resCh := make(chan *events.LambdaFunctionURLStreamingResponse)
194201
errCh := make(chan error)
195-
w := functionURLStreamingResponseWriter{
196-
headers: make(http.Header),
197-
resCh: resCh,
198-
}
199202

200-
go func() {
201-
defer w.body.Close()
203+
go func(resCh chan<- *events.LambdaFunctionURLStreamingResponse, errCh chan<- error) {
204+
defer close(resCh)
205+
defer close(errCh)
206+
207+
w := functionURLStreamingResponseWriter{
208+
headers: make(http.Header),
209+
resCh: resCh,
210+
}
211+
212+
defer w.Close()
202213

203214
if err := adapter(ctx, req, &w); err != nil {
204215
errCh <- err
205216
}
206-
207-
close(errCh)
208-
}()
217+
}(resCh, errCh)
209218

210219
select {
211220
case res := <-resCh:

0 commit comments

Comments
 (0)