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