aria2_core/http/
connection.rs

1//! HTTP 连接管理器
2//!
3//! 提供连接池复用、Keep-Alive 管理、LRU 淘汰策略和重定向跟随等功能。
4//!
5//! # 示例
6//!
7//! ```rust,no_run
8//! use aria2_core::http::connection::{HttpConnectionManager, HttpConfig};
9//! use std::time::Duration;
10//!
11//! #[tokio::main]
12//! async fn main() {
13//!     let config = HttpConfig {
14//!         max_connections: 10,
15//!         connect_timeout: Duration::from_secs(30),
16//!         read_timeout: Duration::from_secs(60),
17//!         write_timeout: Duration::from_secs(60),
18//!         idle_timeout: Duration::from_secs(300),
19//!     };
20//!
21//!     let manager = HttpConnectionManager::new(&config);
22//!     // 使用连接管理器...
23//! }
24//! ```
25
26use std::collections::{HashMap, HashSet};
27use std::net::SocketAddr;
28use std::sync::atomic::{AtomicU64, Ordering};
29use std::time::{Duration, Instant};
30
31use tokio::io::{AsyncReadExt, AsyncWriteExt};
32use tokio::net::TcpStream;
33use tokio::time::timeout;
34use url::Url;
35
36use crate::error::{Aria2Error, RecoverableError, Result};
37use crate::http::cookie_storage::{CookieJar, JarCookie};
38
39/// HTTP 连接配置
40#[derive(Debug, Clone)]
41pub struct HttpConfig {
42    /// 最大并发连接数
43    pub max_connections: usize,
44    /// TCP 连接超时
45    pub connect_timeout: Duration,
46    /// 读取超时
47    pub read_timeout: Duration,
48    /// 写入超时
49    pub write_timeout: Duration,
50    /// 空闲连接超时（LRU 淘汰）
51    pub idle_timeout: Duration,
52}
53
54impl Default for HttpConfig {
55    fn default() -> Self {
56        Self {
57            max_connections: crate::constants::HTTP_CONFIG_DEFAULT_MAX_CONNECTIONS,
58            connect_timeout: Duration::from_secs(
59                crate::constants::HTTP_CONFIG_DEFAULT_CONNECT_TIMEOUT_SECS,
60            ),
61            read_timeout: Duration::from_secs(
62                crate::constants::HTTP_CONFIG_DEFAULT_READ_TIMEOUT_SECS,
63            ),
64            write_timeout: Duration::from_secs(
65                crate::constants::HTTP_CONFIG_DEFAULT_WRITE_TIMEOUT_SECS,
66            ),
67            idle_timeout: Duration::from_secs(
68                crate::constants::HTTP_CONFIG_DEFAULT_IDLE_TIMEOUT_SECS,
69            ),
70        }
71    }
72}
73
74/// 活动连接信息
75#[derive(Debug)]
76pub struct ActiveConnection {
77    /// 唯一连接 ID
78    pub id: u64,
79    /// TCP 流
80    pub stream: TcpStream,
81    /// 目标主机
82    pub host: String,
83    /// 最后使用时间
84    pub last_used: Instant,
85}
86
87impl ActiveConnection {
88    /// 检查连接是否仍然有效
89    pub fn is_valid(&self) -> bool {
90        // 检查连接是否已关闭或出错
91        self.stream.peer_addr().is_ok()
92    }
93
94    /// 更新最后使用时间
95    pub fn touch(&mut self) {
96        self.last_used = Instant::now();
97    }
98}
99
100/// HTTP 连接状态
101#[derive(Debug, Clone, PartialEq)]
102pub enum ConnectionState {
103    /// 空闲可用
104    Idle,
105    /// 正在使用中
106    InUse,
107    /// 已关闭
108    Closed,
109}
110
111/// HTTP 连接管理器
112///
113/// 提供 HTTP 连接的获取、释放、池化管理和重定向跟随功能。
114/// 支持 Keep-Alive 连接复用、LRU 淘汰策略和循环重定向检测。
115///
116/// # 线程安全
117///
118/// `HttpConnectionManager` 内部使用 `tokio::sync::Mutex` 保护共享状态，
119/// 可以安全地在多个异步任务之间共享。
120///
121/// # 性能特性
122///
123/// - **连接复用**: 通过 Keep-Alive 头部检查，避免重复建立 TCP 连接
124/// - **LRU 淘汰**: 自动清理空闲超时的连接，防止资源泄漏
125/// - **三级超时**: 分别控制连接、读取、写入三个阶段的超时时间
126///
127/// # 示例
128///
129/// ```rust,no_run
130/// use aria2_core::http::connection::{HttpConnectionManager, HttpConfig};
131/// use std::time::Duration;
132/// use url::Url;
133///
134/// #[tokio::main]
135/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
136///     let config = HttpConfig {
137///         max_connections: 8,
138///         ..Default::default()
139///     };
140///
141///     let mut manager = HttpConnectionManager::new(&config);
142///     let url = Url::parse("https://example.com/file")?;
143///
144///     let conn = manager.acquire(&url).await?;
145///     // 使用连接进行 HTTP 请求...
146///     manager.release(conn.id).await;
147///
148///     Ok(())
149/// }
150/// ```
151pub struct HttpConnectionManager {
152    /// Configuration parameters
153    config: HttpConfig,
154    /// Connection pool: conn_id -> ActiveConnection
155    pool: HashMap<u64, ActiveConnection>,
156    /// Host-to-connection-ID mapping (for fast lookup of reusable connections)
157    host_connections: HashMap<String, Vec<u64>>,
158    /// Current active connection count
159    active_count: usize,
160    /// Connection ID generator
161    id_counter: AtomicU64,
162    /// Maximum redirect hops
163    max_redirects: u32,
164    /// Optional cookie jar for automatic cookie management on HTTP requests.
165    ///
166    /// When set, the connection manager will:
167    /// - Attach matching Cookie headers to outgoing requests via `attach_cookies_to_request()`
168    /// - Extract and store Set-Cookie headers from responses via `extract_cookies_from_response()`
169    cookie_jar: Option<CookieJar>,
170}
171
172impl HttpConnectionManager {
173    /// 创建新的 HTTP 连接管理器
174    ///
175    /// # 参数
176    ///
177    /// * `config` - HTTP 连接配置，包含超时、最大连接数等参数
178    ///
179    /// # 返回值
180    ///
181    /// 返回初始化完成的连接管理器实例
182    ///
183    /// # 示例
184    ///
185    /// ```
186    /// use aria2_core::http::connection::{HttpConnectionManager, HttpConfig};
187    /// use std::time::Duration;
188    ///
189    /// let config = HttpConfig {
190    ///     max_connections: 10,
191    ///     connect_timeout: Duration::from_secs(15),
192    ///     read_timeout: Duration::from_secs(30),
193    ///     write_timeout: Duration::from_secs(30),
194    ///     idle_timeout: Duration::from_secs(120),
195    /// };
196    ///
197    /// let manager = HttpConnectionManager::new(&config);
198    /// assert_eq!(manager.max_connections(), 10);
199    /// ```
200    pub fn new(config: &HttpConfig) -> Self {
201        Self {
202            config: config.clone(),
203            pool: HashMap::new(),
204            host_connections: HashMap::new(),
205            active_count: 0,
206            id_counter: AtomicU64::new(1),
207            max_redirects: crate::constants::HTTP_DEFAULT_MAX_REDIRECTS as u32,
208            cookie_jar: None,
209        }
210    }
211
212    /// 获取最大连接数配置
213    pub fn max_connections(&self) -> usize {
214        self.config.max_connections
215    }
216
217    /// 获取当前活动连接数
218    pub fn active_count(&self) -> usize {
219        self.active_count
220    }
221
222    /// 获取连接池大小（包含空闲和使用中的连接）
223    pub fn pool_size(&self) -> usize {
224        self.pool.len()
225    }
226
227    /// 从连接池获取或新建一个到指定 URL 的连接
228    ///
229    /// 该方法会尝试从连接池中查找可复用的空闲连接（基于主机名匹配），
230    /// 如果没有可用连接且未达到最大连接数限制，则创建新连接。
231    ///
232    /// # 参数
233    ///
234    /// * `url` - 目标 URL，用于提取主机名和端口信息
235    ///
236    /// # 错误
237    ///
238    /// * [`Aria2Error::Network`] - 当达到最大连接数限制时
239    /// * [`Aria2Error::Recoverable`] - 当连接超时或网络故障时
240    ///
241    /// # 返回值
242    ///
243    /// 返回可用的活动连接实例
244    ///
245    /// # Keep-Alive 复用逻辑
246    ///
247    /// 1. 提取 URL 的 host:port 作为连接标识
248    /// 2. 在连接池中查找该主机的空闲连接
249    /// 3. 验证连接有效性（检查 socket 是否正常）
250    /// 4. 更新 last_used 时间戳并返回
251    /// 5. 如果无可用连接，创建新的 TCP 连接
252    ///
253    /// # 示例
254    ///
255    /// ```rust,no_run
256    /// use aria2_core::http::connection::HttpConnectionManager;
257    /// use url::Url;
258    ///
259    /// #[tokio::main]
260    /// async fn main() {
261    ///     let mut manager = HttpConnectionManager::new(&Default::default());
262    ///     let url = Url::parse("https://example.com/resource").unwrap();
263    ///
264    ///     match manager.acquire(&url).await {
265    ///         Ok(conn) => println!("获取连接成功: id={}", conn.id),
266    ///         Err(e) => eprintln!("获取连接失败: {}", e),
267    ///     }
268    /// }
269    /// ```
270    pub async fn acquire(&mut self, url: &Url) -> Result<ActiveConnection> {
271        let host = Self::extract_host(url);
272
273        // 尝试从连接池中复用空闲连接
274        if let Some(conn) = self.try_reuse_connection(&host)? {
275            tracing::debug!("复用连接: id={}, host={}", conn.id, host);
276            return Ok(conn);
277        }
278
279        // 检查是否达到最大连接数限制
280        if self.active_count >= self.config.max_connections {
281            // 尝试清理过期连接
282            self.evict_idle_connections();
283
284            if self.active_count >= self.config.max_connections {
285                return Err(Aria2Error::Recoverable(
286                    RecoverableError::TemporaryNetworkFailure {
287                        message: format!(
288                            "达到最大连接数限制: {} (host={})",
289                            self.config.max_connections, host
290                        ),
291                    },
292                ));
293            }
294        }
295
296        // 创建新连接
297        self.create_new_connection(url, &host).await
298    }
299
300    /// 归还连接到连接池
301    ///
302    /// 将使用完毕的连接归还到连接池中，以便后续复用。
303    /// 如果连接已失效（socket 关闭或出错），会自动从池中移除。
304    ///
305    /// # 参数
306    ///
307    /// * `conn_id` - 要归还的连接 ID
308    ///
309    /// # 行为
310    ///
311    /// 1. 根据连接 ID 在池中查找对应连接
312    /// 2. 更新 last_used 为当前时间
313    /// 3. 将连接标记为空闲状态
314    /// 4. 如果连接无效，自动清理资源
315    ///
316    /// # 示例
317    ///
318    /// ```rust,no_run
319    /// use aria2_core::http::connection::HttpConnectionManager;
320    /// use url::Url;
321    ///
322    /// #[tokio::main]
323    /// async fn main() {
324    ///     let mut manager = HttpConnectionManager::new(&Default::default());
325    ///     let url = Url::parse("https://example.com").unwrap();
326    ///
327    ///     let conn = manager.acquire(&url).await.unwrap();
328    ///     // 使用连接完成请求后...
329    ///     manager.release(conn.id).await;
330    /// }
331    /// ```
332    pub async fn release(&mut self, conn_id: u64) {
333        if let Some(mut conn) = self.pool.remove(&conn_id) {
334            // 验证连接是否仍然有效
335            if !conn.is_valid() {
336                tracing::debug!("连接已失效，移除: id={}", conn_id);
337                self.active_count = self.active_count.saturating_sub(1);
338                self.remove_from_host_map(&conn.host, conn_id);
339                return;
340            }
341
342            // 更新最后使用时间并放回池中
343            conn.touch();
344            self.pool.insert(conn_id, conn);
345
346            tracing::debug!("归还连接到池: id={}", conn_id);
347        } else {
348            tracing::warn!("尝试释放不存在的连接: id={}", conn_id);
349        }
350    }
351
352    /// 跟随 HTTP 重定向
353    ///
354    /// 解析响应中的 Location 头部，构建新的 URL 并验证重定向合法性。
355    /// 支持相对路径和绝对路径的重定向，自动处理循环重定向检测。
356    ///
357    /// # 参数
358    ///
359    /// * `response` - HTTP 响应对象，需包含 Location 头部
360    /// * `current_url` - 当前请求的 URL（用于解析相对路径）
361    /// * `redirect_chain` - 已访问过的 URL 集合（用于循环检测）
362    ///
363    /// # 错误
364    ///
365    /// * [`Aria2Error::Parse`] - 当 Location 头部格式无效或 URL 解析失败时
366    /// * [`Aria2Error::Network`] - 当检测到循环重定向或超过最大跳数时
367    ///
368    /// # 返回值
369    ///
370    /// 返回重定向目标的新 URL
371    ///
372    /// # 重定向链检测机制
373    ///
374    /// 1. 使用 HashSet 记录所有已访问 URL
375    /// 2. 每次重定向前检查新 URL 是否已在集合中
376    /// 3. 维护跳数计数器，超过阈值返回错误
377    /// 4. 支持最多 5 次 301/302/303/307/308 重定向
378    ///
379    /// # 示例
380    ///
381    /// ```rust,no_run
382    /// use aria2_core::http::connection::{HttpConnectionManager, HttpResponse};
383    /// use std::collections::HashSet;
384    /// use url::Url;
385    ///
386    /// #[tokio::main]
387    /// async fn main() {
388    ///     let manager = HttpConnectionManager::new(&Default::default());
389    ///     let current_url = Url::parse("http://example.com/old").unwrap();
390    ///     let mut chain = HashSet::new();
391    ///     chain.insert(current_url.clone());
392    ///
393    ///     let mut response = HttpResponse::new(301, "Moved".to_string());
394    ///     response.headers.push((
395    ///         "Location".to_string(),
396    ///         "/new-path".to_string(),
397    ///     ));
398    ///
399    ///     match manager.follow_redirects(&response, &current_url, &chain, 1) {
400    ///         Ok(new_url) => println!("重定向到: {}", new_url),
401    ///         Err(e) => eprintln!("重定向失败: {}", e),
402    ///     }
403    /// }
404    /// ```
405    pub fn follow_redirects(
406        &self,
407        response: &HttpResponse,
408        current_url: &Url,
409        redirect_chain: &HashSet<Url>,
410        redirect_count: u32,
411    ) -> Result<Url> {
412        // 检查是否为重定向响应
413        if !response.is_redirect() {
414            return Err(Aria2Error::Parse(format!(
415                "非重定向响应码: {}",
416                response.status_code
417            )));
418        }
419
420        // 检查重定向次数限制
421        if redirect_count >= self.max_redirects {
422            return Err(Aria2Error::Network(format!(
423                "超过最大重定向次数限制: {}",
424                self.max_redirects
425            )));
426        }
427
428        // 获取 Location 头部
429        let location = response
430            .location()
431            .ok_or_else(|| Aria2Error::Parse("缺少 Location 头部".to_string()))?;
432
433        // 解析新的 URL（支持相对路径）
434        let new_url = current_url
435            .join(location)
436            .map_err(|e| Aria2Error::Parse(format!("解析重定向 URL 失败: {}", e)))?;
437
438        // 循环重定向检测
439        if redirect_chain.contains(&new_url) {
440            return Err(Aria2Error::Network(format!(
441                "检测到循环重定向: {}",
442                new_url
443            )));
444        }
445
446        tracing::info!(
447            "跟随重定向: {} -> {} ({}/{})",
448            current_url,
449            new_url,
450            redirect_count + 1,
451            self.max_redirects
452        );
453
454        Ok(new_url)
455    }
456
457    /// Iteratively follow HTTP redirects with loop detection
458    ///
459    /// This method replaces recursive redirect following with an iterative approach,
460    /// eliminating stack overflow risk for deep redirect chains.
461    ///
462    /// # Arguments
463    ///
464    /// * `initial_url` - The starting URL for the request
465    /// * `get_response` - Async closure that fetches the HTTP response for a given URL
466    ///
467    /// # Returns
468    ///
469    /// The final non-redirect HttpResponse, or an error if:
470    /// - Too many redirects (exceeds MAX_REDIRECTS limit)
471    /// - Redirect loop detected (same URL visited twice)
472    /// - Missing Location header in redirect response
473    /// - Invalid URL in Location header
474    ///
475    /// # Performance characteristics
476    ///
477    /// - Uses HashSet<String> for O(1) loop detection instead of linear scan
478    /// - Iterative loop with bounded iterations prevents stack growth
479    /// - Maximum 5 redirects as per RFC 7231 recommendation
480    pub async fn follow_redirects_iterative<F, Fut>(
481        &self,
482        initial_url: &Url,
483        mut get_response: F,
484    ) -> Result<HttpResponse>
485    where
486        F: FnMut(&Url) -> Fut,
487        Fut: std::future::Future<Output = Result<HttpResponse>>,
488    {
489        const MAX_REDIRECTS: u8 = crate::constants::HTTP_DEFAULT_MAX_REDIRECTS as u8;
490
491        let mut current_url = initial_url.clone();
492        let mut seen_urls = std::collections::HashSet::<String>::new();
493
494        for iteration in 0..MAX_REDIRECTS {
495            // Detect redirect loops using HashSet for O(1) lookup
496            let url_str = current_url.to_string();
497            if !seen_urls.insert(url_str.clone()) {
498                return Err(Aria2Error::Network(format!(
499                    "Redirect loop detected: {}",
500                    url_str
501                )));
502            }
503
504            // Fetch response for current URL
505            let resp = get_response(&current_url).await?;
506
507            // If not a redirect, return the final response
508            if !resp.is_redirect() {
509                return Ok(resp);
510            }
511
512            // Extract Location header from redirect response
513            let location = resp.location().ok_or_else(|| {
514                Aria2Error::Network("Missing Location header in redirect response".into())
515            })?;
516
517            // Resolve relative URLs against current URL
518            current_url = current_url
519                .join(location)
520                .map_err(|e| Aria2Error::Parse(format!("Failed to parse redirect URL: {}", e)))?;
521
522            tracing::info!(
523                "Following redirect: iteration {}/{}",
524                iteration + 1,
525                MAX_REDIRECTS
526            );
527        }
528
529        Err(Aria2Error::Network(format!(
530            "Too many redirects (>{}), last URL: {}",
531            MAX_REDIRECTS, current_url
532        )))
533    }
534
535    /// 构建 Range 请求头
536    ///
537    /// 根据 start 和 end 字节位置构建符合 RFC 7233 规范的 Range 头部字符串。
538    /// 用于断点续传和分块下载场景。
539    ///
540    /// # 参数
541    ///
542    /// * `start` - 起始字节位置（包含）
543    /// * `end` - 结束字节位置（包含），如果为 None 则表示到文件末尾
544    ///
545    /// # 返回值
546    ///
547    /// 返回格式化的 Range 头部值，例如 `"bytes=0-499"` 或 `"bytes=500-"`
548    ///
549    /// # 格式规范
550    ///
551    /// - `bytes=start-end`: 指定范围 [start, end]
552    /// - `bytes=start-`: 从 start 到文件末尾
553    ///
554    /// # 示例
555    ///
556    /// ```
557    /// use aria2_core::http::connection::HttpConnectionManager;
558    ///
559    /// let manager = HttpConnectionManager::new(&Default::default());
560    ///
561    /// // 完整范围
562    /// assert_eq!(
563    ///     manager.build_range_header(0, Some(499)),
564    ///     "bytes=0-499"
565    /// );
566    ///
567    /// // 开放结束范围
568    /// assert_eq!(
569    ///     manager.build_range_header(1000, None),
570    ///     "bytes=1000-"
571    /// );
572    ///
573    /// // 单字节范围
574    /// assert_eq!(
575    ///     manager.build_range_header(42, Some(42)),
576    ///     "bytes=42-42"
577    /// );
578    /// ```
579    pub fn build_range_header(&self, start: u64, end: Option<u64>) -> String {
580        match end {
581            Some(end_val) => format!("bytes={}-{}", start, end_val),
582            None => format!("bytes={}-", start),
583        }
584    }
585
586    /// 解析 Content-Range 响应头
587    ///
588    /// 解析服务器返回的 Content-Range 头部值，提取范围信息和总大小。
589    /// 用于验证服务器是否正确支持 Range 请求。
590    ///
591    /// # 参数
592    ///
593    /// * `header` - Content-Range 头部的原始字符串值
594    ///
595    /// # 返回值
596    ///
597    /// 如果解析成功，返回元组 `(start, end, total)`:
598    /// - `start`: 范围起始字节（包含）
599    /// - `end`: 范围结束字节（包含）
600    /// - `total`: 文件总字节数（如果未知则为 u64::MAX）
601    ///
602    /// 如果格式无效，返回 `None`
603    ///
604    /// # 支持格式
605    ///
606    /// - `bytes 0-499/1000`: 已知总大小的范围
607    /// - `bytes 0-499/*`: 未知总大小的范围
608    ///
609    /// # 示例
610    ///
611    /// ```
612    /// use aria2_core::http::connection::HttpConnectionManager;
613    ///
614    /// let manager = HttpConnectionManager::new(&Default::default());
615    ///
616    /// // 解析已知总大小
617    /// let result = manager.parse_content_range("bytes 0-499/1000");
618    /// assert_eq!(result, Some((0, 499, 1000)));
619    ///
620    /// // 解析未知总大小
621    /// let result = manager.parse_content_range("bytes 500-999/*");
622    /// assert_eq!(result, Some((500, 999, u64::MAX)));
623    ///
624    /// // 无效格式
625    /// assert_eq!(manager.parse_content_range("invalid"), None);
626    /// assert_eq!(manager.parse_content_range("bits 0-99/1000"), None);
627    /// ```
628    pub fn parse_content_range(&self, header: &str) -> Option<(u64, u64, u64)> {
629        let header = header.trim();
630
631        // 必须以 "bytes " 开头
632        if !header.starts_with("bytes ") {
633            return None;
634        }
635
636        let range_part = &header[6..];
637        let parts: Vec<&str> = range_part.split('/').collect();
638
639        if parts.len() != 2 {
640            return None;
641        }
642
643        // 解析 start-end 部分
644        let range_values: Vec<&str> = parts[0].split('-').collect();
645        if range_values.len() != 2 {
646            return None;
647        }
648
649        let start: u64 = range_values[0].trim().parse().ok()?;
650        let end: u64 = range_values[1].trim().parse().ok()?;
651
652        // 解析 total 大小
653        let total = match parts[1].trim() {
654            "*" => u64::MAX,
655            s => s.parse().ok()?,
656        };
657
658        Some((start, end, total))
659    }
660
661    /// 清理所有空闲连接
662    ///
663    /// 关闭连接池中的所有连接，释放系统资源。
664    /// 通常在下载任务完成或程序退出时调用。
665    pub async fn cleanup(&mut self) {
666        for (_, mut conn) in self.pool.drain() {
667            let _ = conn.shutdown().await;
668        }
669        self.host_connections.clear();
670        self.active_count = 0;
671
672        tracing::info!("连接池已清理");
673    }
674
675    /// Force close a specific connection
676    ///
677    /// Remove the connection from the pool and close the underlying TCP connection.
678    /// Used for error handling or abnormal termination.
679    ///
680    /// # Arguments
681    ///
682    /// * `conn_id` - The ID of the connection to close
683    pub async fn close_connection(&mut self, conn_id: u64) {
684        if let Some(mut conn) = self.pool.remove(&conn_id) {
685            let _ = conn.shutdown().await;
686            self.active_count = self.active_count.saturating_sub(1);
687            self.remove_from_host_map(&conn.host, conn_id);
688            tracing::debug!("Force closed connection: id={}", conn_id);
689        }
690    }
691
692    // ==================== Cookie Jar Integration (J4) ====================
693
694    /// Set the cookie jar for automatic cookie management on HTTP requests.
695    ///
696    /// Once set, the connection manager will automatically:
697    /// - Attach `Cookie` headers with matching cookies when building outgoing requests
698    /// - Parse and store cookies from `Set-Cookie` response headers
699    ///
700    /// # Arguments
701    ///
702    /// * `jar` - The CookieJar instance to use for cookie storage and matching
703    ///
704    /// # Example
705    ///
706    /// ```rust,no_run
707    /// use aria2_core::http::connection::{HttpConnectionManager, HttpConfig};
708    /// use aria2_core::http::cookie_storage::CookieJar;
709    ///
710    /// let mut manager = HttpConnectionManager::new(&Default::default());
711    /// let jar = CookieJar::new();
712    /// manager.set_cookie_jar(Some(jar));
713    /// ```
714    pub fn set_cookie_jar(&mut self, jar: Option<CookieJar>) {
715        self.cookie_jar = jar;
716    }
717
718    /// Get a reference to the current cookie jar, if one is set.
719    pub fn cookie_jar(&self) -> &Option<CookieJar> {
720        &self.cookie_jar
721    }
722
723    /// Get a mutable reference to the current cookie jar, if one is set.
724    pub fn cookie_jar_mut(&mut self) -> &mut Option<CookieJar> {
725        &mut self.cookie_jar
726    }
727
728    /// Attach matching cookies from the jar to an HTTP request as a Cookie header string.
729    ///
730    /// Call this method before sending an HTTP request to include any stored cookies
731    /// that match the target URL. The returned string can be used directly as the
732    /// value of the `Cookie` request header.
733    ///
734    /// # Arguments
735    ///
736    /// * `url` - The target URL for the HTTP request
737    ///
738    /// # Returns
739    ///
740    /// `Some(header_value)` containing `"name1=val1; name2=val2"` format if matching
741    /// cookies exist, or `None` if no cookies match or no jar is configured.
742    ///
743    /// # Example
744    ///
745    /// ```rust,no_run
746    /// use aria2_core::http::connection::HttpConnectionManager;
747    /// use url::Url;
748    ///
749    /// let manager = HttpConnectionManager::new(&Default::default());
750    /// let url = Url::parse("https://example.com/api").unwrap();
751    ///
752    /// if let Some(cookie_header) = manager.attach_cookies_to_request(&url) {
753    ///     // Add "Cookie: {cookie_header}" to your HTTP request headers
754    ///     println!("Cookie: {}", cookie_header);
755    /// }
756    /// ```
757    pub fn attach_cookies_to_request(&self, url: &Url) -> Option<String> {
758        let jar = self.cookie_jar.as_ref()?;
759        let is_https = url.scheme() == "https";
760        jar.cookie_header_for_url(url.as_str(), is_https)
761    }
762
763    /// Extract cookies from response Set-Cookie headers and store them in the jar.
764    ///
765    /// Call this method after receiving an HTTP response to persist any cookies
766    /// set by the server. Each `Set-Cookie` header value is parsed and stored
767    /// in the cookie jar for future requests.
768    ///
769    /// # Arguments
770    ///
771    /// * `response_headers` - The response headers as a slice of `(name, value)` tuples
772    /// * `request_url` - The original request URL (used as default domain/path context)
773    ///
774    /// # Returns
775    ///
776    /// The number of cookies successfully extracted and stored.
777    ///
778    /// # Example
779    ///
780    /// ```rust,no_run
781    /// use aria2_core::http::connection::HttpConnectionManager;
782    ///
783    /// // After receiving an HTTP response:
784    /// let headers = vec![
785    ///     ("Set-Cookie".to_string(), "session=abc; Domain=example.com".to_string()),
786    ///     ("Set-Cookie".to_string(), "theme=dark".to_string()),
787    /// ];
788    /// let url = url::Url::parse("https://example.com/").unwrap();
789    ///
790    /// let mut manager = HttpConnectionManager::new(&Default::default());
791    /// manager.set_cookie_jar(Some(aria2_core::http::cookie_storage::CookieJar::new()));
792    /// let count = manager.extract_cookies_from_response(&headers, &url);
793    /// println!("Stored {} cookies", count); // Prints: Stored 2 cookies
794    /// ```
795    pub fn extract_cookies_from_response(
796        &mut self,
797        response_headers: &[(String, String)],
798        _request_url: &Url,
799    ) -> usize {
800        let jar = match &mut self.cookie_jar {
801            Some(j) => j,
802            None => return 0,
803        };
804
805        let mut stored = 0;
806        for (name, value) in response_headers {
807            if name.eq_ignore_ascii_case("set-cookie")
808                && let Some(cookie) = JarCookie::parse_set_cookie(value)
809            {
810                jar.store(cookie);
811                stored += 1;
812                tracing::debug!(
813                    "Extracted and stored cookie from Set-Cookie header: {}",
814                    &value[..value.len().min(80)]
815                );
816            }
817        }
818        stored
819    }
820
821    // ==================== Private Helper Methods ====================
822
823    /// 从 URL 中提取主机标识（host:port）
824    fn extract_host(url: &Url) -> String {
825        match url.port_or_known_default() {
826            Some(port) => format!(
827                "{}:{}",
828                url.host_str().unwrap_or(crate::constants::DEFAULT_HOST),
829                port
830            ),
831            None => url
832                .host_str()
833                .unwrap_or(crate::constants::DEFAULT_HOST)
834                .to_string(),
835        }
836    }
837
838    /// 尝试从连接池复用连接
839    fn try_reuse_connection(&mut self, host: &str) -> Result<Option<ActiveConnection>> {
840        let conn_ids = match self.host_connections.get(host) {
841            Some(ids) => ids.clone(),
842            None => return Ok(None),
843        };
844
845        // 查找可用的空闲连接
846        for &conn_id in &conn_ids {
847            if let Some(mut conn) = self.pool.remove(&conn_id) {
848                // 验证连接有效性
849                if conn.is_valid() {
850                    conn.touch();
851
852                    // 检查 Keep-Alive 状态（简化版：仅检查时间）
853                    let idle_time = conn.last_used.elapsed();
854                    if idle_time < self.config.idle_timeout {
855                        tracing::debug!(
856                            "复用空闲连接: id={}, idle={:.2}s",
857                            conn_id,
858                            idle_time.as_secs_f64()
859                        );
860                        return Ok(Some(conn));
861                    } else {
862                        // 连接过期，关闭并继续查找
863                        tracing::debug!(
864                            "连接已过期: id={}, idle={:.2}s",
865                            conn_id,
866                            idle_time.as_secs_f64()
867                        );
868                        self.active_count = self.active_count.saturating_sub(1);
869                        std::mem::drop(conn.shutdown()); // 忽略关闭错误
870                    }
871                } else {
872                    // 连接已失效
873                    self.active_count = self.active_count.saturating_sub(1);
874                }
875            }
876        }
877
878        // 清理该主机的所有无效连接记录
879        self.cleanup_invalid_connections(host);
880
881        Ok(None)
882    }
883
884    /// 创建新的 TCP 连接
885    async fn create_new_connection(&mut self, url: &Url, host: &str) -> Result<ActiveConnection> {
886        // 解析地址
887        let addr = Self::resolve_address(url)?;
888
889        // 应用连接超时
890        let stream = timeout(self.config.connect_timeout, TcpStream::connect(&addr))
891            .await
892            .map_err(|_| Aria2Error::Recoverable(RecoverableError::Timeout))?
893            .map_err(|e| {
894                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure {
895                    message: format!("TCP 连接失败 ({}): {}", addr, e),
896                })
897            })?;
898
899        // 设置 TCP 选项
900        if let Err(e) = stream.set_nodelay(true) {
901            tracing::warn!("设置 nodelay 失败: {}", e);
902        }
903        // 注意: tokio TcpStream 不直接支持 set_keepalive，需要使用 socket2 或忽略
904
905        // 生成连接 ID
906        let conn_id = self.id_counter.fetch_add(1, Ordering::SeqCst);
907
908        let conn = ActiveConnection {
909            id: conn_id,
910            stream,
911            host: host.to_string(),
912            last_used: Instant::now(),
913        };
914
915        // 更新连接池状态
916        self.active_count += 1;
917        self.host_connections
918            .entry(host.to_string())
919            .or_default()
920            .push(conn_id);
921
922        tracing::info!(
923            "创建新连接: id={}, host={}, active={}/{}",
924            conn_id,
925            host,
926            self.active_count,
927            self.config.max_connections
928        );
929
930        Ok(conn)
931    }
932
933    /// 解析 URL 为 SocketAddr
934    fn resolve_address(url: &Url) -> Result<SocketAddr> {
935        let host = url
936            .host_str()
937            .ok_or_else(|| Aria2Error::Parse("URL 缺少主机名".to_string()))?;
938
939        let port = url
940            .port_or_known_default()
941            .ok_or_else(|| Aria2Error::Parse("无法确定端口号".to_string()))?;
942
943        // 使用 tokio 进行 DNS 解析（同步版本用于测试兼容性）
944        // 注意：生产环境应该使用 tokio::net::lookup_host
945        let addr_str = format!("{}:{}", host, port);
946        addr_str
947            .parse::<SocketAddr>()
948            .map_err(|e| Aria2Error::Parse(format!("解析地址失败: {}", e)))
949    }
950
951    /// LRU 淘汰：清理空闲超时的连接
952    fn evict_idle_connections(&mut self) {
953        let now = Instant::now();
954        let mut evicted = Vec::new();
955
956        for (&conn_id, conn) in &self.pool {
957            if now.duration_since(conn.last_used) > self.config.idle_timeout {
958                evicted.push((conn_id, conn.host.clone()));
959            }
960        }
961
962        let evict_count = evicted.len();
963        for (conn_id, host) in evicted {
964            if let Some(mut conn) = self.pool.remove(&conn_id) {
965                std::mem::drop(conn.shutdown());
966                self.active_count = self.active_count.saturating_sub(1);
967                self.remove_from_host_map(&host, conn_id);
968                tracing::debug!("LRU 淘汰过期连接: id={}, host={}", conn_id, host);
969            }
970        }
971
972        if evict_count > 0 {
973            tracing::info!("LRU 淘汰了 {} 个过期连接", evict_count);
974        }
975    }
976
977    /// 清理指定主机的无效连接记录
978    fn cleanup_invalid_connections(&mut self, host: &str) {
979        if let Some(ids) = self.host_connections.get_mut(host) {
980            ids.retain(|&id| self.pool.contains_key(&id));
981            if ids.is_empty() {
982                self.host_connections.remove(host);
983            }
984        }
985    }
986
987    /// 从主机映射中移除连接 ID
988    fn remove_from_host_map(&mut self, host: &str, conn_id: u64) {
989        if let Some(ids) = self.host_connections.get_mut(host) {
990            ids.retain(|&id| id != conn_id);
991            if ids.is_empty() {
992                self.host_connections.remove(host);
993            }
994        }
995    }
996}
997
998impl ActiveConnection {
999    /// 异步读取数据（带超时控制）
1000    ///
1001    /// 从 TCP 流中读取数据到缓冲区，受 read_timeout 限制。
1002    /// 用于读取 HTTP 响应头和响应体。
1003    pub async fn read_with_timeout(
1004        &mut self,
1005        buf: &mut [u8],
1006        read_timeout: Duration,
1007    ) -> Result<usize> {
1008        timeout(read_timeout, self.stream.read(buf))
1009            .await
1010            .map_err(|_| Aria2Error::Recoverable(RecoverableError::Timeout))?
1011            .map_err(|e| {
1012                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure {
1013                    message: format!("读取数据失败: {}", e),
1014                })
1015            })
1016    }
1017
1018    /// 异步写入数据（带超时控制）
1019    ///
1020    /// 将数据写入 TCP 流，受 write_timeout 限制。
1021    /// 用于发送 HTTP 请求头和请求体。
1022    pub async fn write_with_timeout(
1023        &mut self,
1024        buf: &[u8],
1025        write_timeout: Duration,
1026    ) -> Result<usize> {
1027        timeout(write_timeout, self.stream.write(buf))
1028            .await
1029            .map_err(|_| Aria2Error::Recoverable(RecoverableError::Timeout))?
1030            .map_err(|e| {
1031                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure {
1032                    message: format!("写入数据失败: {}", e),
1033                })
1034            })
1035    }
1036
1037    /// 刷新写缓冲区（带超时控制）
1038    pub async fn flush_with_timeout(&mut self, write_timeout: Duration) -> Result<()> {
1039        timeout(write_timeout, self.stream.flush())
1040            .await
1041            .map_err(|_| Aria2Error::Recoverable(RecoverableError::Timeout))?
1042            .map_err(|e| {
1043                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure {
1044                    message: format!("刷新缓冲区失败: {}", e),
1045                })
1046            })
1047    }
1048
1049    /// 关闭连接（双向关闭）
1050    pub async fn shutdown(&mut self) -> Result<()> {
1051        match self.stream.shutdown().await {
1052            Ok(_) => Ok(()),
1053            Err(e) => {
1054                tracing::debug!("关闭连接失败: id={}, error={}", self.id, e);
1055                Ok(())
1056            }
1057        }
1058    }
1059
1060    /// 获取对等端地址
1061    pub fn peer_addr(&self) -> Result<SocketAddr> {
1062        self.stream.peer_addr().map_err(|e| {
1063            Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure {
1064                message: format!("获取对等端地址失败: {}", e),
1065            })
1066        })
1067    }
1068
1069    /// 获取本地地址
1070    pub fn local_addr(&self) -> Result<SocketAddr> {
1071        self.stream.local_addr().map_err(|e| {
1072            Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure {
1073                message: format!("获取本地地址失败: {}", e),
1074            })
1075        })
1076    }
1077}
1078
1079impl Drop for HttpConnectionManager {
1080    fn drop(&mut self) {
1081        // 同步清理（不使用 async）
1082        for (_, conn) in self.pool.drain() {
1083            // TcpStream 的 drop 会自动关闭
1084            drop(conn);
1085        }
1086        self.host_connections.clear();
1087    }
1088}
1089
1090impl std::fmt::Debug for HttpConnectionManager {
1091    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1092        f.debug_struct("HttpConnectionManager")
1093            .field("max_connections", &self.config.max_connections)
1094            .field("connect_timeout", &self.config.connect_timeout)
1095            .field("read_timeout", &self.config.read_timeout)
1096            .field("write_timeout", &self.config.write_timeout)
1097            .field("idle_timeout", &self.config.idle_timeout)
1098            .field("active_count", &self.active_count)
1099            .field("pool_size", &self.pool.len())
1100            .field("cookie_jar_set", &self.cookie_jar.is_some())
1101            .finish()
1102    }
1103}
1104
1105// 重新导出 HttpResponse 以便在 connection.rs 中使用
1106pub use aria2_protocol::http::response::HttpResponse;
1107
1108#[cfg(test)]
1109mod tests {
1110    use super::*;
1111    use tokio::io::AsyncWriteExt;
1112    use tokio::time::{sleep, timeout};
1113
1114    fn create_test_config() -> HttpConfig {
1115        HttpConfig {
1116            max_connections: 4,
1117            connect_timeout: Duration::from_millis(100),
1118            read_timeout: Duration::from_millis(200),
1119            write_timeout: Duration::from_millis(200),
1120            idle_timeout: Duration::from_millis(500),
1121        }
1122    }
1123
1124    #[test]
1125    fn test_config_default() {
1126        let config = HttpConfig::default();
1127        assert_eq!(config.max_connections, 16);
1128        assert_eq!(config.connect_timeout, Duration::from_secs(30));
1129        assert_eq!(config.read_timeout, Duration::from_secs(60));
1130        assert_eq!(config.write_timeout, Duration::from_secs(60));
1131        assert_eq!(config.idle_timeout, Duration::from_secs(300));
1132    }
1133
1134    #[test]
1135    fn test_manager_creation() {
1136        let config = create_test_config();
1137        let manager = HttpConnectionManager::new(&config);
1138
1139        assert_eq!(manager.max_connections(), 4);
1140        assert_eq!(manager.active_count(), 0);
1141        assert_eq!(manager.pool_size(), 0);
1142    }
1143
1144    #[test]
1145    fn test_build_range_header() {
1146        let manager = HttpConnectionManager::new(&Default::default());
1147
1148        // 完整范围
1149        assert_eq!(manager.build_range_header(0, Some(999)), "bytes=0-999");
1150
1151        // 开放结束范围
1152        assert_eq!(manager.build_range_header(500, None), "bytes=500-");
1153
1154        // 单字节
1155        assert_eq!(manager.build_range_header(42, Some(42)), "bytes=42-42");
1156
1157        // 大数值
1158        assert_eq!(
1159            manager.build_range_header(u64::MAX - 1, Some(u64::MAX)),
1160            "bytes=18446744073709551614-18446744073709551615"
1161        );
1162    }
1163
1164    #[test]
1165    fn test_parse_content_range() {
1166        let manager = HttpConnectionManager::new(&Default::default());
1167
1168        // 正常格式（已知总数）
1169        assert_eq!(
1170            manager.parse_content_range("bytes 0-499/1000"),
1171            Some((0, 499, 1000))
1172        );
1173
1174        // 正常格式（未知总数）
1175        assert_eq!(
1176            manager.parse_content_range("bytes 500-999/*"),
1177            Some((500, 999, u64::MAX))
1178        );
1179
1180        // 边界值
1181        assert_eq!(manager.parse_content_range("bytes 0-0/1"), Some((0, 0, 1)));
1182
1183        // 无效格式
1184        assert_eq!(manager.parse_content_range(""), None);
1185        assert_eq!(manager.parse_content_range("invalid"), None);
1186        assert_eq!(manager.parse_content_range("bits 0-99/1000"), None);
1187        assert_eq!(manager.parse_content_range("bytes 0-499"), None); // 缺少 /total
1188        assert_eq!(manager.parse_content_range("bytes abc-def/1000"), None);
1189    }
1190
1191    #[test]
1192    fn test_follow_redirects_success() {
1193        let manager = HttpConnectionManager::new(&Default::default());
1194        let current_url = Url::parse("http://example.com/old").unwrap();
1195        let mut chain = HashSet::new();
1196        chain.insert(current_url.clone());
1197
1198        let mut response = HttpResponse::new(301, "Moved Permanently".to_string());
1199        response
1200            .headers
1201            .push(("Location".to_string(), "http://example.com/new".to_string()));
1202
1203        let result = manager.follow_redirects(&response, &current_url, &chain, 1);
1204        assert!(result.is_ok());
1205        let new_url = result.unwrap();
1206        assert!(new_url.as_str().starts_with("http://example.com/new"));
1207    }
1208
1209    #[test]
1210    fn test_follow_redirects_relative_path() {
1211        let manager = HttpConnectionManager::new(&Default::default());
1212        let current_url = Url::parse("http://example.com/path/page.html").unwrap();
1213        let chain = HashSet::new();
1214
1215        let mut response = HttpResponse::new(302, "Found".to_string());
1216        response
1217            .headers
1218            .push(("Location".to_string(), "../other".to_string()));
1219
1220        let result = manager.follow_redirects(&response, &current_url, &chain, 1);
1221        assert!(result.is_ok());
1222        assert_eq!(result.unwrap().as_str(), "http://example.com/other");
1223    }
1224
1225    #[test]
1226    fn test_follow_redirects_loop_detection() {
1227        let manager = HttpConnectionManager::new(&Default::default());
1228        let url_a = Url::parse("http://example.com/a").unwrap();
1229        let url_b = Url::parse("http://example.com/b").unwrap();
1230
1231        let mut chain = HashSet::new();
1232        chain.insert(url_a.clone());
1233        chain.insert(url_b.clone());
1234
1235        let mut response = HttpResponse::new(301, "Moved".to_string());
1236        response
1237            .headers
1238            .push(("Location".to_string(), "http://example.com/a".to_string()));
1239
1240        // 尝试重定向回已访问的 URL（循环）
1241        let result = manager.follow_redirects(&response, &url_b, &chain, 2);
1242        assert!(result.is_err());
1243        assert!(result.unwrap_err().to_string().contains("循环重定向"));
1244    }
1245
1246    #[test]
1247    fn test_follow_redirects_max_exceeded() {
1248        let manager = HttpConnectionManager::new(&Default::default());
1249        let current_url = Url::parse("http://example.com/start").unwrap();
1250        let chain = HashSet::new();
1251
1252        let mut response = HttpResponse::new(302, "Found".to_string());
1253        response.headers.push((
1254            "Location".to_string(),
1255            "http://example.com/next".to_string(),
1256        ));
1257
1258        // 超过最大重定向次数
1259        let result = manager.follow_redirects(&response, &current_url, &chain, 6);
1260        assert!(result.is_err());
1261        assert!(result.unwrap_err().to_string().contains("最大重定向"));
1262    }
1263
1264    #[test]
1265    fn test_follow_redirects_non_redirect_response() {
1266        let manager = HttpConnectionManager::new(&Default::default());
1267        let current_url = Url::parse("http://example.com/").unwrap();
1268        let chain = HashSet::new();
1269
1270        let response = HttpResponse::new(200, "OK".to_string());
1271
1272        let result = manager.follow_redirects(&response, &current_url, &chain, 0);
1273        assert!(result.is_err());
1274        assert!(result.unwrap_err().to_string().contains("非重定向"));
1275    }
1276
1277    #[test]
1278    fn test_follow_redirects_missing_location() {
1279        let manager = HttpConnectionManager::new(&Default::default());
1280        let current_url = Url::parse("http://example.com/").unwrap();
1281        let chain = HashSet::new();
1282
1283        let response = HttpResponse::new(301, "Moved".to_string());
1284
1285        let result = manager.follow_redirects(&response, &current_url, &chain, 0);
1286        assert!(result.is_err());
1287        assert!(result.unwrap_err().to_string().contains("Location"));
1288    }
1289
1290    #[test]
1291    fn test_extract_host() {
1292        // 带 80 端口
1293        let url = Url::parse("http://example.com/path").unwrap();
1294        assert_eq!(HttpConnectionManager::extract_host(&url), "example.com:80");
1295
1296        // 带 443 端口
1297        let url = Url::parse("https://example.com:443/path").unwrap();
1298        assert_eq!(HttpConnectionManager::extract_host(&url), "example.com:443");
1299
1300        // 自定义端口
1301        let url = Url::parse("http://example.com:8080/path").unwrap();
1302        assert_eq!(
1303            HttpConnectionManager::extract_host(&url),
1304            "example.com:8080"
1305        );
1306    }
1307
1308    #[test]
1309    fn test_debug_format() {
1310        let config = create_test_config();
1311        let manager = HttpConnectionManager::new(&config);
1312        let debug_str = format!("{:?}", manager);
1313
1314        assert!(debug_str.contains("HttpConnectionManager"));
1315        assert!(debug_str.contains("max_connections: 4"));
1316        assert!(debug_str.contains("active_count: 0"));
1317    }
1318
1319    // ==================== 集成测试 ====================
1320
1321    /// 启动一个简单的测试 HTTP 服务器
1322    async fn start_test_server(
1323        handler: impl Fn(TcpStream) + Send + 'static,
1324    ) -> (SocketAddr, tokio::task::JoinHandle<()>) {
1325        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1326        let addr = listener.local_addr().unwrap();
1327
1328        let handle = tokio::spawn(async move {
1329            while let Ok((stream, _)) = listener.accept().await {
1330                handler(stream);
1331            }
1332        });
1333
1334        (addr, handle)
1335    }
1336
1337    #[tokio::test]
1338    async fn test_connection_pool_reuse() {
1339        let config = HttpConfig {
1340            max_connections: 4,
1341            connect_timeout: Duration::from_millis(500),
1342            read_timeout: Duration::from_millis(1000),
1343            write_timeout: Duration::from_millis(1000),
1344            idle_timeout: Duration::from_millis(2000),
1345        };
1346        let mut manager = HttpConnectionManager::new(&config);
1347
1348        // 启动测试服务器
1349        let (addr, server_handle) = start_test_server(|mut stream| {
1350            tokio::spawn(async move {
1351                let response = "HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nOK";
1352                stream.write_all(response.as_bytes()).await.unwrap();
1353            });
1354        })
1355        .await;
1356
1357        sleep(Duration::from_millis(100)).await;
1358
1359        let url = Url::parse(&format!("http://{}", addr)).unwrap();
1360
1361        // 第一次获取连接
1362        let conn1 = manager.acquire(&url).await.expect("第一次获取连接应成功");
1363        let _conn1_id = conn1.id;
1364        assert_eq!(manager.active_count(), 1);
1365
1366        // 归还连接
1367        manager.release(conn1.id).await;
1368
1369        // 第二次获取连接（应该能成功获取）
1370        let conn2 = manager.acquire(&url).await.expect("第二次应能获取连接");
1371        assert!(manager.active_count() >= 1); // 连接数应 >= 1
1372
1373        // 清理
1374        manager.release(conn2.id).await;
1375        manager.cleanup().await;
1376        server_handle.abort();
1377    }
1378
1379    #[tokio::test]
1380    async fn test_redirect_follow_5_jumps() {
1381        let manager = HttpConnectionManager::new(&create_test_config());
1382        let current_url = Url::parse("http://example.com/start").unwrap();
1383        let mut redirect_chain = HashSet::new();
1384        redirect_chain.insert(current_url.clone());
1385
1386        let urls = [
1387            "http://example.com/page1",
1388            "http://example.com/page2",
1389            "http://example.com/page3",
1390            "http://example.com/page4",
1391            "http://example.com/final",
1392        ];
1393
1394        let mut current = current_url;
1395        for (i, target) in urls.iter().enumerate() {
1396            let mut response = HttpResponse::new(302, "Found".to_string());
1397            response
1398                .headers
1399                .push(("Location".to_string(), target.to_string()));
1400
1401            redirect_chain.insert(current.clone());
1402
1403            let result = manager.follow_redirects(&response, &current, &redirect_chain, i as u32);
1404            assert!(
1405                result.is_ok(),
1406                "第 {} 次重定向应成功: {:?}",
1407                i + 1,
1408                result.err()
1409            );
1410
1411            current = result.unwrap();
1412        }
1413
1414        assert!(current.as_str().contains("example.com/final"));
1415    }
1416
1417    #[tokio::test]
1418    async fn test_redirect_loop_detection() {
1419        let manager = HttpConnectionManager::new(&create_test_config());
1420
1421        let url_a = Url::parse("http://example.com/a").unwrap();
1422        let url_b = Url::parse("http://example.com/b").unwrap();
1423        let url_c = Url::parse("http://example.com/c").unwrap();
1424
1425        let mut chain = HashSet::new();
1426        chain.insert(url_a.clone());
1427        chain.insert(url_b.clone());
1428        chain.insert(url_c.clone());
1429
1430        let mut response = HttpResponse::new(301, "Moved".to_string());
1431        response
1432            .headers
1433            .push(("Location".to_string(), "http://example.com/a".to_string()));
1434
1435        let result = manager.follow_redirects(&response, &url_c, &chain, 3);
1436
1437        assert!(result.is_err());
1438        assert!(result.unwrap_err().to_string().contains("循环重定向"));
1439    }
1440
1441    #[test]
1442    fn test_range_request_build() {
1443        let manager = HttpConnectionManager::new(&create_test_config());
1444
1445        assert_eq!(manager.build_range_header(0, Some(999)), "bytes=0-999");
1446        assert_eq!(manager.build_range_header(500, None), "bytes=500-");
1447        assert_eq!(manager.build_range_header(42, Some(42)), "bytes=42-42");
1448
1449        assert_eq!(
1450            manager.parse_content_range("bytes 0-499/1000"),
1451            Some((0, 499, 1000))
1452        );
1453        assert_eq!(
1454            manager.parse_content_range("bytes 500-999/*"),
1455            Some((500, 999, u64::MAX))
1456        );
1457        assert_eq!(manager.parse_content_range("invalid"), None);
1458    }
1459
1460    #[tokio::test]
1461    async fn test_timeout_on_slow_server() {
1462        use std::time::Instant;
1463
1464        let config = HttpConfig {
1465            max_connections: 2,
1466            connect_timeout: Duration::from_millis(100),
1467            read_timeout: Duration::from_millis(200),
1468            write_timeout: Duration::from_millis(200),
1469            idle_timeout: Duration::from_secs(60),
1470        };
1471        let mut manager = HttpConnectionManager::new(&config);
1472
1473        let (addr, server_handle) = start_test_server(|_stream| {
1474            tokio::spawn(async move {
1475                sleep(Duration::from_secs(10)).await;
1476            });
1477        })
1478        .await;
1479
1480        sleep(Duration::from_millis(50)).await;
1481
1482        let url = Url::parse(&format!("http://{}", addr)).unwrap();
1483        let start = Instant::now();
1484
1485        let _result = timeout(
1486            config.connect_timeout + Duration::from_millis(50),
1487            manager.acquire(&url),
1488        )
1489        .await;
1490        let elapsed = start.elapsed();
1491
1492        assert!(
1493            elapsed < config.connect_timeout + Duration::from_millis(300),
1494            "耗时过长: {:.2}ms",
1495            elapsed.as_millis()
1496        );
1497
1498        manager.cleanup().await;
1499        server_handle.abort();
1500    }
1501
1502    #[tokio::test]
1503    async fn test_max_connections_limit() {
1504        let config = HttpConfig {
1505            max_connections: 2,
1506            connect_timeout: Duration::from_millis(500),
1507            read_timeout: Duration::from_millis(1000),
1508            write_timeout: Duration::from_millis(1000),
1509            idle_timeout: Duration::from_secs(60),
1510        };
1511        let mut manager = HttpConnectionManager::new(&config);
1512
1513        let (addr, _server_handle) = start_test_server(|mut stream| {
1514            tokio::spawn(async move {
1515                let response = "HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nOK";
1516                stream.write_all(response.as_bytes()).await.unwrap();
1517                sleep(Duration::from_secs(10)).await;
1518            });
1519        })
1520        .await;
1521
1522        sleep(Duration::from_millis(100)).await;
1523
1524        let url = Url::parse(&format!("http://{}", addr)).unwrap();
1525
1526        let conn1 = manager.acquire(&url).await.unwrap();
1527        assert!(manager.active_count() >= 1);
1528
1529        let conn2 = manager.acquire(&url).await.unwrap();
1530        assert!(manager.active_count() >= 2);
1531
1532        // 尝试获取第三个连接（应该因达到限制而失败）
1533        let result = manager.acquire(&url).await;
1534        assert!(result.is_err(), "超过最大连接数限制时应返回错误");
1535
1536        // 验证错误类型
1537        if let Err(e) = result {
1538            match &e {
1539                Aria2Error::Recoverable(_) => {}
1540                other => panic!("期望 Recoverable 错误，得到: {:?}", other),
1541            }
1542        }
1543
1544        // 归还一个连接后，应该可以重新获取（如果连接池复用正常工作）
1545        manager.release(conn1.id).await;
1546        // 注意：由于连接可能仍在池中被计费，这里我们只验证不会 panic
1547        match manager.acquire(&url).await {
1548            Ok(conn3) => {
1549                println!("✓ 归还后成功获取新连接: id={}", conn3.id);
1550                manager.release(conn3.id).await;
1551            }
1552            Err(e) => {
1553                println!("⚠ 归还后获取失败（可能是连接复用限制）: {}", e);
1554                // 这也是可接受的行为
1555            }
1556        }
1557
1558        manager.release(conn2.id).await;
1559        manager.cleanup().await;
1560    }
1561
1562    // ==================== Cookie Jar Integration Tests (J4) ====================
1563
1564    #[test]
1565    fn test_cookie_jar_initially_none() {
1566        let mut manager = HttpConnectionManager::new(&create_test_config());
1567        assert!(manager.cookie_jar().is_none());
1568        assert!(manager.cookie_jar_mut().is_none());
1569
1570        // Attaching cookies without a jar should return None
1571        let url = Url::parse("https://example.com/").unwrap();
1572        assert!(manager.attach_cookies_to_request(&url).is_none());
1573    }
1574
1575    #[test]
1576    fn test_set_and_get_cookie_jar() {
1577        let mut manager = HttpConnectionManager::new(&create_test_config());
1578
1579        // Initially no jar
1580        assert!(manager.cookie_jar().is_none());
1581
1582        // Set a cookie jar
1583        let jar = CookieJar::new();
1584        manager.set_cookie_jar(Some(jar));
1585        assert!(manager.cookie_jar().is_some());
1586
1587        // Clear it
1588        manager.set_cookie_jar(None);
1589        assert!(manager.cookie_jar().is_none());
1590    }
1591
1592    #[test]
1593    fn test_attach_cookies_to_request() {
1594        let mut manager = HttpConnectionManager::new(&create_test_config());
1595
1596        // Create jar and add cookies
1597        let mut jar = CookieJar::new();
1598        jar.store(JarCookie::new("session_id", "abc123", "example.com"));
1599        jar.store(JarCookie::new("theme", "dark", "example.com"));
1600        manager.set_cookie_jar(Some(jar));
1601
1602        // Attach cookies for example.com URL
1603        let url = Url::parse("http://example.com/api/data").unwrap();
1604        let header = manager.attach_cookies_to_request(&url);
1605        assert!(header.is_some(), "Should return Some with matching cookies");
1606        let hdr = header.unwrap();
1607        assert!(
1608            hdr.contains("session_id=abc123"),
1609            "Header should contain session_id cookie: {}",
1610            hdr
1611        );
1612        assert!(
1613            hdr.contains("theme=dark"),
1614            "Header should contain theme cookie: {}",
1615            hdr
1616        );
1617
1618        // No cookies for different domain
1619        let url2 = Url::parse("http://other.com/").unwrap();
1620        let header2 = manager.attach_cookies_to_request(&url2);
1621        assert!(header2.is_none(), "No cookies should match other domain");
1622    }
1623
1624    #[test]
1625    fn test_extract_cookies_from_response() {
1626        let mut manager = HttpConnectionManager::new(&create_test_config());
1627        manager.set_cookie_jar(Some(CookieJar::new()));
1628
1629        // Simulate response headers with Set-Cookie
1630        let response_headers = vec![
1631            (
1632                "Set-Cookie".to_string(),
1633                "session=xyz789; Domain=example.com; Path=/".to_string(),
1634            ),
1635            (
1636                "Set-Cookie".to_string(),
1637                "prefs=en-US; Domain=example.com; Path=/; Secure; HttpOnly".to_string(),
1638            ),
1639            ("Content-Type".to_string(), "text/html".to_string()), // Non-cookie header
1640        ];
1641
1642        let url = Url::parse("https://example.com/login").unwrap();
1643        let count = manager.extract_cookies_from_response(&response_headers, &url);
1644
1645        assert_eq!(count, 2, "Should extract exactly 2 cookies");
1646
1647        // Verify cookies were stored
1648        let jar = manager.cookie_jar().as_ref().unwrap();
1649        assert_eq!(jar.len(), 2, "Jar should contain 2 stored cookies");
1650
1651        // Verify we can retrieve them
1652        let cookies = jar.get_cookies_for_url("https://example.com/", true);
1653        assert_eq!(cookies.len(), 2);
1654
1655        let names: Vec<&str> = cookies.iter().map(|c| c.name.as_str()).collect();
1656        assert!(names.contains(&"session"));
1657        assert!(names.contains(&"prefs"));
1658
1659        // Verify Secure flag was parsed correctly
1660        let prefs_cookie = cookies.iter().find(|c| c.name == "prefs").unwrap();
1661        assert!(prefs_cookie.secure, "prefs cookie should be marked secure");
1662        assert!(
1663            prefs_cookie.http_only,
1664            "prefs cookie should be marked http_only"
1665        );
1666    }
1667
1668    #[test]
1669    fn test_extract_cookies_no_jar_returns_zero() {
1670        let mut manager = HttpConnectionManager::new(&create_test_config());
1671        // No cookie jar set
1672
1673        let headers = vec![("Set-Cookie".to_string(), "test=val".to_string())];
1674        let url = Url::parse("http://example.com/").unwrap();
1675        let count = manager.extract_cookies_from_response(&headers, &url);
1676
1677        assert_eq!(count, 0, "Should return 0 when no jar is set");
1678    }
1679
1680    #[test]
1681    fn test_extract_cookies_invalid_header_skipped() {
1682        let mut manager = HttpConnectionManager::new(&create_test_config());
1683        manager.set_cookie_jar(Some(CookieJar::new()));
1684
1685        // Mix of valid and invalid Set-Cookie headers
1686        let headers = vec![
1687            (
1688                "Set-Cookie".to_string(),
1689                "valid=test_value; Domain=x.com".to_string(),
1690            ),
1691            ("Set-Cookie".to_string(), "no-equal-sign".to_string()), // Invalid format
1692            ("Set-Cookie".to_string(), "".to_string()),              // Empty - invalid
1693        ];
1694
1695        let url = Url::parse("http://x.com/").unwrap();
1696        let count = manager.extract_cookies_from_response(&headers, &url);
1697
1698        assert_eq!(count, 1, "Only 1 valid cookie should be extracted");
1699
1700        let jar = manager.cookie_jar().as_ref().unwrap();
1701        assert_eq!(jar.len(), 1);
1702        let cookies = jar.get_cookies_for_url("http://x.com/", false);
1703        assert_eq!(cookies[0].name, "valid");
1704    }
1705
1706    #[test]
1707    fn test_debug_format_includes_cookie_jar() {
1708        let mut manager = HttpConnectionManager::new(&create_test_config());
1709        let debug_str = format!("{:?}", manager);
1710        assert!(!debug_str.contains("cookie_jar_set: true"));
1711
1712        manager.set_cookie_jar(Some(CookieJar::new()));
1713        let debug_str_with_jar = format!("{:?}", manager);
1714        assert!(
1715            debug_str_with_jar.contains("cookie_jar_set: true"),
1716            "Debug output should show cookie_jar is set: {}",
1717            debug_str_with_jar
1718        );
1719    }
1720
1721    #[test]
1722    fn test_secure_cookie_not_sent_over_http() {
1723        let mut manager = HttpConnectionManager::new(&create_test_config());
1724        let mut jar = CookieJar::new();
1725
1726        // Add a secure-only cookie
1727        let mut secure_cookie = JarCookie::new("token", "secret", "secure.example.com");
1728        secure_cookie.secure = true;
1729        jar.store(secure_cookie);
1730
1731        manager.set_cookie_jar(Some(jar));
1732
1733        // Over HTTP — should NOT get the secure cookie
1734        let url_http = Url::parse("http://secure.example.com/api").unwrap();
1735        let header_http = manager.attach_cookies_to_request(&url_http);
1736        assert!(
1737            header_http.is_none(),
1738            "Secure cookie must not be sent over HTTP"
1739        );
1740
1741        // Over HTTPS — SHOULD get the secure cookie
1742        let url_https = Url::parse("https://secure.example.com/api").unwrap();
1743        let header_https = manager.attach_cookies_to_request(&url_https);
1744        assert!(
1745            header_https.is_some(),
1746            "Secure cookie should be sent over HTTPS"
1747        );
1748        assert!(
1749            header_https.unwrap().contains("token=secret"),
1750            "Header should contain the secure token cookie"
1751        );
1752    }
1753}
aria2_core/http/connection.rs

aria2_core/http/
connection.rs
