aria2_core/engine/
bt_peer_connection.rs

1use std::collections::HashSet;
2use std::sync::Arc;
3use tokio::sync::Mutex;
4
5use crate::constants;
6use crate::error::{Aria2Error, FatalError, RecoverableError, Result};
7
8#[allow(clippy::large_enum_variant)]
9pub(crate) enum InnerConnection {
10    Plain(aria2_protocol::bittorrent::peer::connection::PeerConnection),
11    Encrypted(aria2_protocol::bittorrent::peer::encrypted_connection::EncryptedConnection),
12    Utp(UtpPeerConnection),
13}
14
15/// uTP peer connection wrapper
16///
17/// Wraps a uTP stream for BitTorrent peer communication.
18/// Provides the same interface as TCP connections but uses UDP-based uTP protocol.
19pub struct UtpPeerConnection {
20    /// uTP socket reference (shared among multiple connections)
21    socket: Arc<Mutex<aria2_protocol::bittorrent::utp::UtpSocket>>,
22    /// Connection ID within the socket
23    conn_id: u16,
24    /// Info hash for the torrent
25    info_hash: [u8; 20],
26    /// Whether handshake is complete
27    handshake_complete: bool,
28    /// Receive buffer for partial messages
29    recv_buffer: Vec<u8>,
30}
31
32impl UtpPeerConnection {
33    /// Create a new uTP peer connection
34    pub fn new(
35        socket: Arc<Mutex<aria2_protocol::bittorrent::utp::UtpSocket>>,
36        conn_id: u16,
37        info_hash: [u8; 20],
38    ) -> Self {
39        Self {
40            socket,
41            conn_id,
42            info_hash,
43            handshake_complete: false,
44            recv_buffer: Vec::new(),
45        }
46    }
47
48    /// Connect to a remote peer via uTP
49    pub async fn connect(addr: std::net::SocketAddr, info_hash: &[u8; 20]) -> Result<Self> {
50        // Create a new uTP socket for this connection
51        let socket = aria2_protocol::bittorrent::utp::UtpSocket::bind_any()
52            .map_err(|e| Aria2Error::Fatal(FatalError::Config(e.to_string())))?;
53
54        let socket = Arc::new(Mutex::new(socket));
55
56        // Initiate uTP connection
57        let conn_id = {
58            let mut sock = socket.lock().await;
59            sock.connect(addr)
60                .map_err(|e| Aria2Error::Fatal(FatalError::Config(e.to_string())))?
61        };
62
63        Ok(Self {
64            socket,
65            conn_id,
66            info_hash: *info_hash,
67            handshake_complete: false,
68            recv_buffer: Vec::new(),
69        })
70    }
71
72    /// Get the connection ID
73    pub fn conn_id(&self) -> u16 {
74        self.conn_id
75    }
76
77    /// Check if connection is established
78    pub fn is_connected(&self) -> bool {
79        // This is a simplified check - in reality we'd need to check state
80        self.handshake_complete
81    }
82
83    /// Perform BitTorrent handshake over uTP
84    pub async fn perform_handshake(&mut self) -> Result<()> {
85        use aria2_protocol::bittorrent::message::handshake::Handshake;
86        use aria2_protocol::bittorrent::peer::id::generate_peer_id;
87
88        // Generate proper peer_id following BEP 20 format (-AR0001- prefix)
89        let peer_id = generate_peer_id();
90        let handshake = Handshake::new(&self.info_hash, &peer_id);
91        let handshake_bytes = handshake.to_bytes();
92
93        // Send handshake
94        {
95            let mut socket = self.socket.lock().await;
96            socket.send(self.conn_id, &handshake_bytes).map_err(|e| {
97                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure {
98                    message: e.to_string(),
99                })
100            })?;
101        }
102
103        // Receive handshake response
104        let mut response_buf = vec![0u8; 68]; // Handshake is 68 bytes
105        let len = {
106            let mut socket = self.socket.lock().await;
107            socket.recv(self.conn_id, &mut response_buf).map_err(|e| {
108                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure {
109                    message: e.to_string(),
110                })
111            })?
112        };
113
114        if len < 68 {
115            return Err(Aria2Error::Fatal(FatalError::Config(
116                "Handshake response too short".to_string(),
117            )));
118        }
119
120        // Parse handshake
121        let response = Handshake::parse(&response_buf[..len])
122            .map_err(|e| Aria2Error::Fatal(FatalError::Config(e)))?;
123
124        // Verify info hash
125        if response.info_hash != self.info_hash {
126            return Err(Aria2Error::Fatal(FatalError::Config(
127                "Info hash mismatch".to_string(),
128            )));
129        }
130
131        self.handshake_complete = true;
132        Ok(())
133    }
134
135    /// Send a BitTorrent message
136    pub async fn send_message(&mut self, message: &[u8]) -> Result<()> {
137        let mut socket = self.socket.lock().await;
138        socket.send(self.conn_id, message).map_err(|e| {
139            Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure {
140                message: e.to_string(),
141            })
142        })?;
143        Ok(())
144    }
145
146    /// Receive a BitTorrent message
147    pub async fn recv_message(&mut self) -> Result<Option<Vec<u8>>> {
148        // Try to receive data
149        let mut buf = vec![0u8; constants::BT_RECEIVE_BUFFER_SIZE];
150        let len = {
151            let mut socket = self.socket.lock().await;
152            match socket.recv(self.conn_id, &mut buf) {
153                Ok(0) => return Ok(None), // No data available
154                Ok(len) => len,
155                Err(e) => {
156                    return Err(Aria2Error::Recoverable(
157                        RecoverableError::TemporaryNetworkFailure {
158                            message: e.to_string(),
159                        },
160                    ));
161                }
162            }
163        };
164
165        // Append to receive buffer
166        self.recv_buffer.extend_from_slice(&buf[..len]);
167
168        // Try to parse a complete message
169        if self.recv_buffer.len() >= 4 {
170            // Read message length (4 bytes, big-endian)
171            let msg_len = u32::from_be_bytes([
172                self.recv_buffer[0],
173                self.recv_buffer[1],
174                self.recv_buffer[2],
175                self.recv_buffer[3],
176            ]) as usize;
177
178            // Check if we have the complete message
179            if self.recv_buffer.len() >= 4 + msg_len {
180                // Extract message
181                let message = self.recv_buffer[4..4 + msg_len].to_vec();
182                self.recv_buffer = self.recv_buffer[4 + msg_len..].to_vec();
183                return Ok(Some(message));
184            }
185        }
186
187        Ok(None)
188    }
189
190    /// Close the connection
191    pub async fn close(&mut self) -> Result<()> {
192        let mut socket = self.socket.lock().await;
193        socket.close_connection(self.conn_id).map_err(|e| {
194            Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure {
195                message: e.to_string(),
196            })
197        })?;
198        Ok(())
199    }
200
201    /// Get connection statistics
202    pub async fn stats(&self) -> Result<aria2_protocol::bittorrent::utp::ConnectionStats> {
203        let socket = self.socket.lock().await;
204        socket
205            .connection_stats(self.conn_id)
206            .map_err(|e| Aria2Error::Fatal(FatalError::Config(e.to_string())))
207    }
208}
209
210/// Peer connection abstraction that supports both plain and encrypted (MSE) connections.
211///
212/// This mirrors the original aria2 C++ architecture where connection management
213/// is separated from the download command logic (see BtRuntime in original).
214/// Now also supports uTP (UDP-based) connections per BEP 29.
215pub struct BtPeerConn {
216    pub(crate) inner: InnerConnection,
217    /// Set of piece indices for which the peer has sent an AllowedFast message.
218    /// Pieces in this set can be requested even when the peer is choked.
219    pub allowed_fast: HashSet<u32>,
220    /// Connection type (TCP or uTP)
221    pub connection_type: ConnectionType,
222}
223
224/// Type of peer connection
225#[derive(Debug, Clone, Copy, PartialEq, Eq)]
226pub enum ConnectionType {
227    /// Standard TCP connection
228    Tcp,
229    /// uTP (UDP-based) connection
230    Utp,
231}
232
233impl BtPeerConn {
234    /// Connect via MSE (Message Stream Encryption) over TCP
235    pub async fn connect_mse(
236        addr: &aria2_protocol::bittorrent::peer::connection::PeerAddr,
237        info_hash: &[u8; 20],
238        require_encryption: bool,
239    ) -> Result<Self> {
240        match aria2_protocol::bittorrent::peer::encrypted_connection::EncryptedConnection::connect_with_mse(addr, info_hash, require_encryption).await {
241            Ok(conn) => Ok(Self {
242                inner: InnerConnection::Encrypted(conn),
243                allowed_fast: HashSet::new(),
244                connection_type: ConnectionType::Tcp,
245            }),
246            Err(e) => Err(Aria2Error::Fatal(FatalError::Config(e))),
247        }
248    }
249
250    /// Connect via plain TCP
251    pub async fn connect_plain(
252        addr: &aria2_protocol::bittorrent::peer::connection::PeerAddr,
253        info_hash: &[u8; 20],
254    ) -> Result<Self> {
255        match aria2_protocol::bittorrent::peer::connection::PeerConnection::connect(addr, info_hash)
256            .await
257        {
258            Ok(conn) => Ok(Self {
259                inner: InnerConnection::Plain(conn),
260                allowed_fast: HashSet::new(),
261                connection_type: ConnectionType::Tcp,
262            }),
263            Err(e) => Err(Aria2Error::Fatal(FatalError::Config(e))),
264        }
265    }
266
267    /// Connect via uTP (Micro Transport Protocol)
268    ///
269    /// uTP is a UDP-based transport protocol that provides:
270    /// - Reliable, ordered delivery
271    /// - LEDBAT congestion control (low priority, background traffic)
272    /// - Better performance on congested networks
273    /// - NAT traversal benefits
274    pub async fn connect_utp(addr: std::net::SocketAddr, info_hash: &[u8; 20]) -> Result<Self> {
275        let utp_conn = UtpPeerConnection::connect(addr, info_hash).await?;
276
277        Ok(Self {
278            inner: InnerConnection::Utp(utp_conn),
279            allowed_fast: HashSet::new(),
280            connection_type: ConnectionType::Utp,
281        })
282    }
283
284    /// Create a uTP connection from an existing socket
285    ///
286    /// Used when accepting incoming uTP connections
287    pub fn from_utp_socket(
288        socket: Arc<Mutex<aria2_protocol::bittorrent::utp::UtpSocket>>,
289        conn_id: u16,
290        info_hash: &[u8; 20],
291    ) -> Self {
292        let utp_conn = UtpPeerConnection::new(socket, conn_id, *info_hash);
293
294        Self {
295            inner: InnerConnection::Utp(utp_conn),
296            allowed_fast: HashSet::new(),
297            connection_type: ConnectionType::Utp,
298        }
299    }
300
301    /// Get the connection type
302    pub fn connection_type(&self) -> ConnectionType {
303        self.connection_type
304    }
305
306    /// Check if this is a uTP connection
307    pub fn is_utp(&self) -> bool {
308        self.connection_type == ConnectionType::Utp
309    }
310
311    /// Add a piece index to the AllowedFast set.
312    ///
313    /// Called when an AllowedFast message is received from this peer.
314    /// Pieces in the allowed_fast set can be requested even when the peer
315    /// is choked (BEP 6 / Fast Extension).
316    pub fn add_allowed_fast(&mut self, index: u32) {
317        self.allowed_fast.insert(index);
318    }
319
320    /// Check whether a piece index is in the AllowedFast set.
321    ///
322    /// Returns true if the peer has granted fast access to this piece,
323    /// meaning a Request can be sent even while the peer is choked.
324    pub fn is_allowed_fast(&self, index: u32) -> bool {
325        self.allowed_fast.contains(&index)
326    }
327
328    /// Get a reference to the full AllowedFast set
329    ///
330    /// Returns all piece indices that this peer has allowed us to request
331    /// via BEP 6 Fast Extension, even when choked.
332    pub fn allowed_fast_set(&self) -> &HashSet<u32> {
333        &self.allowed_fast
334    }
335
336    pub async fn send_unchoke(&mut self) -> Result<()> {
337        match &mut self.inner {
338            InnerConnection::Plain(c) => c.send_unchoke().await.map_err(|e| {
339                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
340            }),
341            InnerConnection::Encrypted(c) => c.send_unchoke().await.map_err(|e| {
342                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
343            }),
344            InnerConnection::Utp(c) => {
345                // uTP uses the same BitTorrent protocol messages
346                use aria2_protocol::bittorrent::message::serializer::serialize;
347                use aria2_protocol::bittorrent::message::types::BtMessage;
348                let msg = BtMessage::Unchoke;
349                c.send_message(&serialize(&msg)).await
350            }
351        }
352    }
353
354    pub async fn send_choke(&mut self) -> Result<()> {
355        match &mut self.inner {
356            InnerConnection::Plain(c) => c.send_choke().await.map_err(|e| {
357                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
358            }),
359            InnerConnection::Encrypted(c) => c.send_choke().await.map_err(|e| {
360                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
361            }),
362            InnerConnection::Utp(c) => {
363                use aria2_protocol::bittorrent::message::serializer::serialize;
364                use aria2_protocol::bittorrent::message::types::BtMessage;
365                let msg = BtMessage::Choke;
366                c.send_message(&serialize(&msg)).await
367            }
368        }
369    }
370
371    pub async fn send_interested(&mut self) -> Result<()> {
372        match &mut self.inner {
373            InnerConnection::Plain(c) => c.send_interested().await.map_err(|e| {
374                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
375            }),
376            InnerConnection::Encrypted(c) => c.send_interested().await.map_err(|e| {
377                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
378            }),
379            InnerConnection::Utp(c) => {
380                use aria2_protocol::bittorrent::message::serializer::serialize;
381                use aria2_protocol::bittorrent::message::types::BtMessage;
382                let msg = BtMessage::Interested;
383                c.send_message(&serialize(&msg)).await
384            }
385        }
386    }
387
388    pub async fn send_not_interested(&mut self) -> Result<()> {
389        match &mut self.inner {
390            InnerConnection::Plain(c) => c.send_not_interested().await.map_err(|e| {
391                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
392            }),
393            InnerConnection::Encrypted(c) => c.send_not_interested().await.map_err(|e| {
394                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
395            }),
396            InnerConnection::Utp(c) => {
397                use aria2_protocol::bittorrent::message::serializer::serialize;
398                use aria2_protocol::bittorrent::message::types::BtMessage;
399                let msg = BtMessage::NotInterested;
400                c.send_message(&serialize(&msg)).await
401            }
402        }
403    }
404
405    pub async fn send_have(&mut self, piece_index: u32) -> Result<()> {
406        match &mut self.inner {
407            InnerConnection::Plain(c) => c.send_have(piece_index).await.map_err(|e| {
408                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
409            }),
410            InnerConnection::Encrypted(c) => c.send_have(piece_index).await.map_err(|e| {
411                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
412            }),
413            InnerConnection::Utp(c) => {
414                use aria2_protocol::bittorrent::message::serializer::serialize;
415                use aria2_protocol::bittorrent::message::types::BtMessage;
416                let msg = BtMessage::Have { piece_index };
417                c.send_message(&serialize(&msg)).await
418            }
419        }
420    }
421
422    pub async fn send_request(
423        &mut self,
424        req: aria2_protocol::bittorrent::message::types::PieceBlockRequest,
425    ) -> Result<()> {
426        match &mut self.inner {
427            InnerConnection::Plain(c) => c.send_request(req).await.map_err(|e| {
428                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
429            }),
430            InnerConnection::Encrypted(c) => c.send_request(req).await.map_err(|e| {
431                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
432            }),
433            InnerConnection::Utp(c) => {
434                use aria2_protocol::bittorrent::message::serializer::serialize;
435                use aria2_protocol::bittorrent::message::types::{BtMessage, PieceBlockRequest};
436                let msg = BtMessage::Request {
437                    request: PieceBlockRequest::new(req.index, req.begin, req.length),
438                };
439                c.send_message(&serialize(&msg)).await
440            }
441        }
442    }
443
444    pub async fn send_cancel(
445        &mut self,
446        req: &aria2_protocol::bittorrent::message::types::PieceBlockRequest,
447    ) -> Result<()> {
448        match &mut self.inner {
449            InnerConnection::Plain(c) => c.send_cancel(req).await.map_err(|e| {
450                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
451            }),
452            InnerConnection::Encrypted(c) => c.send_cancel(req).await.map_err(|e| {
453                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
454            }),
455            InnerConnection::Utp(c) => {
456                use aria2_protocol::bittorrent::message::serializer::serialize;
457                use aria2_protocol::bittorrent::message::types::{BtMessage, PieceBlockRequest};
458                let msg = BtMessage::Cancel {
459                    request: PieceBlockRequest::new(req.index, req.begin, req.length),
460                };
461                c.send_message(&serialize(&msg)).await
462            }
463        }
464    }
465
466    pub async fn send_bitfield(&mut self, bitfield: Vec<u8>) -> Result<()> {
467        match &mut self.inner {
468            InnerConnection::Plain(c) => c.send_bitfield(bitfield).await.map_err(|e| {
469                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
470            }),
471            InnerConnection::Encrypted(c) => c.send_bitfield(bitfield).await.map_err(|e| {
472                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
473            }),
474            InnerConnection::Utp(c) => {
475                use aria2_protocol::bittorrent::message::serializer::serialize;
476                use aria2_protocol::bittorrent::message::types::BtMessage;
477                let msg = BtMessage::Bitfield { data: bitfield };
478                c.send_message(&serialize(&msg)).await
479            }
480        }
481    }
482
483    pub async fn read_message(
484        &mut self,
485    ) -> Result<Option<aria2_protocol::bittorrent::message::types::BtMessage>> {
486        match &mut self.inner {
487            InnerConnection::Plain(c) => c.read_message().await.map_err(|e| {
488                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
489            }),
490            InnerConnection::Encrypted(c) => c.read_message().await.map_err(|e| {
491                Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { message: e })
492            }),
493            InnerConnection::Utp(c) => {
494                // Receive raw message bytes
495                let msg_bytes = c.recv_message().await?;
496                if let Some(bytes) = msg_bytes {
497                    // Parse into BtMessage
498                    use aria2_protocol::bittorrent::message::factory::parse_message;
499                    parse_message(&bytes).map_err(|e| Aria2Error::Fatal(FatalError::Config(e)))
500                } else {
501                    Ok(None)
502                }
503            }
504        }
505    }
506
507    pub fn is_connected(&self) -> bool {
508        match &self.inner {
509            InnerConnection::Plain(c) => c.is_connected(),
510            InnerConnection::Encrypted(c) => c.is_connected(),
511            InnerConnection::Utp(c) => c.is_connected(),
512        }
513    }
514
515    pub fn is_encrypted(&self) -> bool {
516        matches!(self.inner, InnerConnection::Encrypted(_))
517    }
518}
519
520#[cfg(test)]
521mod tests {
522    use super::*;
523
524    #[test]
525    fn test_allowed_fast_set_operations() {
526        let mut set: HashSet<u32> = HashSet::new();
527        assert!(set.is_empty());
528        assert!(!set.contains(&42));
529        set.insert(42);
530        assert!(set.contains(&42));
531        set.insert(10);
532        set.insert(99);
533        assert_eq!(set.len(), 3);
534        assert!(!set.contains(&999));
535        set.insert(42);
536        assert_eq!(set.len(), 3);
537    }
538
539    #[test]
540    fn test_allowed_fast_multiple_indices() {
541        let mut set: HashSet<u32> = HashSet::new();
542        for i in 0..100u32 {
543            set.insert(i);
544        }
545        assert_eq!(set.len(), 100);
546        for i in 0..100u32 {
547            assert!(set.contains(&i));
548        }
549        assert!(!set.contains(&100));
550    }
551}
aria2_core/engine/bt_peer_connection.rs

aria2_core/engine/
bt_peer_connection.rs
