aria2_core/engine/
bt_peer_interaction.rs

1//! BT Peer Interaction Manager - Peer connection and initialization
2//!
3//! This module manages the interaction with BitTorrent peers, including:
4//! - Connection establishment (plain and encrypted)
5//! - Initial handshake and bitfield exchange
6//! - Waiting for unchoke messages
7//! - Peer statistics tracking
8//!
9//! # Architecture Reference
10//!
11//! Based on original aria2 C++ structure:
12//! - `src/PeerConnection.cc/h` - Peer connection management
13//! - `src/PeerInteractionCommand.h` - Interaction logic
14//! - `src/BtSetup.cc/h` - BT setup and initialization
15
16use std::time::Duration;
17
18use crate::constants;
19use crate::engine::bt_peer_connection::BtPeerConn;
20use crate::error::{Aria2Error, RecoverableError, Result};
21use tracing::{debug, error, info, warn};
22
23/// Delay between peer connection setup and message reading (milliseconds)
24pub const PEER_CONNECTION_DELAY_MS: u64 = constants::BT_PEER_CONNECTION_DELAY_MS;
25
26/// Maximum attempts to wait for unchoke from a peer
27pub const MAX_UNCHOKE_WAIT_ATTEMPTS: u32 = constants::BT_MAX_UNCHOKE_WAIT_ATTEMPTS as u32;
28
29/// Timeout for each message read from peer (seconds)
30pub const PEER_MESSAGE_TIMEOUT_SECS: u64 = constants::BT_PEER_MESSAGE_TIMEOUT_SECS;
31
32/// Result of peer connection attempt
33pub struct PeerConnectionResult {
34    /// Successfully connected peers
35    pub connections: Vec<BtPeerConn>,
36    /// Number of failed connections
37    pub failed_count: usize,
38}
39
40/// BT Peer Interaction Manager
41///
42/// Handles the lifecycle of peer connections from initial connection
43/// through the handshake phase until they're ready for data transfer.
44pub struct BtPeerInteraction;
45
46impl BtPeerInteraction {
47    /// Connect to multiple peers with automatic fallback strategies
48    ///
49    /// Attempts to connect to all provided peer addresses using:
50    /// 1. MSE encryption if required or forced
51    /// 2. Plain connection as fallback
52    ///
53    /// For each successful connection:
54    /// - Sends initial unchoke and interested messages
55    /// - Exchanges bitfields
56    /// - Waits for unchoke from the peer
57    ///
58    /// # Arguments
59    /// * `peer_addrs` - List of peer addresses to connect to
60    /// * `info_hash_raw` - Torrent info hash for handshake
61    /// * `num_pieces` - Total number of pieces (for bitfield size)
62    /// * `require_crypto` - Whether to require encrypted connections
63    /// * `force_encrypt` - Whether to force encryption (fallback to plain)
64    ///
65    /// # Returns
66    /// * `PeerConnectionResult` containing connected peers and failure count
67    pub async fn connect_to_peers(
68        peer_addrs: &[aria2_protocol::bittorrent::peer::connection::PeerAddr],
69        info_hash_raw: &[u8; 20],
70        num_pieces: u32,
71        require_crypto: bool,
72        force_encrypt: bool,
73    ) -> Result<PeerConnectionResult> {
74        info!("[BT] Connecting to {} peers...", peer_addrs.len());
75
76        let mut active_connections: Vec<BtPeerConn> = Vec::new();
77        let mut failed_count = 0usize;
78
79        for addr in peer_addrs {
80            debug!("[BT] Connecting to peer {}:{}", addr.ip, addr.port);
81
82            let conn_result =
83                Self::connect_single_peer(addr, info_hash_raw, require_crypto, force_encrypt).await;
84
85            match conn_result {
86                Ok(mut conn) => {
87                    info!(
88                        "[BT] Connected to peer {}:{} (encrypted={})",
89                        addr.ip,
90                        addr.port,
91                        conn.is_encrypted()
92                    );
93
94                    // Initialize the connection
95                    if let Err(e) = Self::initialize_connection(&mut conn, num_pieces).await {
96                        warn!("[BT] Failed to initialize peer {}: {}", addr.ip, e);
97                        failed_count += 1;
98                        continue;
99                    }
100
101                    // Wait for unchoke
102                    match Self::wait_for_unchoke(&mut conn, addr).await {
103                        Ok(()) => {
104                            active_connections.push(conn);
105                        }
106                        Err(e) => {
107                            warn!("[BT] No unchoke from peer {}: {}", addr.ip, e);
108                            // Still add the connection even without unchoke
109                            // (it might unchoke later)
110                            active_connections.push(conn);
111                        }
112                    }
113                }
114                Err(e) => {
115                    error!("[BT] Failed to connect peer {}: {}", addr.ip, e);
116                    failed_count += 1;
117                    continue;
118                }
119            }
120        }
121
122        info!("[BT] Active connections: {}", active_connections.len());
123
124        if active_connections.is_empty() {
125            return Err(Aria2Error::Recoverable(
126                RecoverableError::TemporaryNetworkFailure {
127                    message: "All peer connections failed".into(),
128                },
129            ));
130        }
131
132        Ok(PeerConnectionResult {
133            connections: active_connections,
134            failed_count,
135        })
136    }
137
138    /// Connect to a single peer with encryption fallback logic
139    async fn connect_single_peer(
140        addr: &aria2_protocol::bittorrent::peer::connection::PeerAddr,
141        info_hash_raw: &[u8; 20],
142        require_crypto: bool,
143        force_encrypt: bool,
144    ) -> Result<BtPeerConn> {
145        if force_encrypt || require_crypto {
146            // Try MSE encrypted connection
147            BtPeerConn::connect_mse(addr, info_hash_raw, require_crypto).await
148        } else {
149            // Try MSE first, fall back to plain
150            match BtPeerConn::connect_mse(addr, info_hash_raw, false).await {
151                Ok(conn) => Ok(conn),
152                Err(_) => {
153                    debug!("[BT] MSE failed, trying plain connection");
154                    BtPeerConn::connect_plain(addr, info_hash_raw).await
155                }
156            }
157        }
158    }
159
160    /// Initialize a newly established connection
161    ///
162    /// Sends initial protocol messages:
163    /// - Unchoke (we allow them to request from us)
164    /// - Interested (we want to download from them)
165    /// - Bitfield (our current piece possession status)
166    async fn initialize_connection(conn: &mut BtPeerConn, num_pieces: u32) -> Result<()> {
167        // Send initial messages
168        conn.send_unchoke().await?;
169        conn.send_interested().await?;
170
171        // Send empty bitfield (we have nothing yet)
172        let bf_len = (num_pieces as usize).div_ceil(8);
173        let empty_bf = vec![0u8; bf_len];
174        conn.send_bitfield(empty_bf).await?;
175
176        // Small delay to allow processing
177        tokio::time::sleep(Duration::from_millis(PEER_CONNECTION_DELAY_MS)).await;
178
179        Ok(())
180    }
181
182    /// Wait for an unchoke message from a peer
183    ///
184    /// Polls the connection for messages until we receive an Unchoke
185    /// or hit the timeout/attempts limit.
186    async fn wait_for_unchoke(
187        conn: &mut BtPeerConn,
188        addr: &aria2_protocol::bittorrent::peer::connection::PeerAddr,
189    ) -> Result<()> {
190        debug!("[BT] Waiting for unchoke from {}:{}", addr.ip, addr.port);
191
192        for _ in 0..MAX_UNCHOKE_WAIT_ATTEMPTS {
193            match tokio::time::timeout(
194                Duration::from_secs(PEER_MESSAGE_TIMEOUT_SECS),
195                conn.read_message(),
196            )
197            .await
198            {
199                Ok(Ok(Some(msg))) => {
200                    use aria2_protocol::bittorrent::message::types::BtMessage;
201                    if matches!(msg, BtMessage::Unchoke) {
202                        info!("[BT] Got unchoke from {}:{}", addr.ip, addr.port);
203                        return Ok(());
204                    }
205                    debug!("[BT] Got message while waiting for unchoke: {:?}", msg);
206                }
207                Ok(Ok(None)) => {
208                    warn!("[BT] EOF from peer while waiting for unchoke");
209                    return Err(Aria2Error::Recoverable(
210                        RecoverableError::TemporaryNetworkFailure {
211                            message: "Peer closed connection".into(),
212                        },
213                    ));
214                }
215                Ok(Err(e)) => {
216                    error!("[BT] Error reading from peer: {}", e);
217                    return Err(Aria2Error::Recoverable(
218                        RecoverableError::TemporaryNetworkFailure {
219                            message: format!("Read error: {}", e),
220                        },
221                    ));
222                }
223                Err(_) => {
224                    debug!("[BT] Timeout reading from peer, retrying...");
225                }
226            }
227        }
228
229        warn!(
230            "[BT] Did not receive unchoke from {}:{} after {} attempts",
231            addr.ip, addr.port, MAX_UNCHOKE_WAIT_ATTEMPTS
232        );
233        Ok(()) // Continue anyway, might get unchoke later
234    }
235
236    /// Broadcast a HAVE message to all connected peers
237    ///
238    /// Notifies all peers that we have completed downloading a piece.
239    ///
240    /// # Arguments
241    /// * `connections` - Mutable slice of active peer connections
242    /// * `piece_index` - Index of the completed piece
243    pub async fn broadcast_have(connections: &mut [BtPeerConn], piece_index: u32) {
244        for conn in connections.iter_mut() {
245            if let Err(e) = conn.send_have(piece_index).await {
246                warn!("[BT] Failed to send HAVE to peer: {}", e);
247            }
248        }
249    }
250
251    /// Initialize peer bitfield tracker for all connections
252    ///
253    /// Sets up tracking of which pieces each peer claims to have.
254    ///
255    /// # Arguments
256    /// * `connections` - Slice of active peer connections
257    /// * `num_pieces` - Total number of pieces in the torrent
258    /// * `peer_tracker` - Mutable reference to the peer bitfield tracker
259    pub fn initialize_peer_tracking(
260        connections: &[BtPeerConn],
261        num_pieces: u32,
262        peer_tracker: &mut aria2_protocol::bittorrent::piece::peer_tracker::PeerBitfieldTracker,
263    ) {
264        for (i, _conn) in connections.iter().enumerate() {
265            let empty_bf = vec![0xFFu8; (num_pieces as usize).div_ceil(8)];
266            peer_tracker.update_peer_bitfield(&format!("peer_{}", i), &empty_bf);
267        }
268
269        debug!(
270            "[BT] Initialized peer tracking for {} peers",
271            connections.len()
272        );
273    }
274
275    /// Clean up peer connections (drop them properly)
276    ///
277    /// Ensures all connections are properly closed.
278    ///
279    /// # Arguments
280    /// * `connections` - Mutable slice of peer connections to close
281    pub fn cleanup_connections(connections: &mut [BtPeerConn]) {
282        for conn in connections.iter_mut() {
283            let _ = conn;
284        }
285        debug!("[BT] Cleaned up {} connections", connections.len());
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292
293    #[test]
294    fn test_constants_are_reasonable() {
295        const _: () = {
296            assert!(PEER_CONNECTION_DELAY_MS >= 10);
297            assert!(PEER_CONNECTION_DELAY_MS <= 1000);
298            assert!(MAX_UNCHOKE_WAIT_ATTEMPTS >= 10);
299            assert!(MAX_UNCHOKE_WAIT_ATTEMPTS <= 100);
300            assert!(PEER_MESSAGE_TIMEOUT_SECS >= 1);
301            assert!(PEER_MESSAGE_TIMEOUT_SECS <= 30);
302        };
303    }
304
305    #[test]
306    fn test_peer_connection_result_default() {
307        let result = PeerConnectionResult {
308            connections: Vec::new(),
309            failed_count: 0,
310        };
311        assert!(result.connections.is_empty());
312        assert_eq!(result.failed_count, 0);
313    }
314}
aria2_core/engine/bt_peer_interaction.rs

aria2_core/engine/
bt_peer_interaction.rs
