File tree Expand file tree Collapse file tree 2 files changed +56
-2
lines changed Expand file tree Collapse file tree 2 files changed +56
-2
lines changed Original file line number Diff line number Diff line change @@ -9,6 +9,7 @@ package http
9
9
import (
10
10
"bytes"
11
11
"internal/testenv"
12
+ "io"
12
13
"io/fs"
13
14
"net/url"
14
15
"os"
@@ -189,6 +190,45 @@ func TestNoUnicodeStrings(t *testing.T) {
189
190
}
190
191
}
191
192
193
+ type requestTooLargerResponseWriter struct {
194
+ called bool
195
+ }
196
+
197
+ func (rw * requestTooLargerResponseWriter ) Header () Header { return Header {} }
198
+ func (rw * requestTooLargerResponseWriter ) Write (b []byte ) (int , error ) { return len (b ), nil }
199
+ func (rw * requestTooLargerResponseWriter ) WriteHeader (statusCode int ) {}
200
+ func (rw * requestTooLargerResponseWriter ) requestTooLarge () {
201
+ rw .called = true
202
+ }
203
+
204
+ type wrapper struct {
205
+ ResponseWriter
206
+ }
207
+
208
+ func (w * wrapper ) Unwrap () ResponseWriter {
209
+ return w .ResponseWriter
210
+ }
211
+
212
+ func TestMaxBytesReaderUnwrapTriggersRequestTooLarge (t * testing.T ) {
213
+ body := strings .NewReader ("123456" )
214
+ limit := int64 (5 )
215
+
216
+ innerRw := & requestTooLargerResponseWriter {}
217
+ wrappedRw := & wrapper {ResponseWriter : innerRw }
218
+
219
+ l := MaxBytesReader (wrappedRw , io .NopCloser (body ), limit )
220
+
221
+ buf := make ([]byte , 10 )
222
+ _ , err := l .Read (buf )
223
+
224
+ if _ , ok := err .(* MaxBytesError ); ! ok {
225
+ t .Errorf ("expected MaxBytesError, got %T" , err )
226
+ }
227
+ if ! innerRw .called {
228
+ t .Errorf ("expected requestTooLarge to be called, but it wasn't" )
229
+ }
230
+ }
231
+
192
232
func TestProtocols (t * testing.T ) {
193
233
var p Protocols
194
234
if p .HTTP1 () {
Original file line number Diff line number Diff line change @@ -1243,9 +1243,23 @@ func (l *maxBytesReader) Read(p []byte) (n int, err error) {
1243
1243
type requestTooLarger interface {
1244
1244
requestTooLarge ()
1245
1245
}
1246
- if res , ok := l .w .(requestTooLarger ); ok {
1247
- res .requestTooLarge ()
1246
+ // Unwrap the ResponseWriter wrappers until we find one that implements
1247
+ // the server-only requestTooLarger interface, then call requestTooLarge().
1248
+ // This ensures that even if the ResponseWriter is wrapped by a custom implementation,
1249
+ // the underlying server writer can be notified when the request body is too large.
1250
+ rw := l .w
1251
+ for {
1252
+ if res , ok := rw .(requestTooLarger ); ok {
1253
+ res .requestTooLarge ()
1254
+ break
1255
+ }
1256
+ unwrapper , ok := rw .(rwUnwrapper )
1257
+ if ! ok {
1258
+ break
1259
+ }
1260
+ rw = unwrapper .Unwrap ()
1248
1261
}
1262
+
1249
1263
l .err = & MaxBytesError {l .i }
1250
1264
return n , l .err
1251
1265
}
You can’t perform that action at this time.
0 commit comments