@@ -52,11 +52,18 @@ pub struct Response {
5252
5353impl Response {
5454 #[ cfg( feature = "std" ) ]
55- pub ( crate ) fn create ( mut parent : ResponseLazy , is_head : bool ) -> Result < Response , Error > {
55+ pub ( crate ) fn create (
56+ mut parent : ResponseLazy ,
57+ is_head : bool ,
58+ max_body_size : Option < usize > ,
59+ ) -> Result < Response , Error > {
5660 let mut body = Vec :: new ( ) ;
5761 if !is_head && parent. status_code != 204 && parent. status_code != 304 {
5862 for byte in & mut parent {
5963 let ( byte, length) = byte?;
64+ if max_body_size. is_some_and ( |max| body. len ( ) . saturating_add ( length) > max) {
65+ return Err ( Error :: BodyOverflow ) ;
66+ }
6067 body. reserve ( length) ;
6168 body. push ( byte) ;
6269 }
@@ -79,6 +86,7 @@ impl Response {
7986 is_head : bool ,
8087 max_headers_size : Option < usize > ,
8188 max_status_line_len : Option < usize > ,
89+ max_body_size : Option < usize > ,
8290 ) -> Result < Response , Error > {
8391 use HttpStreamState :: * ;
8492
@@ -98,6 +106,10 @@ impl Response {
98106 EndOnClose => {
99107 while let Some ( byte_result) = read_until_closed_async ( & mut stream) . await {
100108 let ( byte, length) = byte_result?;
109+ if max_body_size. is_some_and ( |max| body. len ( ) . saturating_add ( length) > max)
110+ {
111+ return Err ( Error :: BodyOverflow ) ;
112+ }
101113 body. reserve ( length) ;
102114 body. push ( byte) ;
103115 }
@@ -107,6 +119,11 @@ impl Response {
107119 read_with_content_length_async ( & mut stream, & mut length) . await
108120 {
109121 let ( byte, expected_length) = byte_result?;
122+ if max_body_size
123+ . is_some_and ( |max| body. len ( ) . saturating_add ( expected_length) > max)
124+ {
125+ return Err ( Error :: BodyOverflow ) ;
126+ }
110127 body. reserve ( expected_length) ;
111128 body. push ( byte) ;
112129 }
@@ -123,6 +140,10 @@ impl Response {
123140 . await
124141 {
125142 let ( byte, length) = byte_result?;
143+ if max_body_size. is_some_and ( |max| body. len ( ) . saturating_add ( length) > max)
144+ {
145+ return Err ( Error :: BodyOverflow ) ;
146+ }
126147 body. reserve ( length) ;
127148 body. push ( byte) ;
128149 } ,
@@ -290,6 +311,8 @@ pub struct ResponseLazy {
290311 stream : HttpStreamBytes ,
291312 state : HttpStreamState ,
292313 max_trailing_headers_size : Option < usize > ,
314+ max_body_size : Option < usize > ,
315+ bytes_read : usize ,
293316}
294317
295318#[ cfg( feature = "std" ) ]
@@ -301,6 +324,7 @@ impl ResponseLazy {
301324 stream : HttpStream ,
302325 max_headers_size : Option < usize > ,
303326 max_status_line_len : Option < usize > ,
327+ max_body_size : Option < usize > ,
304328 ) -> Result < ResponseLazy , Error > {
305329 let mut stream = BufReader :: with_capacity ( BACKING_READ_BUFFER_LENGTH , stream) . bytes ( ) ;
306330 let ResponseMetadata {
@@ -319,6 +343,8 @@ impl ResponseLazy {
319343 stream,
320344 state,
321345 max_trailing_headers_size,
346+ max_body_size,
347+ bytes_read : 0 ,
322348 } )
323349 }
324350
@@ -333,6 +359,9 @@ impl ResponseLazy {
333359 stream : BufReader :: with_capacity ( 1 , http_stream) . bytes ( ) ,
334360 state : HttpStreamState :: EndOnClose ,
335361 max_trailing_headers_size : None ,
362+ // Body was already fully loaded and size-checked by send_async
363+ max_body_size : None ,
364+ bytes_read : 0 ,
336365 }
337366 }
338367}
@@ -343,7 +372,7 @@ impl Iterator for ResponseLazy {
343372
344373 fn next ( & mut self ) -> Option < Self :: Item > {
345374 use HttpStreamState :: * ;
346- match self . state {
375+ let result = match self . state {
347376 EndOnClose => read_until_closed ( & mut self . stream ) ,
348377 ContentLength ( ref mut length) => read_with_content_length ( & mut self . stream , length) ,
349378 Chunked ( ref mut expecting_chunks, ref mut length, ref mut content_length) =>
@@ -355,7 +384,17 @@ impl Iterator for ResponseLazy {
355384 content_length,
356385 self . max_trailing_headers_size ,
357386 ) ,
387+ } ;
388+
389+ // Check body size limit before returning the byte
390+ if let Some ( Ok ( ( _, expected_length) ) ) = & result {
391+ if self . max_body_size . is_some_and ( |max| self . bytes_read + expected_length > max) {
392+ return Some ( Err ( Error :: BodyOverflow ) ) ;
393+ }
394+ self . bytes_read += 1 ;
358395 }
396+
397+ result
359398 }
360399}
361400
@@ -537,7 +576,7 @@ macro_rules! define_read_methods {
537576 return None ;
538577 }
539578 * chunk_length = incoming_length;
540- * content_length += incoming_length;
579+ * content_length = content_length . saturating_add ( incoming_length) ;
541580 }
542581
543582 if * chunk_length > 0 {
0 commit comments