aria2_protocol/http/
client.rs

1use std::time::Duration;
2
3use reqwest::{Certificate, Client, ClientBuilder, redirect};
4use tracing::{debug, info};
5
6use crate::http::request::HttpRequest;
7use crate::http::response::HttpResponse;
8
9#[derive(Debug, Clone)]
10pub struct HttpClientOptions {
11    pub connect_timeout: Duration,
12    pub timeout: Duration,
13    pub max_redirects: usize,
14    pub user_agent: String,
15    pub accept_gzip: bool,
16    pub verify_tls: bool,
17    pub ca_cert_path: Option<String>,
18}
19
20impl Default for HttpClientOptions {
21    fn default() -> Self {
22        Self {
23            connect_timeout: Duration::from_secs(30),
24            timeout: Duration::from_secs(300),
25            max_redirects: 5,
26            user_agent: "aria2/1.37.0-Rust".to_string(),
27            accept_gzip: true,
28            verify_tls: true,
29            ca_cert_path: None,
30        }
31    }
32}
33
34pub struct HttpClient {
35    inner: Client,
36    options: HttpClientOptions,
37}
38
39impl HttpClient {
40    pub fn new(options: HttpClientOptions) -> Result<Self, String> {
41        let mut builder = ClientBuilder::new()
42            .connect_timeout(options.connect_timeout)
43            .timeout(options.timeout)
44            .user_agent(&options.user_agent)
45            .redirect(redirect::Policy::limited(options.max_redirects));
46
47        if options.accept_gzip {
48            builder = builder.gzip(true);
49        }
50
51        if !options.verify_tls {
52            builder = builder.danger_accept_invalid_certs(true);
53        }
54
55        if let Some(ref ca_path) = options.ca_cert_path {
56            match std::fs::read(ca_path) {
57                Ok(cert_bytes) => {
58                    let cert = Certificate::from_pem(&cert_bytes)
59                        .map_err(|e| format!("Failed to load CA certificate: {}", e))?;
60                    builder = builder.add_root_certificate(cert);
61                }
62                Err(e) => return Err(format!("Failed to read CA certificate file: {}", e)),
63            }
64        }
65
66        let inner = builder
67            .build()
68            .map_err(|e| format!("Failed to create HTTP client: {}", e))?;
69
70        info!(
71            "HttpClient initialized (timeout={:?}, max_redirects={}, verify_tls={})",
72            options.timeout, options.max_redirects, options.verify_tls
73        );
74
75        Ok(Self { inner, options })
76    }
77
78    pub fn default_client() -> Result<Self, String> {
79        Self::new(HttpClientOptions::default())
80    }
81
82    pub async fn execute(&self, request: HttpRequest) -> Result<HttpResponse, String> {
83        debug!("Sending HTTP request: {} {}", request.method, request.url);
84
85        let mut reqwest_request = match request.method.to_uppercase().as_str() {
86            "GET" => self.inner.get(&request.url),
87            "POST" => self.inner.post(&request.url),
88            "HEAD" => self.inner.head(&request.url),
89            "PUT" => self.inner.put(&request.url),
90            _ => return Err(format!("Unsupported HTTP method: {}", request.method)),
91        };
92
93        if let Some(ref headers) = request.headers {
94            for (key, value) in headers.iter() {
95                reqwest_request = reqwest_request.header(
96                    key.as_str()
97                        .parse::<reqwest::header::HeaderName>()
98                        .map_err(|e| format!("Invalid header name: {}", e))?,
99                    value.as_str(),
100                );
101            }
102        }
103
104        if let Some(body) = request.body {
105            reqwest_request = reqwest_request.body(body);
106        }
107
108        let response = reqwest_request
109            .send()
110            .await
111            .map_err(|e| format!("HTTP request failed: {}", e))?;
112
113        let status = response.status();
114        let status_code = status.as_u16();
115        debug!("Received HTTP response: status_code={}", status_code);
116
117        let headers_map: Vec<(String, String)> = response
118            .headers()
119            .iter()
120            .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap_or("").to_string()))
121            .collect();
122
123        let body_bytes = response
124            .bytes()
125            .await
126            .map_err(|e| format!("Failed to read response body: {}", e))?
127            .to_vec();
128
129        Ok(HttpResponse {
130            status_code,
131            status_text: status.canonical_reason().unwrap_or("Unknown").to_string(),
132            headers: headers_map,
133            body: body_bytes,
134        })
135    }
136
137    pub fn get<U: Into<String>>(&self, url: U) -> HttpRequestBuilder<'_> {
138        HttpRequestBuilder::new(self, "GET", url.into())
139    }
140
141    pub fn post<U: Into<String>>(&self, url: U) -> HttpRequestBuilder<'_> {
142        HttpRequestBuilder::new(self, "POST", url.into())
143    }
144
145    pub fn head<U: Into<String>>(&self, url: U) -> HttpRequestBuilder<'_> {
146        HttpRequestBuilder::new(self, "HEAD", url.into())
147    }
148
149    pub fn options_ref(&self) -> &HttpClientOptions {
150        &self.options
151    }
152}
153
154pub struct HttpRequestBuilder<'a> {
155    client: &'a HttpClient,
156    method: String,
157    url: String,
158    headers: Option<reqwest::header::HeaderMap>,
159    body: Option<Vec<u8>>,
160}
161
162impl<'a> HttpRequestBuilder<'a> {
163    fn new(client: &'a HttpClient, method: &str, url: String) -> Self {
164        Self {
165            client,
166            method: method.to_string(),
167            url,
168            headers: None,
169            body: None,
170        }
171    }
172
173    pub fn header(mut self, name: &str, value: &str) -> Self {
174        let mut headers = self.headers.take().unwrap_or_default();
175        headers.insert(
176            name.parse::<reqwest::header::HeaderName>()
177                .expect("Invalid header name"),
178            value
179                .parse::<reqwest::header::HeaderValue>()
180                .expect("Invalid header value"),
181        );
182        self.headers = Some(headers);
183        self
184    }
185
186    pub fn header_raw<K, V>(mut self, key: K, value: V) -> Self
187    where
188        K: TryInto<reqwest::header::HeaderName>,
189        V: TryInto<reqwest::header::HeaderValue>,
190    {
191        let mut headers = self.headers.take().unwrap_or_default();
192        if let (Ok(k), Ok(v)) = (key.try_into(), value.try_into()) {
193            headers.insert(k, v);
194        }
195        self.headers = Some(headers);
196        self
197    }
198
199    pub fn range(self, start: u64, end: Option<u64>) -> Self {
200        let range_value = match end {
201            Some(e) => format!("bytes={}-{}", start, e),
202            None => format!("bytes={}-", start),
203        };
204        self.header("Range", &range_value)
205    }
206
207    pub fn body<B: Into<Vec<u8>>>(mut self, body: B) -> Self {
208        self.body = Some(body.into());
209        self
210    }
211
212    pub fn user_agent(self, ua: &str) -> Self {
213        self.header("User-Agent", ua)
214    }
215
216    pub fn referer(self, referer: &str) -> Self {
217        self.header("Referer", referer)
218    }
219
220    pub async fn send(self) -> Result<HttpResponse, String> {
221        let headers_map = self.headers.map(|h| {
222            h.iter()
223                .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap_or("").to_string()))
224                .collect::<Vec<_>>()
225        });
226
227        let request = HttpRequest {
228            method: self.method.clone(),
229            url: self.url.clone(),
230            headers: headers_map,
231            body: self.body,
232        };
233
234        self.client.execute(request).await
235    }
236}
237
238#[derive(Debug, Clone, PartialEq)]
239pub enum RedirectPolicy {
240    Follow,
241    Limit(usize),
242    None,
243}
244
245impl Default for RedirectPolicy {
246    fn default() -> Self {
247        Self::Limit(5)
248    }
249}
250
251pub struct RedirectHandler;
252
253impl RedirectHandler {
254    pub fn should_follow_redirect(
255        status_code: u16,
256        _method: &str,
257        current_redirects: usize,
258        max_redirects: usize,
259    ) -> Option<RedirectAction> {
260        if current_redirects >= max_redirects {
261            debug!("Maximum redirect limit reached: {}", max_redirects);
262            return None;
263        }
264
265        match status_code {
266            301 => Some(RedirectAction::FollowKeepMethod),
267            302 | 303 => Some(RedirectAction::FollowChangeToGet),
268            307 | 308 => Some(RedirectAction::FollowKeepMethod),
269            _ => None,
270        }
271    }
272
273    pub fn resolve_redirect_url(current_url: &str, location: &str) -> Result<String, String> {
274        let location = location.trim();
275        if location.is_empty() {
276            return Err("Location header is empty".to_string());
277        }
278
279        if location.starts_with("http://") || location.starts_with("https://") {
280            return Ok(location.to_string());
281        }
282
283        if location.starts_with("/") {
284            if let Some(base_end) = current_url[8..].find('/') {
285                Ok(format!("{}{}", &current_url[..8 + base_end], location))
286            } else {
287                Ok(format!("{}{}", current_url, location))
288            }
289        } else {
290            let last_slash = current_url.rfind('/').unwrap_or(current_url.len());
291            if last_slash < 8 {
292                Ok(format!("{}/{}", current_url, location))
293            } else {
294                Ok(format!("{}{}", &current_url[..last_slash + 1], location))
295            }
296        }
297    }
298
299    pub fn sanitize_redirect_url(url: &str) -> Result<String, String> {
300        if !url.starts_with("http://") && !url.starts_with("https://") {
301            return Err(format!("Unsafe redirect URL protocol: {}", url));
302        }
303        Ok(url.to_string())
304    }
305}
306
307pub enum RedirectAction {
308    FollowKeepMethod,
309    FollowChangeToGet,
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315
316    #[test]
317    fn test_should_follow_301() {
318        let action = RedirectHandler::should_follow_redirect(301, "GET", 0, 5);
319        assert!(action.is_some());
320
321        let no_action = RedirectHandler::should_follow_redirect(301, "GET", 5, 5);
322        assert!(no_action.is_none());
323    }
324
325    #[test]
326    fn test_should_not_follow_200() {
327        let action = RedirectHandler::should_follow_redirect(200, "GET", 0, 5);
328        assert!(action.is_none());
329    }
330
331    #[test]
332    fn test_resolve_absolute_redirect() {
333        let url = RedirectHandler::resolve_redirect_url(
334            "http://example.com/page",
335            "http://other.example.com/new",
336        )
337        .unwrap();
338        assert_eq!(url, "http://other.example.com/new");
339    }
340
341    #[test]
342    fn test_resolve_relative_redirect() {
343        let url = RedirectHandler::resolve_redirect_url("http://example.com/old/path", "/new/path")
344            .unwrap();
345        assert_eq!(url, "http://example.com/new/path");
346    }
347
348    #[test]
349    fn test_resolve_relative_path_redirect() {
350        let url = RedirectHandler::resolve_redirect_url(
351            "http://example.com/old/page.html",
352            "new-page.html",
353        )
354        .unwrap();
355        assert_eq!(url, "http://example.com/old/new-page.html");
356    }
357
358    #[test]
359    fn test_302_changes_to_get() {
360        let action = RedirectHandler::should_follow_redirect(302, "POST", 0, 5);
361        assert!(action.is_some());
362        if let Some(action) = action {
363            assert!(matches!(action, RedirectAction::FollowChangeToGet));
364        }
365    }
366
367    #[test]
368    fn test_307_keeps_method() {
369        let action = RedirectHandler::should_follow_redirect(307, "POST", 0, 5);
370        assert!(action.is_some());
371        if let Some(action) = action {
372            assert!(matches!(action, RedirectAction::FollowKeepMethod));
373        }
374    }
375}
aria2_protocol/http/client.rs

aria2_protocol/http/
client.rs
