1use bytes::Bytes;
2use futures_util::{Stream, StreamExt};
3use tokio::io::{AsyncWrite, AsyncWriteExt};
4
5use super::prefix::stop_reason;
6use super::{PrefixRead, Response, ResponseBody, StopReason};
7use crate::core::body_stream::BodyStream;
8use crate::core::deadline::{Deadline, within};
9use crate::core::error::{Error, Kind, Result};
10use crate::core::session::decompress::{
11 BodyLimit, Decoder, body_read_error, content_codings, decompress_body, drain_stream_into_vec,
12};
13
14impl Response {
15 pub async fn text(self) -> crate::core::Result<String> {
16 self.text_with_charset("utf-8").await
17 }
18
19 #[cfg(feature = "charset")]
20 pub async fn text_with_charset(self, default_encoding: &str) -> crate::core::Result<String> {
21 let label = self.charset_label().map(str::to_owned);
22 let bytes = self.bytes().await?;
23 let encoding = encoding_rs::Encoding::for_label(
24 label.as_deref().unwrap_or(default_encoding).as_bytes(),
25 )
26 .unwrap_or(encoding_rs::UTF_8);
27 Ok(encoding.decode(&bytes).0.into_owned())
28 }
29
30 #[cfg(not(feature = "charset"))]
31 pub async fn text_with_charset(self, _default_encoding: &str) -> crate::core::Result<String> {
32 Ok(String::from_utf8_lossy(&self.bytes().await?).into_owned())
33 }
34
35 #[cfg(any(feature = "charset", feature = "html"))]
36 fn content_type(&self) -> Option<(&str, &str)> {
37 let ct = self.header("content-type")?;
38 let (media, params) = ct.split_once(';').unwrap_or((ct, ""));
39 Some((media.trim(), params))
40 }
41
42 #[cfg(feature = "html")]
43 pub(crate) fn media_type(&self) -> Option<&str> {
44 self.content_type().map(|(media, _)| media)
45 }
46
47 #[cfg(feature = "charset")]
48 fn charset_label(&self) -> Option<&str> {
49 let (_, params) = self.content_type()?;
50 params.split(';').find_map(|param| {
51 let (k, v) = param.split_once('=')?;
52 k.trim()
53 .eq_ignore_ascii_case("charset")
54 .then(|| v.trim().trim_matches('"'))
55 })
56 }
57
58 pub async fn bytes(self) -> crate::core::Result<Bytes> {
59 let stream = match self.body {
60 ResponseBody::Buffered(b) => return Ok(Bytes::from(b)),
61 ResponseBody::Taken => {
62 return Err(Error::new(Kind::Body).with_message(
63 "response body stream was taken by `into_stream`; read the bytes from that stream",
64 ));
65 }
66 ResponseBody::Streaming(s) => s,
67 };
68 let buf = drain_stream_into_vec(stream, self.compression.max_body_size).await?;
69 let encoding = content_codings(self.headers.get_all(http::header::CONTENT_ENCODING));
70 let (buf, _) = decompress_body(buf, encoding.as_deref(), &self.compression)?;
71 Ok(Bytes::from(buf))
72 }
73
74 pub async fn json<T: serde::de::DeserializeOwned>(self) -> crate::core::Result<T> {
75 serde_json::from_slice(&self.bytes().await?).map_err(Error::from_json)
76 }
77
78 pub fn into_stream(mut self) -> Result<BodyStream> {
79 match std::mem::replace(&mut self.body, ResponseBody::Taken) {
80 ResponseBody::Streaming(s) => Ok(s),
81 ResponseBody::Buffered(b) => Ok(BodyStream::from_bytes(bytes::Bytes::from(b))),
82 ResponseBody::Taken => Err(Error::new(Kind::Body)
83 .with_message("response body has already been taken as a stream")),
84 }
85 }
86
87 pub async fn copy_to<W>(self, writer: &mut W) -> Result<u64>
88 where
89 W: tokio::io::AsyncWrite + Unpin,
90 {
91 write_stream(self.into_stream()?, writer, Error::from).await
92 }
93
94 pub fn into_decoded_stream(self, limit: Option<u64>) -> Result<BodyStream> {
95 let encoding = content_codings(self.headers.get_all(http::header::CONTENT_ENCODING));
96 if encoding.is_none() {
97 self.refuse_declared_overflow(limit)?;
98 }
99 let decoder = Decoder::capped(encoding.as_deref(), &self.compression, limit)?;
100 Ok(BodyStream::decoded(self.into_stream()?, decoder))
101 }
102
103 fn refuse_declared_overflow(&self, limit: Option<u64>) -> Result<()> {
104 let cap = BodyLimit::session(self.compression.max_body_size).tighter(limit);
105 match self.content_length() {
106 Some(declared) if declared > u64::try_from(cap.bytes).unwrap_or(u64::MAX) => {
107 Err(cap.error())
108 }
109 _ => Ok(()),
110 }
111 }
112
113 pub async fn copy_decoded_to<W>(self, writer: &mut W, limit: Option<u64>) -> Result<u64>
114 where
115 W: AsyncWrite + Unpin,
116 {
117 write_stream(self.into_decoded_stream(limit)?, writer, body_read_error).await
118 }
119
120 pub(crate) async fn error_for_status_with_body(self, deadline: &Deadline) -> Result<Self> {
121 let Some(err) = self.status_error() else {
122 return Ok(self);
123 };
124 let limit = self.compression.max_error_body;
125 let declared = self.content_length();
126 let mut body = Vec::new();
127 let mut decodes = false;
128 let read = within(
129 Some(deadline.error_body_wait()),
130 self.read_prefix_into(limit, |_, _| false, &mut body, &mut decodes),
131 )
132 .await;
133 if !matches!(read, Ok(Ok(_))) && body.is_empty() {
134 return Err(err);
135 }
136 let complete = declared == Some(body.len() as u64);
137 let err = err.with_body(body);
138 Err(match (decodes, complete) {
139 (true, _) => err.without_content_coding(),
140 (false, false) => err.without_content_length(),
141 (false, true) => err,
142 })
143 }
144
145 pub async fn read_until<F>(self, limit: usize, done: F) -> Result<Vec<u8>>
146 where
147 F: FnMut(&[u8], usize) -> bool,
148 {
149 Ok(self.read_prefix(limit, done).await?.bytes)
150 }
151
152 pub(crate) async fn read_prefix<F>(self, limit: usize, done: F) -> Result<PrefixRead>
153 where
154 F: FnMut(&[u8], usize) -> bool,
155 {
156 let mut bytes = Vec::new();
157 let stopped_by = self
158 .read_prefix_into(limit, done, &mut bytes, &mut false)
159 .await?;
160 Ok(PrefixRead { bytes, stopped_by })
161 }
162
163 async fn read_prefix_into<F>(
164 self,
165 limit: usize,
166 mut done: F,
167 out: &mut Vec<u8>,
168 decodes: &mut bool,
169 ) -> Result<StopReason>
170 where
171 F: FnMut(&[u8], usize) -> bool,
172 {
173 let limit = limit.min(self.compression.max_body_size);
174 if limit == 0 {
175 return Ok(StopReason::LimitReached);
176 }
177 let encoding = content_codings(self.headers.get_all(http::header::CONTENT_ENCODING));
178 let mut decoder = Decoder::truncated(encoding.as_deref(), &self.compression, limit)?;
179 *decodes = decoder.is_some();
180 let mut stream = self.into_stream()?;
181 while let Some(chunk) = stream.next().await {
182 let chunk = chunk.map_err(Error::from)?;
183 let from = out.len();
184 match decoder.as_mut() {
185 Some(decoder) => decoder.feed(&chunk, out)?,
186 None => out.extend_from_slice(&chunk),
187 }
188 out.truncate(limit);
189 if let Some(reason) = stop_reason(out, from, limit, &mut done) {
190 return Ok(reason);
191 }
192 }
193 let from = out.len();
194 if let Some(decoder) = decoder.as_mut() {
195 decoder.finish(out)?;
196 out.truncate(limit);
197 }
198 Ok(stop_reason(out, from, limit, &mut done).unwrap_or(StopReason::EndOfBody))
199 }
200}
201
202async fn write_stream<S, W>(
203 mut stream: S,
204 writer: &mut W,
205 map: fn(std::io::Error) -> Error,
206) -> Result<u64>
207where
208 S: Stream<Item = std::io::Result<Bytes>> + Unpin,
209 W: AsyncWrite + Unpin,
210{
211 let mut total: u64 = 0;
212 while let Some(chunk) = stream.next().await {
213 let chunk = chunk.map_err(map)?;
214 writer.write_all(&chunk).await.map_err(Error::from)?;
215 total += chunk.len() as u64;
216 }
217 writer.flush().await.map_err(Error::from)?;
218 Ok(total)
219}