async_httype/
body.rs

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}