1use std::fmt;
2use std::collections::HashMap;
3use async_std::io::{Read, Write};
4use crate::{Error, read_chunked_stream, read_sized_stream, write_to_stream, flush_stream};
5
6pub struct Body {
7 bytes: Vec<u8>,
8 length: usize,
9 length_limit: Option<usize>,
10}
11
12impl Body {
13
14 pub fn new() -> Self {
15 Self {
16 bytes: Vec::new(),
17 length: 0,
18 length_limit: None,
19 }
20 }
21
22 pub fn bytes(&self) -> &Vec<u8> {
23 &self.bytes
24 }
25
26 pub fn length(&self) -> usize {
27 self.length
28 }
29
30 pub fn length_limit(&self) -> Option<usize> {
31 self.length_limit
32 }
33
34 pub fn has_length_limit(&self) -> bool {
35 self.length_limit.is_some()
36 }
37
38 pub fn set_length_limit(&mut self, limit: usize) {
39 self.length_limit = Some(limit);
40 }
41
42 pub fn remove_length_limit(&mut self) {
43 self.length_limit = None;
44 }
45
46 pub async fn read<I>(&mut self, stream: &mut I, res: &HashMap<String, String>) -> Result<usize, Error>
47 where
48 I: Read + Unpin,
49 {
50 let length = res.get("Content-Length");
51 let encoding = res.get("Transfer-Encoding");
52
53 if encoding.is_some() && encoding.unwrap().contains(&String::from("chunked")) {
54 self.read_chunked(stream).await
55 } else {
56 let length = match length {
57 Some(length) => match length.parse::<usize>() {
58 Ok(length) => length,
59 Err(_) => return Err(Error::InvalidHeader(String::from("Content-Length"))),
60 },
61 None => return Err(Error::InvalidHeader(String::from("Content-Length"))),
62 };
63 self.read_sized(stream, length).await
64 }
65 }
66
67 pub async fn read_chunked<I>(&mut self, stream: &mut I) -> Result<usize, Error>
68 where
69 I: Read + Unpin,
70 {
71 let limit = match self.length_limit {
72 Some(limit) => match limit == 0 {
73 true => return Err(Error::SizeLimitExceeded(limit)),
74 false => Some(limit - self.length),
75 },
76 None => None,
77 };
78
79 let length = read_chunked_stream(stream, &mut self.bytes, limit).await?;
80 self.length += length;
81
82 Ok(length)
83 }
84
85 pub async fn read_sized<I>(&mut self, stream: &mut I, length: usize) -> Result<usize, Error>
86 where
87 I: Read + Unpin,
88 {
89 match self.length_limit {
90 Some(limit) => match length + self.length > limit {
91 true => return Err(Error::SizeLimitExceeded(limit)),
92 false => (),
93 },
94 None => (),
95 };
96
97 let length = read_sized_stream(stream, &mut self.bytes, length).await?;
98 self.length += length;
99
100 Ok(length)
101 }
102
103 pub async fn write<I>(&mut self, stream: &mut I) -> Result<usize, Error>
104 where
105 I: Write + Unpin,
106 {
107 let size = write_to_stream(stream, &self.bytes()).await?;
108 flush_stream(stream).await?;
109 Ok(size)
110 }
111
112 pub fn clear(&mut self) {
113 self.bytes.clear();
114 self.length = 0;
115 self.length_limit = None;
116 }
117}
118
119impl fmt::Display for Body {
120 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
121 write!(fmt, "{:?}", self.bytes())
122 }
123}