1use reqwest::{Response, StatusCode};
7
8pub const USER_AGENT: &str = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0 Safari/537.36";
15
16pub const MAX_BODY_BYTES: usize = 5 * 1024 * 1024;
25
26fn push_capped(buf: &mut Vec<u8>, chunk: &[u8], max: usize) -> bool {
29 let remaining = max.saturating_sub(buf.len());
30 if chunk.len() >= remaining {
31 buf.extend_from_slice(&chunk[..remaining]);
32 true
33 } else {
34 buf.extend_from_slice(chunk);
35 false
36 }
37}
38
39pub async fn read_body_capped(resp: Response) -> Result<String, (anyhow::Error, bool)> {
46 let bytes = read_body_capped_bytes(resp).await?;
47 Ok(String::from_utf8_lossy(&bytes).into_owned())
48}
49
50pub async fn read_body_capped_bytes(mut resp: Response) -> Result<Vec<u8>, (anyhow::Error, bool)> {
53 let mut buf: Vec<u8> = Vec::new();
54 if let Some(len) = resp.content_length() {
56 buf.reserve(len.min(MAX_BODY_BYTES as u64) as usize);
57 }
58 loop {
59 match resp.chunk().await {
60 Ok(Some(chunk)) => {
61 if push_capped(&mut buf, &chunk, MAX_BODY_BYTES) {
62 break;
63 }
64 }
65 Ok(None) => break,
66 Err(e) => {
67 let transient = e.is_timeout();
68 return Err((e.into(), transient));
69 }
70 }
71 }
72 Ok(buf)
73}
74
75pub fn transient_send_error(e: &reqwest::Error) -> bool {
77 e.is_timeout() || e.is_connect() || e.is_request()
78}
79
80pub fn transient_status(status: StatusCode) -> bool {
82 status.is_server_error() || status.as_u16() == 429
83}
84
85#[cfg(test)]
86mod tests {
87 use super::*;
88
89 #[test]
90 fn push_capped_truncates_oversized_chunk() {
91 let mut buf = Vec::new();
92 let stopped = push_capped(&mut buf, &[b'x'; 10], 4);
94 assert!(stopped);
95 assert_eq!(buf.len(), 4);
96 }
97
98 #[test]
99 fn push_capped_accumulates_until_cap() {
100 let mut buf = Vec::new();
101 assert!(!push_capped(&mut buf, b"abc", 8));
102 assert!(!push_capped(&mut buf, b"de", 8));
103 assert_eq!(buf, b"abcde");
104 let stopped = push_capped(&mut buf, b"fghij", 8);
106 assert!(stopped);
107 assert_eq!(buf.len(), 8);
108 assert_eq!(buf, b"abcdefgh");
109 }
110
111 #[test]
112 fn push_capped_small_body_unaffected() {
113 let mut buf = Vec::new();
114 let stopped = push_capped(&mut buf, b"hello", 1024);
115 assert!(!stopped);
116 assert_eq!(buf, b"hello");
117 }
118
119 #[test]
120 fn transient_status_covers_5xx_and_429() {
121 assert!(transient_status(StatusCode::INTERNAL_SERVER_ERROR));
122 assert!(transient_status(StatusCode::TOO_MANY_REQUESTS));
123 assert!(!transient_status(StatusCode::NOT_FOUND));
124 assert!(!transient_status(StatusCode::FORBIDDEN));
125 }
126}