1use bytes::Bytes;
21use futures_util::Stream;
22use std::pin::Pin;
23use std::sync::atomic::{AtomicBool, Ordering};
24use std::sync::Arc;
25use std::task::{Context, Poll};
26
27use super::request_body_error::RequestBodyError;
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum BodyState {
32 Unread,
34 Streaming,
36 Complete,
38 Error,
40}
41
42pub struct RequestBody {
68 inner: Option<BodyInner>,
69 declared_length: Option<u64>,
70 bytes_received: u64,
71 state: BodyState,
72 max_bytes: u64,
73 consumed: Arc<AtomicBool>,
77}
78
79#[allow(dead_code)]
81enum BodyInner {
82 Incoming {
84 stream: Pin<Box<dyn Stream<Item = Result<Bytes, IncomingError>> + Send + 'static>>,
85 },
86 Fixed { data: Bytes, offset: usize },
88 Empty,
90}
91
92#[derive(Debug)]
94pub struct IncomingError(pub(crate) String);
95
96impl std::fmt::Display for IncomingError {
97 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
98 write!(f, "incoming body error: {}", self.0)
99 }
100}
101
102impl std::error::Error for IncomingError {}
103
104impl From<IncomingError> for RequestBodyError {
105 fn from(e: IncomingError) -> Self {
106 Self::Transport(e.0)
107 }
108}
109
110impl RequestBody {
111 pub fn empty() -> Self {
113 Self {
114 inner: Some(BodyInner::Empty),
115 declared_length: None,
116 bytes_received: 0,
117 state: BodyState::Unread,
118 max_bytes: u64::MAX,
119 consumed: Arc::new(AtomicBool::new(true)),
120 }
121 }
122
123 pub fn from_bytes(data: impl Into<Bytes>, max_bytes: u64) -> Self {
128 let data = data.into();
129 let len = data.len() as u64;
130 Self {
131 inner: Some(BodyInner::Fixed { data, offset: 0 }),
132 declared_length: Some(len),
133 bytes_received: 0,
134 state: BodyState::Unread,
135 max_bytes,
136 consumed: Arc::new(AtomicBool::new(false)),
137 }
138 }
139
140 #[allow(dead_code)]
146 pub(crate) fn from_incoming(
147 stream: impl Stream<Item = Result<Bytes, IncomingError>> + Send + 'static,
148 declared_length: Option<u64>,
149 max_bytes: u64,
150 ) -> Self {
151 Self {
152 inner: Some(BodyInner::Incoming {
153 stream: Box::pin(stream),
154 }),
155 declared_length,
156 bytes_received: 0,
157 state: BodyState::Unread,
158 max_bytes,
159 consumed: Arc::new(AtomicBool::new(false)),
160 }
161 }
162
163 pub fn declared_length(&self) -> Option<u64> {
165 self.declared_length
166 }
167
168 pub fn bytes_received(&self) -> u64 {
170 self.bytes_received
171 }
172
173 pub fn is_complete(&self) -> bool {
175 self.state == BodyState::Complete
176 }
177
178 pub fn state(&self) -> BodyState {
180 self.state
181 }
182
183 pub fn max_bytes(&self) -> u64 {
185 self.max_bytes
186 }
187
188 pub(crate) fn consumed_flag(&self) -> Arc<AtomicBool> {
194 self.consumed.clone()
195 }
196
197 pub(crate) fn was_fully_consumed(&self) -> bool {
200 self.consumed.load(Ordering::Acquire)
201 }
202
203 fn mark_consumed(&self) {
205 self.consumed.store(true, Ordering::Release);
206 }
207
208 pub async fn read_all(mut self) -> Result<Bytes, RequestBodyError> {
218 if self.state == BodyState::Complete || self.state == BodyState::Error {
219 return Err(RequestBodyError::AlreadyConsumed);
220 }
221 if self.state == BodyState::Streaming {
222 return Err(RequestBodyError::MixedConsumptionMode);
223 }
224
225 let inner = self.inner.take().ok_or(RequestBodyError::AlreadyConsumed)?;
226
227 match inner {
228 BodyInner::Empty => {
229 self.state = BodyState::Complete;
230 self.mark_consumed();
231 Ok(Bytes::new())
232 }
233 BodyInner::Fixed { data, offset } => {
234 let remaining = &data[offset..];
235 let total = self
236 .bytes_received
237 .checked_add(remaining.len() as u64)
238 .ok_or(RequestBodyError::LimitExceeded {
239 limit: self.max_bytes,
240 received: u64::MAX,
241 })?;
242 if total > self.max_bytes {
243 self.state = BodyState::Error;
244 return Err(RequestBodyError::LimitExceeded {
245 limit: self.max_bytes,
246 received: total,
247 });
248 }
249 self.bytes_received = total;
250 self.state = BodyState::Complete;
251 self.mark_consumed();
252 Ok(data.slice(offset..))
253 }
254 BodyInner::Incoming { mut stream } => {
255 let mut buf = Vec::new();
256 use futures_util::StreamExt;
257 while let Some(item) = stream.next().await {
258 let chunk = item.map_err(|e| RequestBodyError::Transport(e.0))?;
259 let new_total = self.bytes_received.checked_add(chunk.len() as u64).ok_or(
260 RequestBodyError::LimitExceeded {
261 limit: self.max_bytes,
262 received: u64::MAX,
263 },
264 )?;
265 if new_total > self.max_bytes {
266 self.state = BodyState::Error;
267 return Err(RequestBodyError::LimitExceeded {
268 limit: self.max_bytes,
269 received: new_total,
270 });
271 }
272 self.bytes_received = new_total;
273 buf.extend_from_slice(&chunk);
274 }
275 self.state = BodyState::Complete;
276 if let Some(declared) = self.declared_length {
278 if self.bytes_received < declared {
279 let received = self.bytes_received;
280 return Err(RequestBodyError::PrematureEof {
281 received,
282 expected: Some(declared),
283 });
284 }
285 }
286 self.mark_consumed();
287 Ok(Bytes::from(buf))
288 }
289 }
290 }
291
292 pub async fn next_chunk(&mut self) -> Result<Option<Bytes>, RequestBodyError> {
306 if self.state == BodyState::Error {
307 return Err(RequestBodyError::AlreadyConsumed);
308 }
309 if self.state == BodyState::Complete {
310 return Ok(None);
311 }
312
313 if self.state == BodyState::Unread {
315 self.state = BodyState::Streaming;
316 }
317
318 let inner = self
319 .inner
320 .as_mut()
321 .ok_or(RequestBodyError::AlreadyConsumed)?;
322
323 match inner {
324 BodyInner::Empty => {
325 self.state = BodyState::Complete;
326 self.mark_consumed();
327 Ok(None)
328 }
329 BodyInner::Fixed { data, offset } => {
330 if *offset >= data.len() {
331 self.state = BodyState::Complete;
332 self.mark_consumed();
333 return Ok(None);
334 }
335 let remaining = &data[*offset..];
336 let chunk_size = remaining.len().min(8192);
337 let new_total = self.bytes_received.checked_add(chunk_size as u64).ok_or(
338 RequestBodyError::LimitExceeded {
339 limit: self.max_bytes,
340 received: u64::MAX,
341 },
342 )?;
343 if new_total > self.max_bytes {
344 self.state = BodyState::Error;
345 return Err(RequestBodyError::LimitExceeded {
346 limit: self.max_bytes,
347 received: new_total,
348 });
349 }
350 let chunk = &data[*offset..*offset + chunk_size];
351 *offset += chunk_size;
352 self.bytes_received = new_total;
353 if *offset >= data.len() {
354 self.state = BodyState::Complete;
355 }
356 Ok(Some(Bytes::copy_from_slice(chunk)))
357 }
358 BodyInner::Incoming { stream } => {
359 use futures_util::StreamExt;
360 match stream.next().await {
361 Some(Ok(chunk)) => {
362 let new_total = self.bytes_received.checked_add(chunk.len() as u64).ok_or(
363 RequestBodyError::LimitExceeded {
364 limit: self.max_bytes,
365 received: u64::MAX,
366 },
367 )?;
368 if new_total > self.max_bytes {
369 self.state = BodyState::Error;
370 return Err(RequestBodyError::LimitExceeded {
371 limit: self.max_bytes,
372 received: new_total,
373 });
374 }
375 self.bytes_received = new_total;
376 Ok(Some(chunk))
377 }
378 Some(Err(e)) => {
379 self.state = BodyState::Error;
380 Err(RequestBodyError::Transport(e.0))
381 }
382 None => {
383 self.state = BodyState::Complete;
384 if let Some(declared) = self.declared_length {
386 if self.bytes_received < declared {
387 let received = self.bytes_received;
388 return Err(RequestBodyError::PrematureEof {
389 received,
390 expected: Some(declared),
391 });
392 }
393 }
394 self.mark_consumed();
395 Ok(None)
396 }
397 }
398 }
399 }
400 }
401}
402
403impl std::fmt::Debug for RequestBody {
404 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
405 f.debug_struct("RequestBody")
406 .field("declared_length", &self.declared_length)
407 .field("bytes_received", &self.bytes_received)
408 .field("state", &self.state)
409 .field("max_bytes", &self.max_bytes)
410 .field("consumed", &self.was_fully_consumed())
411 .finish()
412 }
413}
414
415impl Stream for RequestBody {
416 type Item = Result<Bytes, RequestBodyError>;
417
418 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
419 if self.state == BodyState::Complete || self.state == BodyState::Error {
421 return Poll::Ready(None);
422 }
423
424 if self.bytes_received >= self.max_bytes {
426 self.state = BodyState::Error;
427 return Poll::Ready(Some(Err(RequestBodyError::LimitExceeded {
428 limit: self.max_bytes,
429 received: self.bytes_received,
430 })));
431 }
432
433 let max_bytes = self.max_bytes;
434 let bytes_received = self.bytes_received;
435
436 let inner = match self.inner.as_mut() {
437 Some(i) => i,
438 None => return Poll::Ready(None),
439 };
440
441 match inner {
442 BodyInner::Empty => {
443 self.state = BodyState::Complete;
444 self.mark_consumed();
445 Poll::Ready(None)
446 }
447 BodyInner::Fixed { data, offset } => {
448 if *offset >= data.len() {
449 self.state = BodyState::Complete;
450 self.mark_consumed();
451 Poll::Ready(None)
452 } else {
453 let remaining = &data[*offset..];
454 let chunk_size = remaining.len().min(8192);
455 let new_total = match bytes_received.checked_add(chunk_size as u64) {
456 Some(v) => v,
457 None => {
458 self.state = BodyState::Error;
459 return Poll::Ready(Some(Err(RequestBodyError::LimitExceeded {
460 limit: max_bytes,
461 received: u64::MAX,
462 })));
463 }
464 };
465 if new_total > max_bytes {
466 self.state = BodyState::Error;
467 return Poll::Ready(Some(Err(RequestBodyError::LimitExceeded {
468 limit: max_bytes,
469 received: new_total,
470 })));
471 }
472 let chunk = Bytes::copy_from_slice(&data[*offset..*offset + chunk_size]);
473 *offset += chunk_size;
474 self.bytes_received = new_total;
475 if self.state == BodyState::Unread {
476 self.state = BodyState::Streaming;
477 }
478 Poll::Ready(Some(Ok(chunk)))
479 }
480 }
481 BodyInner::Incoming { stream } => match stream.as_mut().poll_next(cx) {
482 Poll::Ready(Some(Ok(chunk))) => {
483 let new_total = match bytes_received.checked_add(chunk.len() as u64) {
484 Some(v) => v,
485 None => {
486 self.state = BodyState::Error;
487 return Poll::Ready(Some(Err(RequestBodyError::LimitExceeded {
488 limit: max_bytes,
489 received: u64::MAX,
490 })));
491 }
492 };
493 if new_total > max_bytes {
494 self.state = BodyState::Error;
495 Poll::Ready(Some(Err(RequestBodyError::LimitExceeded {
496 limit: max_bytes,
497 received: new_total,
498 })))
499 } else {
500 self.bytes_received = new_total;
501 if self.state == BodyState::Unread {
502 self.state = BodyState::Streaming;
503 }
504 Poll::Ready(Some(Ok(chunk)))
505 }
506 }
507 Poll::Ready(Some(Err(e))) => {
508 self.state = BodyState::Error;
509 Poll::Ready(Some(Err(RequestBodyError::Transport(e.0))))
510 }
511 Poll::Ready(None) => {
512 self.state = BodyState::Complete;
513 if let Some(declared) = self.declared_length {
515 if self.bytes_received < declared {
516 let received = self.bytes_received;
517 self.state = BodyState::Error;
518 return Poll::Ready(Some(Err(RequestBodyError::PrematureEof {
519 received,
520 expected: Some(declared),
521 })));
522 }
523 }
524 self.mark_consumed();
525 Poll::Ready(None)
526 }
527 Poll::Pending => Poll::Pending,
528 },
529 }
530 }
531}
532
533#[cfg(test)]
534mod tests {
535 use super::*;
536 use futures_util::StreamExt;
537 use proptest::prelude::*;
538
539 #[tokio::test]
540 async fn empty_body_read_all() {
541 let body = RequestBody::empty();
542 assert_eq!(body.declared_length(), None);
543 assert_eq!(body.bytes_received(), 0);
544 let data = body.read_all().await.unwrap();
545 assert!(data.is_empty());
546 }
547
548 #[tokio::test]
549 async fn fixed_body_read_all() {
550 let body = RequestBody::from_bytes(b"hello".to_vec(), u64::MAX);
551 assert_eq!(body.declared_length(), Some(5));
552 let data = body.read_all().await.unwrap();
553 assert_eq!(&data[..], b"hello");
554 }
555
556 #[tokio::test]
557 async fn fixed_body_streaming() {
558 let mut body = RequestBody::from_bytes(b"hello world".to_vec(), u64::MAX);
559 let mut chunks = Vec::new();
560 while let Some(chunk) = body.next_chunk().await.unwrap() {
561 chunks.push(chunk);
562 }
563 let total: Vec<u8> = chunks.iter().flat_map(|c| c.iter().copied()).collect();
564 assert_eq!(&total[..], b"hello world");
565 }
566
567 #[tokio::test]
568 async fn limit_exceeded_on_read_all() {
569 let body = RequestBody::from_bytes(b"hello".to_vec(), 3);
570 let err = body.read_all().await.unwrap_err();
571 assert!(err.is_limit_exceeded());
572 }
573
574 #[tokio::test]
575 async fn limit_exceeded_on_stream() {
576 let mut body = RequestBody::from_bytes(b"hello".to_vec(), 3);
577 let err = body.next_chunk().await.unwrap_err();
578 assert!(err.is_limit_exceeded());
579 }
580
581 #[tokio::test]
582 async fn already_consumed_after_read_all() {
583 let body = RequestBody::from_bytes(b"hello".to_vec(), u64::MAX);
584 let _data = body.read_all().await.unwrap();
585 }
587
588 #[tokio::test]
589 async fn mixed_consumption_mode() {
590 let large_body = vec![0u8; 16384];
592 let mut body = RequestBody::from_bytes(large_body, u64::MAX);
593 let _chunk = body.next_chunk().await.unwrap();
594 assert_eq!(body.state(), BodyState::Streaming);
596 }
597
598 #[tokio::test]
599 async fn zero_length_body() {
600 let body = RequestBody::from_bytes(Vec::new(), u64::MAX);
601 assert_eq!(body.declared_length(), Some(0));
602 let data = body.read_all().await.unwrap();
603 assert!(data.is_empty());
604 }
605
606 #[tokio::test]
607 async fn stream_trait_works() {
608 let body = RequestBody::from_bytes(b"abc".to_vec(), u64::MAX);
609 let mut stream = body;
610 let mut all = Vec::new();
611 while let Some(chunk) = stream.next().await.transpose().unwrap() {
612 all.extend_from_slice(&chunk);
613 }
614 assert_eq!(&all[..], b"abc");
615 }
616
617 #[test]
618 fn body_state_debug() {
619 assert_eq!(format!("{:?}", BodyState::Unread), "Unread");
620 assert_eq!(format!("{:?}", BodyState::Streaming), "Streaming");
621 assert_eq!(format!("{:?}", BodyState::Complete), "Complete");
622 assert_eq!(format!("{:?}", BodyState::Error), "Error");
623 }
624
625 #[test]
626 fn request_body_debug() {
627 let body = RequestBody::empty();
628 let dbg = format!("{:?}", body);
629 assert!(dbg.contains("RequestBody"));
630 assert!(dbg.contains("Unread"));
631 }
632
633 #[tokio::test]
634 async fn premature_eof_returns_error() {
635 use futures_util::stream;
636 let short_data = b"hi";
638 let declared = 10u64;
639 let body_stream =
640 stream::once(async move { Ok::<_, IncomingError>(Bytes::copy_from_slice(short_data)) });
641 let body = RequestBody::from_incoming(body_stream, Some(declared), u64::MAX);
642 let result = body.read_all().await;
643 assert!(result.is_err());
644 match result.unwrap_err() {
645 RequestBodyError::PrematureEof { received, expected } => {
646 assert_eq!(received, 2);
647 assert_eq!(expected, Some(10));
648 }
649 other => panic!("expected PrematureEof, got: {:?}", other),
650 }
651 }
652
653 #[tokio::test]
654 async fn premature_eof_streaming_returns_error() {
655 use futures_util::stream;
656 let short_data = b"hi";
658 let declared = 10u64;
659 let body_stream =
660 stream::once(async move { Ok::<_, IncomingError>(Bytes::copy_from_slice(short_data)) });
661 let mut body = RequestBody::from_incoming(body_stream, Some(declared), u64::MAX);
662 let chunk = body.next_chunk().await.unwrap();
664 assert!(chunk.is_some());
665 let result = body.next_chunk().await;
667 match result {
668 Err(RequestBodyError::PrematureEof { received, expected }) => {
669 assert_eq!(received, 2);
670 assert_eq!(expected, Some(10));
671 }
672 Ok(None) => {
673 panic!("expected PrematureEof, got Ok(None)");
674 }
675 other => panic!("expected PrematureEof, got: {:?}", other),
676 }
677 }
678
679 #[tokio::test]
680 async fn exact_declared_length_succeeds() {
681 use futures_util::stream;
682 let data = b"hello";
683 let declared = 5u64;
684 let body_stream =
685 stream::once(
686 async move { Ok::<_, IncomingError>(Bytes::copy_from_slice(data.as_slice())) },
687 );
688 let body = RequestBody::from_incoming(body_stream, Some(declared), u64::MAX);
689 let result = body.read_all().await;
690 assert!(result.is_ok());
691 assert_eq!(result.unwrap().as_ref(), b"hello");
692 }
693
694 #[tokio::test]
695 async fn checked_add_overflow_returns_error() {
696 let body = RequestBody::from_bytes(vec![0u8; 200], 100);
697 let result = body.read_all().await;
698 assert!(result.is_err());
699 assert!(result.unwrap_err().is_limit_exceeded());
700 }
701
702 #[test]
703 fn state_transitions_unread_to_complete_on_read() {
704 proptest::proptest!(|(data in prop::collection::vec(any::<u8>(), 0..500))| {
705 let rt = tokio::runtime::Runtime::new().unwrap();
706 let body = RequestBody::from_bytes(data, u64::MAX);
707 prop_assert_eq!(body.state(), BodyState::Unread);
708 let flag = body.consumed_flag();
709 let _ = rt.block_on(body.read_all());
710 prop_assert!(flag.load(std::sync::atomic::Ordering::Acquire));
711 });
712 }
713
714 #[test]
715 fn consumed_flag_set_after_read() {
716 proptest::proptest!(|(data in prop::collection::vec(any::<u8>(), 0..500))| {
717 let rt = tokio::runtime::Runtime::new().unwrap();
718 let body = RequestBody::from_bytes(data, u64::MAX);
719 let flag = body.consumed_flag();
720 prop_assert!(!flag.load(std::sync::atomic::Ordering::Acquire));
721 let _ = rt.block_on(body.read_all());
722 prop_assert!(flag.load(std::sync::atomic::Ordering::Acquire));
723 });
724 }
725
726 #[test]
727 fn chunked_body_via_stream_succeeds() {
728 proptest::proptest!(|(data in prop::collection::vec(any::<u8>(), 0..1000))| {
729 use futures_util::stream;
730
731 let rt = tokio::runtime::Runtime::new().unwrap();
732 let chunk_size = if data.is_empty() { 1 } else { (data[0] as usize % 64) + 1 };
733 let chunks: Vec<Result<Bytes, IncomingError>> = data.chunks(chunk_size)
734 .map(|c| Ok(Bytes::copy_from_slice(c)))
735 .collect();
736 let body_stream = stream::iter(chunks);
737 let body = RequestBody::from_incoming(body_stream, Some(data.len() as u64), u64::MAX);
738 let result = rt.block_on(body.read_all());
739 if let Ok(val) = result {
740 prop_assert_eq!(val.len(), data.len());
741 }
742 });
743 }
744
745 #[test]
746 fn premature_eof_detected() {
747 proptest::proptest!(|(data in prop::collection::vec(any::<u8>(), 0..100))| {
748 use futures_util::stream;
749
750 let rt = tokio::runtime::Runtime::new().unwrap();
751 let declared = (data.len() as u64) + 100;
752 let data_owned = data.clone();
753 let body_stream = stream::once(async move {
754 Ok::<_, IncomingError>(Bytes::from(data_owned))
755 });
756 let body = RequestBody::from_incoming(body_stream, Some(declared), u64::MAX);
757 let result = rt.block_on(body.read_all());
758 if let Err(e) = result {
759 prop_assert!(
760 matches!(e, RequestBodyError::PrematureEof { .. }),
761 "expected PrematureEof, got: {:?}",
762 e
763 );
764 }
765 });
766 }
767}