Skip to main content

leyline/core/
response.rs

1use std::fmt;
2use std::sync::{Arc, OnceLock};
3
4use bytes::Bytes;
5use http::{HeaderMap, HeaderName, HeaderValue, StatusCode};
6use url::Url;
7
8use crate::core::body_stream::BodyStream;
9use crate::core::error::{Error, Kind, Result};
10use crate::core::session::decompress::{
11    Decoder, content_codings, decompress_body, drain_stream_into_vec,
12};
13use crate::trace::masked;
14use crate::util::redact;
15
16pub(crate) enum ResponseBody {
17    Buffered(Vec<u8>),
18    Streaming(BodyStream),
19    Taken,
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23#[non_exhaustive]
24pub enum HttpVersion {
25    Http1_1,
26    Http2,
27    Http3,
28}
29
30impl HttpVersion {
31    pub(crate) fn ja4h_token(self) -> &'static str {
32        match self {
33            Self::Http1_1 => "1",
34            Self::Http2 => "2",
35            Self::Http3 => "3",
36        }
37    }
38
39    pub fn as_str(self) -> &'static str {
40        match self {
41            Self::Http1_1 => "HTTP/1.1",
42            Self::Http2 => "HTTP/2",
43            Self::Http3 => "HTTP/3",
44        }
45    }
46}
47
48#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
49#[non_exhaustive]
50pub struct ResponseTiming {
51    pub reused: bool,
52    pub connect_ms: Option<u32>,
53    pub send_ms: u32,
54    pub total_ms: u32,
55}
56
57impl ResponseTiming {
58    pub(crate) fn accumulator() -> Self {
59        Self {
60            reused: true,
61            connect_ms: None,
62            send_ms: 0,
63            total_ms: 0,
64        }
65    }
66
67    pub(crate) fn leg(started: std::time::Instant, connect_ms: Option<u32>) -> Self {
68        let total_ms = Self::millis(started);
69        Self {
70            reused: connect_ms.is_none(),
71            connect_ms,
72            send_ms: total_ms.saturating_sub(connect_ms.unwrap_or(0)),
73            total_ms,
74        }
75    }
76
77    pub(crate) fn millis(start: std::time::Instant) -> u32 {
78        u32::try_from(start.elapsed().as_millis()).unwrap_or(u32::MAX)
79    }
80
81    pub(crate) fn add_leg(&mut self, leg: &ResponseTiming) {
82        self.total_ms = self.total_ms.saturating_add(leg.total_ms);
83        self.send_ms = self.send_ms.saturating_add(leg.send_ms);
84        if let Some(c) = leg.connect_ms {
85            self.connect_ms = Some(self.connect_ms.unwrap_or(0).saturating_add(c));
86        }
87        self.reused &= leg.reused;
88    }
89}
90
91pub struct Response {
92    pub(crate) status: StatusCode,
93    pub(crate) version: HttpVersion,
94    pub(crate) headers: HeaderMap,
95    pub(crate) trailers: Vec<(HeaderName, HeaderValue)>,
96    pub(crate) body: ResponseBody,
97    pub(crate) url: Url,
98    pub(crate) redirect_chain: Vec<Url>,
99    pub(crate) request_headers: Vec<(String, String)>,
100    pub(crate) tls: Option<crate::pool::TlsInfo>,
101    pub(crate) request_method: String,
102    pub(crate) timing: ResponseTiming,
103    pub(crate) audit_tls: Option<Arc<crate::audit::AuditTlsCache>>,
104    pub(crate) audit_cache: OnceLock<crate::audit::AuditData>,
105    pub(crate) compression: crate::core::CompressionConfig,
106}
107
108impl Response {
109    pub fn status(&self) -> StatusCode {
110        self.status
111    }
112
113    pub fn version(&self) -> HttpVersion {
114        self.version
115    }
116
117    pub fn timing(&self) -> &ResponseTiming {
118        &self.timing
119    }
120
121    pub fn url(&self) -> &Url {
122        &self.url
123    }
124
125    pub fn redirect_chain(&self) -> &[Url] {
126        &self.redirect_chain
127    }
128
129    pub fn headers(&self) -> &HeaderMap {
130        &self.headers
131    }
132
133    pub fn trailers(&self) -> impl Iterator<Item = (&HeaderName, &HeaderValue)> {
134        self.trailers.iter().map(|(k, v)| (k, v))
135    }
136
137    pub fn cookies(&self) -> impl Iterator<Item = crate::cookie::Cookie> + '_ {
138        let url = &self.url;
139        self.headers
140            .get_all(http::header::SET_COOKIE)
141            .iter()
142            .filter_map(move |value| {
143                crate::cookie::parse::parse_set_cookie(value.to_str().ok()?, url)
144            })
145    }
146
147    pub fn tls(&self) -> Option<&crate::pool::TlsInfo> {
148        self.tls.as_ref()
149    }
150
151    pub fn request_headers(&self) -> impl Iterator<Item = (&str, &str)> {
152        self.request_headers
153            .iter()
154            .map(|(k, v)| (k.as_str(), v.as_str()))
155    }
156
157    pub async fn text(self) -> crate::core::Result<String> {
158        self.text_with_charset("utf-8").await
159    }
160
161    #[cfg(feature = "charset")]
162    pub async fn text_with_charset(self, default_encoding: &str) -> crate::core::Result<String> {
163        let label = self.charset_label().map(str::to_owned);
164        let bytes = self.bytes().await?;
165        let encoding = encoding_rs::Encoding::for_label(
166            label.as_deref().unwrap_or(default_encoding).as_bytes(),
167        )
168        .unwrap_or(encoding_rs::UTF_8);
169        Ok(encoding.decode(&bytes).0.into_owned())
170    }
171
172    #[cfg(not(feature = "charset"))]
173    pub async fn text_with_charset(self, _default_encoding: &str) -> crate::core::Result<String> {
174        Ok(String::from_utf8_lossy(&self.bytes().await?).into_owned())
175    }
176
177    #[cfg(feature = "charset")]
178    fn charset_label(&self) -> Option<&str> {
179        let ct = self.header("content-type")?;
180        ct.split(';').skip(1).find_map(|param| {
181            let (k, v) = param.split_once('=')?;
182            k.trim()
183                .eq_ignore_ascii_case("charset")
184                .then(|| v.trim().trim_matches('"'))
185        })
186    }
187
188    pub async fn bytes(self) -> crate::core::Result<Bytes> {
189        let stream = match self.body {
190            ResponseBody::Buffered(b) => return Ok(Bytes::from(b)),
191            ResponseBody::Taken => {
192                return Err(Error::new(Kind::Body).with_message(
193                    "response body stream was taken by `into_stream`; read the bytes from that stream",
194                ));
195            }
196            ResponseBody::Streaming(s) => s,
197        };
198        let buf = drain_stream_into_vec(stream, self.compression.max_body_size).await?;
199        let encoding = content_codings(self.headers.get_all(http::header::CONTENT_ENCODING));
200        let (buf, _) = decompress_body(buf, encoding.as_deref(), &self.compression)?;
201        Ok(Bytes::from(buf))
202    }
203
204    pub async fn json<T: serde::de::DeserializeOwned>(self) -> crate::core::Result<T> {
205        serde_json::from_slice(&self.bytes().await?).map_err(Error::from_json)
206    }
207
208    pub fn into_stream(mut self) -> Result<BodyStream> {
209        match std::mem::replace(&mut self.body, ResponseBody::Taken) {
210            ResponseBody::Streaming(s) => Ok(s),
211            ResponseBody::Buffered(b) => Ok(BodyStream::from_bytes(bytes::Bytes::from(b))),
212            ResponseBody::Taken => Err(Error::new(Kind::Body)
213                .with_message("response body has already been taken as a stream")),
214        }
215    }
216
217    pub async fn copy_to<W>(self, writer: &mut W) -> Result<u64>
218    where
219        W: tokio::io::AsyncWrite + Unpin,
220    {
221        use futures_util::StreamExt;
222        use tokio::io::AsyncWriteExt;
223
224        let mut stream = self.into_stream()?;
225        let mut total: u64 = 0;
226        while let Some(chunk) = stream.next().await {
227            let chunk = chunk.map_err(Error::from)?;
228            writer.write_all(&chunk).await.map_err(Error::from)?;
229            total += chunk.len() as u64;
230        }
231        writer.flush().await.map_err(Error::from)?;
232        Ok(total)
233    }
234
235    pub fn content_length(&self) -> Option<u64> {
236        self.header("content-length").and_then(|v| v.parse().ok())
237    }
238
239    pub fn error_for_status(self) -> crate::core::Result<Self> {
240        match self.status_error() {
241            Some(err) => Err(err),
242            None => Ok(self),
243        }
244    }
245
246    pub fn error_for_status_ref(&self) -> crate::core::Result<&Self> {
247        match self.status_error() {
248            Some(err) => Err(err),
249            None => Ok(self),
250        }
251    }
252
253    fn status_error(&self) -> Option<Error> {
254        if self.status.as_u16() < 400 {
255            return None;
256        }
257        Some(
258            Error::new(Kind::Status)
259                .with_status(self.status)
260                .with_url(self.url.clone()),
261        )
262    }
263
264    pub fn audit(&self) -> Option<&crate::audit::AuditData> {
265        let tls = self.audit_tls.as_ref()?;
266        Some(self.audit_cache.get_or_init(|| {
267            let ja4h = crate::audit::compute_ja4h(&crate::audit::Ja4hInput {
268                method: &self.request_method,
269                http_version: self.version.ja4h_token(),
270                headers: &self.request_headers,
271            });
272            crate::audit::AuditData {
273                ja4: tls.ja4.clone(),
274                ja3: tls.ja3.clone(),
275                h2_fingerprint: tls.h2_fingerprint.clone(),
276                ja4t: tls.ja4t.clone(),
277                ja4h,
278            }
279        }))
280    }
281
282    pub fn header(&self, name: &str) -> Option<&str> {
283        self.headers.get(name).and_then(|v| v.to_str().ok())
284    }
285
286    pub async fn read_until<F>(self, limit: usize, mut done: F) -> Result<Vec<u8>>
287    where
288        F: FnMut(&[u8], usize) -> bool,
289    {
290        use futures_util::StreamExt;
291
292        let encoding = content_codings(self.headers.get_all(http::header::CONTENT_ENCODING));
293        let mut decoder = Decoder::new(encoding.as_deref(), &self.compression)?;
294        let mut stream = self.into_stream()?;
295        let mut out = Vec::new();
296        while let Some(chunk) = stream.next().await {
297            let chunk = chunk.map_err(Error::from)?;
298            let from = out.len();
299            match decoder.as_mut() {
300                Some(decoder) => decoder.feed(&chunk, &mut out)?,
301                None => out.extend_from_slice(&chunk),
302            }
303            out.truncate(limit);
304            if done(&out, from) || out.len() >= limit {
305                return Ok(out);
306            }
307        }
308        if let Some(decoder) = decoder.as_mut() {
309            decoder.finish(&mut out)?;
310            out.truncate(limit);
311        }
312        Ok(out)
313    }
314}
315
316impl fmt::Debug for Response {
317    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
318        let chain: Vec<String> = self
319            .redirect_chain
320            .iter()
321            .map(|url| redact(url.as_str()))
322            .collect();
323        f.debug_struct("Response")
324            .field("status", &self.status)
325            .field("version", &self.version)
326            .field("url", &redact(self.url.as_str()))
327            .field("redirect_chain", &chain)
328            .field(
329                "headers",
330                &masked(self.headers.iter().map(|(k, v)| (k.as_str(), v.as_bytes()))),
331            )
332            .field(
333                "trailers",
334                &masked(
335                    self.trailers
336                        .iter()
337                        .map(|(k, v)| (k.as_str(), v.as_bytes())),
338                ),
339            )
340            .field(
341                "request_headers",
342                &masked(
343                    self.request_headers
344                        .iter()
345                        .map(|(k, v)| (k.as_str(), v.as_bytes())),
346                ),
347            )
348            .field("timing", &self.timing)
349            .finish_non_exhaustive()
350    }
351}
352
353#[cfg(test)]
354mod tests;