async_fetch/
response.rs

1use std::fmt;
2use std::pin::Pin;
3use std::collections::HashMap;
4use std::collections::hash_map::RandomState;
5use std::io::{Error, ErrorKind};
6use std::str::FromStr;
7use async_std::io::{Read};
8use async_httplib::{Status, Version, read_exact, read_chunks};
9use crate::{read_content_length, read_transfer_encoding};
10
11pub struct Response<'a> {
12    status: Status,
13    version: Version,
14    headers: HashMap<String, String>,
15    reader: Pin<Box<dyn Read + Send + Unpin + 'a>>,
16    chunkline_limit: Option<usize>,
17    body_limit: Option<usize>,
18}
19
20impl<'a> Response<'a> {
21
22    pub fn default() -> Self {
23        Self {
24            status: Status::Ok,
25            version: Version::Http1_1,
26            headers: HashMap::with_hasher(RandomState::new()),
27            reader: Box::pin("".as_bytes()),
28            chunkline_limit: None,
29            body_limit: None,
30        }
31    }
32
33    pub fn with_reader<R>(reader: R) -> Self
34        where
35        R: Read + Send + Unpin + 'a,
36    {
37        let mut res = Self::default();
38        res.set_reader(reader);
39        res
40    }
41
42    pub fn status(&self) -> &Status {
43        &self.status
44    }
45
46    pub fn version(&self) -> &Version {
47        &self.version
48    }
49
50    pub fn headers(&self) -> &HashMap<String, String> {
51        &self.headers
52    }
53
54    pub fn header<N: Into<String>>(&self, name: N) -> Option<&String> {
55        self.headers.get(&name.into())
56    }
57
58    pub fn reader(&self) -> &Pin<Box<dyn Read + Send + Unpin + 'a>> {
59        &self.reader
60    }
61
62    pub fn chunkline_limit(&self) -> &Option<usize> {
63        &self.chunkline_limit
64    }
65
66    pub fn body_limit(&self) -> &Option<usize> {
67        &self.body_limit
68    }
69
70    pub fn has_status(&self, value: Status) -> bool {
71        self.status == value
72    }
73
74    pub fn has_version(&self, value: Version) -> bool {
75        self.version == value
76    }
77
78    pub fn has_headers(&self) -> bool {
79        !self.headers.is_empty()
80    }
81
82    pub fn has_header<N: Into<String>>(&self, name: N) -> bool {
83        self.headers.contains_key(&name.into())
84    }
85
86    pub fn has_chunkline_limit(&self) -> bool {
87        self.chunkline_limit.is_some()
88    }
89
90    pub fn has_body_limit(&self) -> bool {
91        self.body_limit.is_some()
92    }
93
94    pub fn set_status(&mut self, value: Status) {
95        self.status = value;
96    }
97
98    pub fn set_status_str(&mut self, value: &str) -> Result<(), Error> {
99        self.status = Status::from_str(value)?;
100        Ok(())
101    }
102
103    pub fn set_version(&mut self, value: Version) {
104        self.version = value;
105    }
106
107    pub fn set_version_str(&mut self, value: &str) -> Result<(), Error> {
108        self.version = Version::from_str(value)?;
109        Ok(())
110    }
111
112    pub fn set_header<N: Into<String>, V: Into<String>>(&mut self, name: N, value: V) {
113        self.headers.insert(name.into(), value.into());
114    }
115
116    pub fn set_reader<R>(&mut self, reader: R)
117        where
118        R: Read + Send + Unpin + 'a,
119    {
120        self.reader = Box::pin(reader);
121    }
122
123    pub fn set_chunkline_limit(&mut self, length: usize) {
124        self.chunkline_limit = Some(length);
125    }
126
127    pub fn set_body_limit(&mut self, length: usize) {
128        self.body_limit = Some(length);
129    }
130
131    pub fn remove_header<N: Into<String>>(&mut self, name: N) {
132        self.headers.remove(&name.into());
133    }
134
135    pub fn clear_headers(&mut self) {
136        self.headers.clear();
137    }
138
139    pub fn to_proto_string(&self) -> String {
140        let mut output = String::new();
141        if !self.has_version(Version::Http0_9) {
142            output.push_str(&format!("{} {} {}\r\n", self.version, self.status, self.status.reason()));
143
144            for (name, value) in self.headers.iter() {
145                output.push_str(&format!("{}: {}\r\n", name, value));
146            }
147
148            output.push_str("\r\n");
149        }
150        output
151    }
152
153    pub async fn recv(&mut self) -> Result<Vec<u8>, Error> {
154        let mut data = Vec::new();
155
156        if read_transfer_encoding(&self.headers) == "chunked" {
157            read_chunks(&mut self.reader, &mut data, (self.chunkline_limit, self.body_limit)).await?;
158        } else if self.has_header("Content-Length") {
159            let length = read_content_length(&self.headers, self.body_limit)?;
160            read_exact(&mut self.reader, &mut data, length).await?;
161        }
162
163        Ok(data)
164    }
165
166    pub async fn recv_string(&mut self) -> Result<String, Error> {
167        let data = self.recv().await?;
168        let txt = match String::from_utf8(data) {
169            Ok(txt) => txt,
170            Err(e) => return Err(Error::new(ErrorKind::InvalidData, e.to_string())),
171        };
172        Ok(txt)
173    }
174
175    #[cfg(feature = "json")]
176    pub async fn recv_json(&mut self) -> Result<serde_json::Value, Error> {
177        let mut data = self.recv().await?;
178        if data.is_empty() {
179            data = "{}".as_bytes().to_vec();
180        }
181        let json: serde_json::Value = match serde_json::from_slice(&data) {
182            Ok(json) => json,
183            Err(e) => return Err(Error::new(ErrorKind::InvalidData, e.to_string())),
184        };
185        Ok(json)
186    }
187}
188
189impl fmt::Display for Response<'_> {
190    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
191        write!(fmt, "{}", self.to_proto_string())
192    }
193}
194
195impl From<Response<'_>> for String {
196    fn from(item: Response) -> String {
197        item.to_proto_string()
198    }
199}