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}