blooym.dev
[MIRROR ONLY] A correct and efficient ATProto blob proxy for secure content delivery.
codeberg.org/Blooym/porxie
1use bytes::{Bytes, BytesMut};
2use futures::StreamExt;
3use reqwest::redirect::Policy;
4use std::{num::NonZeroU64, time::Duration};
5use thiserror::Error;
6
7#[inline]
8pub fn build_http_client(
9 timeout: Duration,
10 connect_timeout: Duration,
11 https_only: bool,
12) -> Result<reqwest::Client, reqwest::Error> {
13 reqwest::Client::builder()
14 .user_agent(concat!(
15 env!("CARGO_PKG_NAME"),
16 "/",
17 env!("CARGO_PKG_VERSION_MAJOR"),
18 ".",
19 env!("CARGO_PKG_VERSION_MINOR"),
20 " (",
21 env!("CARGO_PKG_REPOSITORY"),
22 ")"
23 ))
24 .https_only(https_only)
25 .redirect(Policy::limited(3))
26 .gzip(true)
27 .brotli(true)
28 .zstd(true)
29 .deflate(true)
30 .connect_timeout(connect_timeout)
31 .timeout(timeout)
32 .build()
33}
34
35#[derive(Debug, Error)]
36pub enum BytesStreamCappedError {
37 /// The response content length exceeded the size limit.
38 #[error("content exceeded the maximum size")]
39 TooLarge,
40 /// An internal client error occurred whilst processing the request,
41 /// see [`reqwest::Error`].
42 #[error(transparent)]
43 ClientError(#[from] reqwest::Error),
44}
45
46/// A wrapper around `Response::bytes_stream()` that acts like `Response::bytes()`
47/// but enforces a maximum size limit while streaming the response.
48pub async fn bytes_stream_capped(
49 response: reqwest::Response,
50 max_size: NonZeroU64,
51) -> Result<Bytes, BytesStreamCappedError> {
52 if let Some(content_length) = response.content_length()
53 && content_length > max_size.get()
54 {
55 return Err(BytesStreamCappedError::TooLarge);
56 }
57
58 let mut buffer = BytesMut::with_capacity(
59 response
60 .content_length()
61 .unwrap_or(64 * 1024)
62 .min(max_size.get())
63 .try_into()
64 .unwrap_or(usize::MAX),
65 );
66 let mut stream = response.bytes_stream();
67 while let Some(chunk) = stream.next().await {
68 let chunk = chunk.map_err(BytesStreamCappedError::ClientError)?;
69 if buffer.len() as u64 + chunk.len() as u64 > max_size.get() {
70 return Err(BytesStreamCappedError::TooLarge);
71 }
72 buffer.extend_from_slice(&chunk);
73 }
74
75 Ok(buffer.freeze())
76}