aria2_core/engine/
bt_mse_handshake.rs

1//! MSE (Message Stream Encryption) encrypted handshake module
2//!
3//! Implements BitTorrent BEP 10 encryption layer protocol, overlaying optional encrypted handshake on BT protocol handshake.
4//! Contains three-phase handshake: Method Selection, PAD/DH key exchange, SKEY/SVC verification.
5
6use rc4::{KeyInit, Rc4 as Rc4Cipher, StreamCipher};
7use ring::agreement::{self, EphemeralPrivateKey, UnparsedPublicKey};
8use ring::rand::SystemRandom;
9use sha1::{Digest, Sha1};
10use std::sync::Mutex;
11
12use crate::error::{Aria2Error, Result};
13
14/// MSE encryption method
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum CryptoMethod {
17    /// Plaintext transmission, no encryption
18    Plain = 0x0001,
19    /// RC4 stream cipher
20    Rc4 = 0x0002,
21    /// AES-128-CBC (optional)
22    Aes128Cbc = 0x0003,
23}
24
25impl CryptoMethod {
26    /// Create CryptoMethod from u16 value
27    pub fn from_u16(value: u16) -> Option<Self> {
28        match value {
29            0x0001 => Some(CryptoMethod::Plain),
30            0x0002 => Some(CryptoMethod::Rc4),
31            0x0003 => Some(CryptoMethod::Aes128Cbc),
32            _ => None,
33        }
34    }
35
36    /// Convert to u16
37    pub fn to_u16(self) -> u16 {
38        self as u16
39    }
40}
41
42/// MSE handshake state machine
43#[derive(Debug, Clone, PartialEq)]
44#[allow(clippy::large_enum_variant)]
45pub enum MseState {
46    /// Idle state
47    Idle,
48    /// Method selection sent, waiting for remote response
49    MethodSelectionSent,
50    /// DH parameter exchange in progress
51    KeyExchangeInProgress,
52    /// Waiting for SKEY verification
53    VerificationPending,
54    /// Handshake complete
55    Established(MseCryptoContext),
56    /// Handshake failed
57    Failed(String),
58}
59
60/// MSE encryption context (used for encrypting/decrypting BT messages after handshake complete)
61pub struct MseCryptoContext {
62    send_key: Vec<u8>,
63    recv_key: Vec<u8>,
64    crypto_method: CryptoMethod,
65    rc4_send: Option<Rc4Cipher>,
66    rc4_recv: Option<Rc4Cipher>,
67}
68
69impl std::fmt::Debug for MseCryptoContext {
70    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
71        f.debug_struct("MseCryptoContext")
72            .field("send_key", &self.send_key)
73            .field("recv_key", &self.recv_key)
74            .field("crypto_method", &self.crypto_method)
75            .field("rc4_send", &self.rc4_send.is_some())
76            .field("rc4_recv", &self.rc4_recv.is_some())
77            .finish()
78    }
79}
80
81impl Clone for MseCryptoContext {
82    fn clone(&self) -> Self {
83        // Note: RC4 state cannot truly be cloned (will lose state), here creating new instance
84        match self.crypto_method {
85            CryptoMethod::Rc4 => {
86                // Attempt to sync state (imperfect but usable)
87                Self::new(&self.send_key, &self.recv_key, self.crypto_method)
88            }
89            _ => Self {
90                send_key: self.send_key.clone(),
91                recv_key: self.recv_key.clone(),
92                crypto_method: self.crypto_method,
93                rc4_send: None,
94                rc4_recv: None,
95            },
96        }
97    }
98}
99
100impl PartialEq for MseCryptoContext {
101    fn eq(&self, other: &Self) -> bool {
102        self.send_key == other.send_key
103            && self.recv_key == other.recv_key
104            && self.crypto_method == other.crypto_method
105    }
106}
107
108impl MseCryptoContext {
109    /// Create new encryption context using derived keys
110    ///
111    /// For RC4 method, initializes two RC4 instances and discards first 1024 bytes keystream
112    /// (MSE spec section 5.2 to prevent keystream attack)
113    pub fn new(send_key: &[u8], recv_key: &[u8], method: CryptoMethod) -> Self {
114        match method {
115            CryptoMethod::Rc4 => {
116                // Initialize send direction RC4 and discard 1024 bytes keystream
117                let mut rc4_send = Rc4Cipher::new_from_slice(send_key).unwrap();
118                let mut discard = vec![0u8; 1024];
119                rc4_send.apply_keystream(&mut discard);
120
121                // Initialize receive direction RC4 and discard 1024 bytes keystream
122                let mut rc4_recv = Rc4Cipher::new_from_slice(recv_key).unwrap();
123                let mut discard = vec![0u8; 1024];
124                rc4_recv.apply_keystream(&mut discard);
125
126                Self {
127                    send_key: send_key.to_vec(),
128                    recv_key: recv_key.to_vec(),
129                    crypto_method: method,
130                    rc4_send: Some(rc4_send),
131                    rc4_recv: Some(rc4_recv),
132                }
133            }
134            _ => Self {
135                send_key: send_key.to_vec(),
136                recv_key: recv_key.to_vec(),
137                crypto_method: method,
138                rc4_send: None,
139                rc4_recv: None,
140            },
141        }
142    }
143
144    /// Encrypt data
145    pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>> {
146        match self.crypto_method {
147            CryptoMethod::Rc4 => {
148                if let Some(ref mut rc4) = self.rc4_send {
149                    let mut data = plaintext.to_vec();
150                    rc4.apply_keystream(&mut data);
151                    Ok(data)
152                } else {
153                    Err(Aria2Error::Fatal(crate::error::FatalError::Config(
154                        "RC4 cipher not initialized".to_string(),
155                    )))
156                }
157            }
158            CryptoMethod::Plain | CryptoMethod::Aes128Cbc => {
159                // Plain and AES modes currently not supported or return plaintext directly
160                Ok(plaintext.to_vec())
161            }
162        }
163    }
164
165    /// Decrypt data
166    pub fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>> {
167        match self.crypto_method {
168            CryptoMethod::Rc4 => {
169                if let Some(ref mut rc4) = self.rc4_recv {
170                    let mut data = ciphertext.to_vec();
171                    rc4.apply_keystream(&mut data);
172                    Ok(data)
173                } else {
174                    Err(Aria2Error::Fatal(crate::error::FatalError::Config(
175                        "RC4 cipher not initialized".to_string(),
176                    )))
177                }
178            }
179            CryptoMethod::Plain | CryptoMethod::Aes128Cbc => {
180                // Plain and AES modes currently not supported or return plaintext directly
181                Ok(ciphertext.to_vec())
182            }
183        }
184    }
185
186    /// Get current encryption method
187    pub fn crypto_method(&self) -> CryptoMethod {
188        self.crypto_method
189    }
190
191    /// Whether using encryption
192    pub fn is_encrypted(&self) -> bool {
193        self.crypto_method != CryptoMethod::Plain
194    }
195}
196
197/// Plaintext fallback: create unencrypted context
198impl Default for MseCryptoContext {
199    fn default() -> Self {
200        Self {
201            send_key: vec![],
202            recv_key: vec![],
203            crypto_method: CryptoMethod::Plain,
204            rc4_send: None,
205            rc4_recv: None,
206        }
207    }
208}
209
210/// MSE handshake manager
211pub struct MseHandshakeManager {
212    /// Handshake state machine.
213    /// Uses std::sync::Mutex because the lock is only held for short synchronous
214    /// reads/writes and never across .await points.
215    state: Mutex<MseState>,
216    local_dh_private_key: Option<EphemeralPrivateKey>,
217    local_dh_pubkey: Vec<u8>,
218    remote_dh_pubkey: Option<Vec<u8>>,
219    pub(crate) shared_secret: Option<Vec<u8>>,
220    info_hash: [u8; 20],
221    pad_length: u16,
222}
223
224impl MseHandshakeManager {
225    /// Create new MSE handshake instance
226    pub fn new(info_hash: [u8; 20]) -> Result<Self> {
227        let rng = SystemRandom::new();
228
229        // Generate X25519 key pair
230        let private_key = EphemeralPrivateKey::generate(&agreement::X25519, &rng).map_err(|e| {
231            Aria2Error::Fatal(crate::error::FatalError::Config(format!(
232                "Failed to generate DH key: {}",
233                e
234            )))
235        })?;
236
237        // Get public key (X25519 public key is fixed at 32 bytes)
238        let public_key = private_key.compute_public_key().map_err(|e| {
239            Aria2Error::Fatal(crate::error::FatalError::Config(format!(
240                "Failed to compute public key: {}",
241                e
242            )))
243        })?;
244
245        let pubkey_slice = public_key.as_ref();
246        let mut local_dh_pubkey = vec![0u8; pubkey_slice.len()];
247        local_dh_pubkey.copy_from_slice(pubkey_slice);
248
249        // Random PAD length (0-512 bytes)
250        use rand::RngCore;
251        let mut rng_core = rand::thread_rng();
252        let pad_length = rng_core.next_u32() as u16 % 513;
253
254        Ok(Self {
255            state: Mutex::new(MseState::Idle),
256            local_dh_private_key: Some(private_key),
257            local_dh_pubkey,
258            remote_dh_pubkey: None,
259            shared_secret: None,
260            info_hash,
261            pad_length,
262        })
263    }
264
265    /// Phase 1: Build Method Selection payload
266    ///
267    /// Returns `\x13MSegadd` (supports encryption) or `\x00` (plaintext only)
268    pub fn build_method_selection(&self) -> Vec<u8> {
269        b"\x13MSegadd".to_vec()
270    }
271
272    /// Parse remote Method Selection
273    ///
274    /// data: Method selection bytes received from remote
275    pub fn parse_remote_method_selection(data: &[u8]) -> Result<CryptoMethod> {
276        if data.is_empty() {
277            return Err(Aria2Error::Parse("Empty method selection".to_string()));
278        }
279
280        // Check if MSegadd (\x13MSegadd)
281        if data == b"\x13MSegadd" {
282            return Ok(CryptoMethod::Rc4); // Supports encryption
283        }
284
285        // Check if plaintext mode (\x00)
286        if data == b"\x00" || data[0] == 0x00 {
287            return Ok(CryptoMethod::Plain);
288        }
289
290        Err(Aria2Error::Parse(format!(
291            "Invalid method selection: {:?}",
292            data
293        )))
294    }
295
296    /// Phase 2: Build PAD + DH public key + CryptoProvisions payload
297    ///
298    /// Format:
299    /// ```text
300    /// ┌────────┬──────────┬─────────────┬────────────┐
301    /// │ PAD_D  │ PAD_LEN  │ Crypto_Pro  │ ICB/IV     │
302    /// │ (2B BE)│ (2B BE)  │ (2B BE)     │ (optional 16B) │
303    /// ├────────┼──────────┼─────────────┼────────────┤
304    /// │          DH Public Key (X25519 = 32 bytes)                │
305    /// └────────┴──────────┴─────────────┴────────────┘
306    /// ```
307    pub fn build_key_exchange_payload(&self, crypto_methods: &[CryptoMethod]) -> Result<Vec<u8>> {
308        let mut payload = Vec::new();
309
310        // PAD_D (2 bytes big-endian): Random padding length
311        payload.extend_from_slice(&self.pad_length.to_be_bytes());
312
313        // PAD_LEN (2 bytes big-endian): Same as PAD_D
314        payload.extend_from_slice(&self.pad_length.to_be_bytes());
315
316        // Crypto_Provisions (2 bytes big-endian): Supported encryption method bitmask
317        let mut crypto_provisions: u16 = 0;
318        for method in crypto_methods {
319            crypto_provisions |= method.to_u16();
320        }
321        payload.extend_from_slice(&crypto_provisions.to_be_bytes());
322
323        // Random PAD data
324        use rand::RngCore;
325        let mut rng_core = rand::thread_rng();
326        let mut pad_data = vec![0u8; self.pad_length as usize];
327        rng_core.fill_bytes(&mut pad_data);
328        payload.extend_from_slice(&pad_data);
329
330        // DH Public Key (32 bytes for X25519)
331        payload.extend_from_slice(&self.local_dh_pubkey);
332
333        Ok(payload)
334    }
335
336    /// Parse remote key exchange payload, compute shared secret
337    pub fn process_remote_key_exchange(&mut self, data: &[u8]) -> Result<()> {
338        // Minimum payload = 6-byte header (PAD_D + PAD_LEN + Crypto_Pro) + 32-byte
339        // DH public key = 38 bytes. PAD data length is variable (0..=512) and is
340        // validated separately below against the embedded PAD_LEN field.
341        if data.len() < 6 + 32 {
342            return Err(Aria2Error::Parse(
343                "Key exchange payload too short".to_string(),
344            ));
345        }
346
347        // Parse fields
348        let _pad_d = u16::from_be_bytes([data[0], data[1]]);
349        let _pad_len = u16::from_be_bytes([data[2], data[3]]);
350        let _crypto_pro = u16::from_be_bytes([data[4], data[5]]);
351        // Skip possible ICB/IV (6-22)
352
353        // Calculate PAD end position
354        let pad_end = 6 + (_pad_len as usize);
355        if data.len() < pad_end + 32 {
356            return Err(Aria2Error::Parse(
357                "Key exchange payload truncated".to_string(),
358            ));
359        }
360
361        // Extract remote DH public key (last 32 bytes)
362        let remote_pubkey_start = data.len() - 32;
363        let remote_pubkey = &data[remote_pubkey_start..];
364
365        // Store remote public key
366        self.remote_dh_pubkey = Some(remote_pubkey.to_vec());
367
368        // Compute shared secret
369        let peer_public_key = UnparsedPublicKey::new(&agreement::X25519, remote_pubkey);
370
371        if let Some(private_key) = self.local_dh_private_key.take() {
372            let shared_secret: Vec<u8> = agreement::agree_ephemeral(
373                private_key,
374                &peer_public_key,
375                |shared_secret: &[u8]| shared_secret.to_vec(),
376            )
377            .map_err(|e| {
378                Aria2Error::Fatal(crate::error::FatalError::Config(format!(
379                    "DH key agreement failed: {}",
380                    e
381                )))
382            })?;
383
384            self.shared_secret = Some(shared_secret);
385        } else {
386            return Err(Aria2Error::Fatal(crate::error::FatalError::Config(
387                "Local DH private key not set".to_string(),
388            )));
389        }
390
391        Ok(())
392    }
393
394    /// Phase 3: Build SKEY + VC + CryptoSelect verification payload
395    ///
396    /// Format:
397    /// - SKEY hash (20 bytes): SHA-1(info_hash + shared_secret)
398    /// - VC (2 bytes): version cipher (0x0001 for RC4)
399    /// - CryptoSelect (2 bytes): Selected encryption method
400    /// - len(I) (2 bytes): Optional initial data length (usually 0)
401    pub fn build_verification_payload(&self, selected_method: CryptoMethod) -> Result<Vec<u8>> {
402        let skey = self.compute_skey()?;
403        let mut payload = Vec::new();
404
405        // SKEY hash (20 bytes)
406        payload.extend_from_slice(&skey);
407
408        // VC (2 bytes): Version cipher
409        payload.extend_from_slice(&0x0001u16.to_be_bytes());
410
411        // CryptoSelect (2 bytes): Selected encryption method
412        payload.extend_from_slice(&selected_method.to_u16().to_be_bytes());
413
414        // len(I) (2 bytes): Initial data length
415        payload.extend_from_slice(&0x0000u16.to_be_bytes());
416
417        Ok(payload)
418    }
419
420    /// Parse remote verification payload, verify SKEY, complete handshake
421    pub fn process_remote_verification(&mut self, data: &[u8]) -> Result<MseCryptoContext> {
422        if data.len() < 26 {
423            return Err(Aria2Error::Parse(
424                "Verification payload too short".to_string(),
425            ));
426        }
427
428        // Parse fields
429        let remote_skey = &data[..20];
430        let vc = u16::from_be_bytes([data[20], data[21]]);
431        let crypto_select = u16::from_be_bytes([data[22], data[23]]);
432        // len(I) at [24,25]
433
434        // Validate VC
435        if vc != 0x0001 {
436            return Err(Aria2Error::Parse(format!("Invalid VC value: {:#06x}", vc)));
437        }
438
439        // Validate SKEY
440        let expected_skey = self.compute_skey()?;
441        if remote_skey != expected_skey.as_slice() {
442            return Err(Aria2Error::Checksum("SKEY verification failed".to_string()));
443        }
444
445        // Parse selected encryption method
446        let selected_method = CryptoMethod::from_u16(crypto_select).ok_or_else(|| {
447            Aria2Error::Parse(format!(
448                "Unknown crypto method selected: {:#06x}",
449                crypto_select
450            ))
451        })?;
452
453        // Derive encryption keys
454        let shared_secret = self.shared_secret.as_ref().ok_or_else(|| {
455            Aria2Error::Fatal(crate::error::FatalError::Config(
456                "Shared secret not computed".to_string(),
457            ))
458        })?;
459
460        let (send_key, recv_key) = Self::derive_keys(&expected_skey, shared_secret);
461
462        // Create encryption context
463        let ctx = MseCryptoContext::new(&send_key, &recv_key, selected_method);
464
465        Ok(ctx)
466    }
467
468    /// Internal: Compute SKEY = SHA-1(info_hash || shared_secret)
469    pub(crate) fn compute_skey(&self) -> Result<[u8; 20]> {
470        let shared_secret = self.shared_secret.as_ref().ok_or_else(|| {
471            Aria2Error::Fatal(crate::error::FatalError::Config(
472                "Shared secret not computed yet".to_string(),
473            ))
474        })?;
475
476        let mut hasher = Sha1::new();
477        hasher.update(self.info_hash);
478        hasher.update(shared_secret);
479        let result = hasher.finalize();
480
481        let mut skey = [0u8; 20];
482        skey.copy_from_slice(&result);
483        Ok(skey)
484    }
485
486    /// Internal: Derive encryption keys
487    ///
488    /// ```text
489    /// send_key = SHA-1(SKEY + "keyA" + shared_secret)[:16]
490    /// recv_key = SHA-1(SKEY + "keyB" + shared_secret)[:16]
491    /// ```
492    pub(crate) fn derive_keys(skey: &[u8], shared_secret: &[u8]) -> (Vec<u8>, Vec<u8>) {
493        // Send direction key
494        let mut hasher_a = Sha1::new();
495        hasher_a.update(skey);
496        hasher_a.update(b"keyA");
497        hasher_a.update(shared_secret);
498        let result_a = hasher_a.finalize();
499
500        // Receive direction key
501        let mut hasher_b = Sha1::new();
502        hasher_b.update(skey);
503        hasher_b.update(b"keyB");
504        hasher_b.update(shared_secret);
505        let result_b = hasher_b.finalize();
506
507        // Take first 16 bytes
508        (result_a[..16].to_vec(), result_b[..16].to_vec())
509    }
510
511    /// Plaintext fallback: create unencrypted context
512    pub fn plaintext_fallback() -> MseCryptoContext {
513        MseCryptoContext::default()
514    }
515
516    /// Get current state
517    pub fn state(&self) -> MseState {
518        self.state.lock().unwrap_or_else(|e| e.into_inner()).clone()
519    }
520
521    /// Update state
522    pub fn set_state(&self, state: MseState) {
523        *self.state.lock().unwrap_or_else(|e| e.into_inner()) = state;
524    }
525
526    /// Get local DH public key (for debugging)
527    pub fn local_dh_pubkey(&self) -> &[u8] {
528        &self.local_dh_pubkey
529    }
530
531    /// Get shared secret (for debugging)
532    pub fn shared_secret(&self) -> Option<&[u8]> {
533        self.shared_secret.as_deref()
534    }
535}
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540
541    #[test]
542    fn test_crypto_method_conversions() {
543        assert_eq!(CryptoMethod::from_u16(0x0001), Some(CryptoMethod::Plain));
544        assert_eq!(CryptoMethod::from_u16(0x0002), Some(CryptoMethod::Rc4));
545        assert_eq!(
546            CryptoMethod::from_u16(0x0003),
547            Some(CryptoMethod::Aes128Cbc)
548        );
549        assert_eq!(CryptoMethod::from_u16(0x9999), None);
550
551        assert_eq!(CryptoMethod::Plain.to_u16(), 0x0001);
552        assert_eq!(CryptoMethod::Rc4.to_u16(), 0x0002);
553        assert_eq!(CryptoMethod::Aes128Cbc.to_u16(), 0x0003);
554    }
555}
aria2_core/engine/bt_mse_handshake.rs

aria2_core/engine/
bt_mse_handshake.rs
