diff --git a/.idea/.gitignore b/.idea/.gitignore new file mode 100644 index 000000000..f6906f2ee --- /dev/null +++ b/.idea/.gitignore @@ -0,0 +1,10 @@ +# 默认忽略的文件 +/shelf/ +/workspace.xml +# 基于编辑器的 HTTP 客户端请求 +/httpRequests/ +# 已忽略包含查询文件的默认文件夹 +/queries/ +# Datasource local storage ignored files +/dataSources/ +/dataSources.local.xml diff --git a/.idea/modules.xml b/.idea/modules.xml new file mode 100644 index 000000000..ca1ed76ec --- /dev/null +++ b/.idea/modules.xml @@ -0,0 +1,8 @@ + + + + + + + + \ No newline at end of file diff --git a/.idea/rqbit.iml b/.idea/rqbit.iml new file mode 100644 index 000000000..97118da4f --- /dev/null +++ b/.idea/rqbit.iml @@ -0,0 +1,30 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml new file mode 100644 index 000000000..35eb1ddfb --- /dev/null +++ b/.idea/vcs.xml @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 7b4f8f903..c41af1af8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -721,6 +721,17 @@ dependencies = [ "half", ] +[[package]] +name = "cipher" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8cf2a2c93cd704877c0858356ed03480ff301ee950b43f1cbe4573b088bfa6c" +dependencies = [ + "block-buffer 0.12.1", + "crypto-common 0.2.2", + "inout", +] + [[package]] name = "clap" version = "4.6.6" @@ -935,6 +946,12 @@ dependencies = [ "libc", ] +[[package]] +name = "cpubits" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15b85f9c39137c3a891689859392b1bd49812121d0d61c9caf00d46ed5ce06ae" + [[package]] name = "cpufeatures" version = "0.2.17" @@ -1061,6 +1078,18 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "crypto-bigint" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a52aa3fcda4e6302a9f48734f234d35d4721b96f8fe07d073f07ce9df4f0271" +dependencies = [ + "cpubits", + "ctutils", + "num-traits", + "rand_core 0.10.1", +] + [[package]] name = "crypto-common" version = "0.1.7" @@ -2670,6 +2699,15 @@ dependencies = [ "libc", ] +[[package]] +name = "inout" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4250ce6452e92010fdf7268ccc5d14faa80bb12fc741938534c58f16804e03c7" +dependencies = [ + "hybrid-array", +] + [[package]] name = "intervaltree" version = "0.2.7" @@ -3071,6 +3109,7 @@ dependencies = [ "byteorder", "bytes", "console-subscriber", + "crypto-bigint", "dashmap 6.2.1", "futures", "governor", @@ -3101,6 +3140,7 @@ dependencies = [ "parking_lot", "pollster", "rand 0.10.2", + "rc4", "regex", "reqwest", "rlimit", @@ -4737,6 +4777,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "rc4" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "840038b674daa9f7a7957440d937951d15c0143c056e631e529141fd780e0c92" +dependencies = [ + "cipher", +] + [[package]] name = "redox_syscall" version = "0.5.18" diff --git a/Cargo.toml b/Cargo.toml index a5bee8318..d082893ff 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,6 +83,7 @@ axum-extra = "0.12" backon = "1.5" base64 = "0.23" bitvec = "1" +crypto-bigint = "0.7" bstr = "1" byteorder = "1" bytes = "1" diff --git a/MSE-PR633-REFACTOR-PLAN.md b/MSE-PR633-REFACTOR-PLAN.md new file mode 100644 index 000000000..f34096c4b --- /dev/null +++ b/MSE-PR633-REFACTOR-PLAN.md @@ -0,0 +1,58 @@ +# MSE PR #633 重构（已实施）：概要与决策记录 + +## 背景 + +rqbit PR #633 实现 MSE（消息流加密）。上游维护者 ikatson 提交 CHANGES_REQUESTED，本文件记录重构的完整梗概：维护者意见、已实施的改动、关键设计决策与验证结果。分支 `feat/mse-crypto-primitives`，单 commit `0402294c`（基于 `de2b107e`）。 + +## 一、维护者意见（CHANGES_REQUESTED） + +主 review：**合并为单 commit**（"It's not too big, so let's just shrink into one commit"）+ 两大担忧： +1. **自写 crypto 太多**——应尽量用 crate（openssl 等，含硬件加速） +2. **spaghetti**——MSE 侵入 session/peer_connection 核心逻辑，建议抽象为 `connect_with_handshake` / `accept` + +6 条 review comments：error.rs `MseForced` 去 SocketAddr、移除 Tcp-only 限制、去冗余行、RC4/DH 用 crate、session 抽象、uTP/socks 透明支持。 + +## 二、已实施的改动 + +### Phase 1 — 抽象层（解决 spaghetti） +- **出站**：`StreamConnector::connect_with_handshake()`（stream_connect.rs）——封装连接 + MSE 决策 + 握手发送，返回 `OutgoingHandshake { kind, read, write, mse_applied }` +- **入站**：`accept_with_handshake()` 自由函数 + `IncomingHandshake` 结构 +- `session.rs`/`peer_connection.rs` 不再携带任何 MSE 细节（`IncomingOutcome`/`OutgoingOutcome`/`Sha1`/`AsyncReadExt` 全部收进 stream_connect） +- 移除 Tcp-only 限制：uTP/socks 透明支持 MSE +- `MseForced` 去 SocketAddr + +### Phase 2 — 加密原语（解决自写 crypto） +- **DH-768 → `crypto-bigint`**：`Uint<12>` + `FixedMontyForm::pow`，删除手写 mod_reduce hack。MSE 固定 prime 硬编码（与 libtorrent/transmission/aria2 同值）。已验证外部 bigint 向量一致 +- **RC4 → `rc4` crate**： + - 删除自实现 `mse/rc4.rs`（140 行） + - `Rc4Writer` 重写为**内循环方案**（参考 libtorrent `rc4_handler`）：`poll_write` 加密整 buffer → 存 pending → 循环 flush，返回 `Ok(full_len)` 或 `Pending`，**绝不返回短写**（消除双加密 bug）；Pending/短写密文缓存，重试不重复加密 + - `Rc4Reader` 用 rc4 crate + - **PadB 扫描改 pattern-search**（libtorrent `read_pe_syncvc` / transmission `read_vc` 同款）：独立实例算 VC 密文模式 → 字节搜索 → 解密实际 VC 验证，不再需要 clone +- SHA-1 用现有 `sha1w`（crate + 可选硬件加速后端） + +### Phase 4 — squash +- `git reset --soft de2b107e` → 单 commit `0402294c`，21 文件 +1843/−32 + +## 三、关键设计决策 + +### RC4 保留自实现 vs 用 crate 的决策过程 +1. 最初倾向保留自实现：RustCrypto `rc4` crate **无 `Clone`**，无法实现 `Rc4Writer` 短写保护（Pending 重试需试探状态） +2. 调研发现：libtorrent/transmission **都不用 clone**——用 **pattern-search**（搜索加密后的 VC 密文模式）而非"试探解密"；`Rc4Writer` 短写用"加密与发送分离 + pending 缓冲" +3. 据此改用 rc4 crate：PadB 用 pattern-search，Rc4Writer 用内循环方案——**彻底绕开 clone 需求** + +### 写错误处理（与三引擎一致） +libtorrent `disconnect(error, sock_write)`、transmission `call_error_callback`（非可重试）、aria2 `throw DL_RETRY_EX`（非 WOULDBLOCK）——**致命写错误一律断连**。rqbit 的 `poll_write` 对底层 `Poll::Pending` 返回 Pending（可重试，对应 EAGAIN），对 `Err` 传播（断连）。原测试 `write_error_does_not_advance_state`（错误后状态不推进）脱离真实场景，改为 `write_error_propagates`（错误传播 + sink 空）。 + +## 四、验证结果 + +- `cargo test -p librqbit --lib`：**45 passed / 0 failed / 5 ignored** + - stream：pending/short-write 状态保持、写错误传播 + - mse：rc4 向量、dh768 外部向量、stream 包装、duplex 握手、明文嗅探回退、零长度 IA、分片前缀重放、fresh-redial、Disabled 单连接、Forced 失败、默认 Disabled 锁定 +- `cargo check --all-targets`：0 error +- `cargo clippy`：新增改动 0 warning（剩余 5 个为基线既有 mod.rs cast/large-size 警告） + +## 五、状态 + +- 单 commit `0402294c`（基于 `de2b107e`），工作区干净 +- **未 push、未回复 PR**（按用户指示） +- 待办：回复维护者（含 pattern-search / 内循环方案说明）、force push diff --git a/crates/librqbit/Cargo.toml b/crates/librqbit/Cargo.toml index 1623e2254..e1d3cd588 100644 --- a/crates/librqbit/Cargo.toml +++ b/crates/librqbit/Cargo.toml @@ -54,6 +54,8 @@ clone_to_owned.workspace = true peer_binary_protocol.workspace = true sha1w.workspace = true dht.workspace = true +crypto-bigint.workspace = true +rc4 = "0.2" librqbit-upnp.workspace = true upnp-serve = { workspace = true, optional = true } diff --git a/crates/librqbit/examples/simulate_traffic.rs b/crates/librqbit/examples/simulate_traffic.rs index 9e4286e48..5d91afa7a 100644 --- a/crates/librqbit/examples/simulate_traffic.rs +++ b/crates/librqbit/examples/simulate_traffic.rs @@ -186,6 +186,7 @@ impl TestHarness { connect_timeout: Some(Duration::from_secs(1)), read_write_timeout: Some(Duration::from_secs(32)), keep_alive_interval: None, + ..Default::default() }), }), ..Default::default() @@ -274,6 +275,7 @@ impl TestHarness { connect_timeout: Some(Duration::from_secs(1)), read_write_timeout: Some(Duration::from_secs(32)), keep_alive_interval: None, + ..Default::default() }), }), ..Default::default() diff --git a/crates/librqbit/src/error.rs b/crates/librqbit/src/error.rs index c95250381..530e8665f 100644 --- a/crates/librqbit/src/error.rs +++ b/crates/librqbit/src/error.rs @@ -30,6 +30,8 @@ pub enum Error { WrongInfoHash, #[error("connecting to ourselves")] ConnectingToOurselves, + #[error("MSE is forced, but the peer did not complete the MSE handshake")] + MseForced, #[error("error writing handshake: {0:#}")] WriteHandshake(#[source] std::io::Error), diff --git a/crates/librqbit/src/lib.rs b/crates/librqbit/src/lib.rs index 29843e008..5447971ef 100644 --- a/crates/librqbit/src/lib.rs +++ b/crates/librqbit/src/lib.rs @@ -61,6 +61,7 @@ mod ip_ranges; pub mod limits; mod listen; mod merge_streams; +mod mse; mod peer_connection; mod peer_info_reader; mod piece_tracker; @@ -90,6 +91,7 @@ pub use create_torrent_file::{CreateTorrentOptions, CreateTorrentResult, create_ pub use dht; pub use librqbit_core::spawn_utils::spawn as librqbit_spawn; pub use listen::{ListenerMode, ListenerOptions}; +pub use mse::MseMode; pub use peer_connection::PeerConnectionOptions; pub use session::{ AddTorrent, AddTorrentOptions, AddTorrentResponse, DhtSessionConfig, ListOnlyResponse, diff --git a/crates/librqbit/src/mse/dh768.rs b/crates/librqbit/src/mse/dh768.rs new file mode 100644 index 000000000..477fb5f97 --- /dev/null +++ b/crates/librqbit/src/mse/dh768.rs @@ -0,0 +1,166 @@ +//! 768-bit Diffie-Hellman key exchange for BitTorrent MSE. +//! +//! The MSE spec fixes a 768-bit MODP group (prime + generator 2) that is not +//! one of the standard groups shipped by crypto libraries (OpenSSL etc.), so +//! the group parameters are hardcoded here (matching libtorrent/transmission/ +//! aria2, which do the same) and the modular exponentiation is delegated to +//! `crypto-bigint` (pure Rust, no C dependency, works with rqbit's `rust-tls` +//! build). The private exponent is 160 bits, matching common MSE peers. +//! +//! Note: `FixedMontyParams::new_vartime` is not constant-time. This is +//! acceptable here because the private exponent is a fresh ephemeral session +//! key generated per connection (not a long-lived secret); libtorrent's +//! non-openssl path uses boost `cpp_int::powm` (also variable-time) for the +//! same reason. + +use crypto_bigint::{ + Odd, Uint, + modular::{FixedMontyForm, FixedMontyParams}, +}; +use rand::Rng; + +/// 768-bit modulus: 12 limbs of 64 bits. +type U768 = Uint<12>; + +/// MSE-specified prime (RFC 2409-style 768-bit group, same hex as libtorrent +/// `pe_crypto.cpp` / transmission `peer-mse.cc` / aria2 `MSEHandshake.cc`). +const DH_PRIME_HEX: &str = concat!( + "FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1", + "29024E088A67CC74020BBEA63B139B22514A08798E3404DD", + "EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245", + "E485B576625E7EC6F44C42E9A63A36210000000000090563" +); + +const TWO: U768 = { + let mut value = [0u64; 12]; + value[0] = 2; + U768::from_words(value) +}; + +fn prime() -> U768 { + U768::from_be_hex(DH_PRIME_HEX) +} + +fn bytes_to_u768(bytes: &[u8; 96]) -> U768 { + U768::from_be_slice(bytes) +} + +fn u768_to_bytes(value: &U768) -> [u8; 96] { + let encoded: [u8; 96] = value.to_be_bytes().into(); + encoded +} + +fn powm(base: &U768, exponent: &[u8; 20]) -> U768 { + let exp = secret_to_u768(exponent); + powm_u(base, &exp) +} + +fn secret_to_u768(secret: &[u8; 20]) -> U768 { + U768::from_be_slice_truncated(secret, 160) +} + +fn powm_u(base: &U768, exponent: &U768) -> U768 { + let p_odd: Odd = Odd::new(prime()).expect("MSE prime is odd"); + // Variable-time Montgomery params: the private exponent is a fresh + // per-connection ephemeral key (not a long-lived secret), so constant-time + // is not required here; see the module note. + let params = FixedMontyParams::new_vartime(p_odd); + let monty_base = FixedMontyForm::new(base, ¶ms); + monty_base.pow(exponent).retrieve() +} + +pub struct Dh768 { + secret: [u8; 20], + public: [u8; 96], +} + +impl Dh768 { + pub fn generate(rng: &mut impl Rng) -> Self { + let mut secret = [0u8; 20]; + while secret.iter().all(|byte| *byte == 0) { + rng.fill_bytes(&mut secret); + } + Self::from_secret(secret) + } + + pub(super) fn from_secret(secret: [u8; 20]) -> Self { + let public = powm(&TWO, &secret); + Self { + secret, + public: u768_to_bytes(&public), + } + } + + pub fn public_key_bytes(&self) -> [u8; 96] { + self.public + } + + pub fn shared_secret(&self, remote: &[u8; 96]) -> Option<[u8; 96]> { + let remote = bytes_to_u768(remote); + // Reject degenerate keys (0, 1, or >= p-1), matching the old hand-rolled + // bounds check. + let two = U768::from(2u64); + let p = prime(); + let p_minus_one = p.wrapping_sub(&U768::ONE); + if remote < two || remote >= p_minus_one { + return None; + } + let secret_u = secret_to_u768(&self.secret); + Some(u768_to_bytes(&powm_u(&remote, &secret_u))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use anyhow::{Context, Result}; + use rand::SeedableRng; + use rand::rngs::SmallRng; + + fn decode(text: &str) -> Result<[u8; N]> { + let bytes = hex::decode(text).context("invalid test vector hex")?; + bytes + .try_into() + .map_err(|_| anyhow::anyhow!("test vector has the wrong length")) + } + + #[test] + fn matches_external_bigint_vectors() -> Result<()> { + let secret_a = decode::<20>("000102030405060708090a0b0c0d0e0f10111213")?; + let public_a = decode::<96>( + "7fba71c678158bd55ef1cc04a919d1b05f79f9da403c67e82bb1a99a7b4bc4ec221cca6c3a78171a40f2cc12e3d9d4454338f7e4b9b33de5e82ab04e86f5cd43aaf9dad923988501c371d3159935de5499e5d726e740b1eabbf4a3dd03c68071", + )?; + let secret_b = decode::<20>("f0e0d0c0b0a09080706050403020100011223344")?; + let public_b = decode::<96>( + "f9fe7e1c27aee331ab8ff8a6183cfcc7bd08dc593fc4d52bc9a2694b7b787daa12e3b2695e3e9febf994447cefa427f9f5da34a4d3cd6c231a8d6517e7130de00a8a09e753ca12648ec18da389e68eeb66f8308b19cc60dfeaadb2540a821f53", + )?; + let shared = decode::<96>( + "909ea4557d5b9f43dafdc5b598850045b8689e4d652af58a63730b00c574bbe4962ab9c78b2f295e3ddb3b456f20a4c65761751bf5d79ec4dba8470fe66ed22b4a25f13528a9575607c77586785a36d560f8556b66e9c16deb87fed185ee07a7", + )?; + + let a = Dh768::from_secret(secret_a); + let b = Dh768::from_secret(secret_b); + assert_eq!(a.public_key_bytes(), public_a); + assert_eq!(b.public_key_bytes(), public_b); + assert_eq!(a.shared_secret(&public_b), Some(shared)); + assert_eq!(b.shared_secret(&public_a), Some(shared)); + Ok(()) + } + + #[test] + fn generated_secret_is_nonzero() { + let mut rng = SmallRng::seed_from_u64(0x5eed); + let dh = Dh768::generate(&mut rng); + assert!(dh.secret.iter().any(|byte| *byte != 0)); + } + + #[test] + fn rejects_degenerate_remote_keys() { + let dh = Dh768::from_secret([1u8; 20]); + assert!(dh.shared_secret(&[0u8; 96]).is_none()); + assert!(dh.shared_secret(&[0xffu8; 96]).is_none()); + let mut one = [0u8; 96]; + one[95] = 1; + assert!(dh.shared_secret(&one).is_none()); + } +} diff --git a/crates/librqbit/src/mse/mod.rs b/crates/librqbit/src/mse/mod.rs new file mode 100644 index 000000000..70c4252bb --- /dev/null +++ b/crates/librqbit/src/mse/mod.rs @@ -0,0 +1,453 @@ +//! Message Stream Encryption (MSE) support. +//! +//! MSE obfuscates the BitTorrent handshake and stream with RC4 and a DH-768 +//! key exchange (Azureus protocol encryption). Some peers (Xunlei, BitComet, +//! etc.) refuse plaintext handshakes, so MSE is required to talk to them. +//! +//! This module is built top-down: the [`MseMode`] config and the peer +//! connection wiring landed first, then the incoming (responder) handshake, +//! and the outgoing (initiator) handshake in a follow-up commit. + +pub mod dh768; +pub mod stream; + +use std::time::Duration; + +use anyhow::{Context, Result, bail}; +use rand::{Rng, RngExt}; +use sha1w::{ISha1, Sha1}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tracing::{debug, trace, warn}; + +use ::rc4::{KeyInit, Rc4, StreamCipher}; +use dh768::Dh768; +use stream::{Rc4Reader, Rc4Writer}; + +const BT_PROTOCOL_PREFIX: &[u8; 20] = b"\x13BitTorrent protocol"; +const BT_HANDSHAKE_LEN: usize = 68; +const MAX_PAD: usize = 512; +const VC_LEN: usize = 8; +const CRYPTO_RC4: u32 = 2; + +/// How long to wait (after sending Ya + PadA) for a peer to send its first 20 +/// bytes before treating it as a silent MSE responder. A plaintext peer sends +/// its 68-byte BT handshake immediately, so we detect it here instead of +/// waiting for the full read/write timeout. +const PLAINTEXT_SNIFF_TIMEOUT: Duration = Duration::from_secs(2); + +/// How MSE is applied to peer connections. +/// +/// `Disabled` (the default) skips MSE entirely (plaintext only). `Enabled` +/// prefers MSE and falls back to a plaintext redial on failure. `Forced` +/// requires MSE: any peer that does not complete the MSE handshake is dropped. +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub enum MseMode { + Disabled, + Enabled, + Forced, +} + +// Manual impl rather than `#[derive(Default)]` + `#[default]`: keeping the +// default explicit here avoids silently carrying a `#[default]` annotation +// over if the default is ever flipped to another variant. +#[allow(clippy::derivable_impls)] +impl Default for MseMode { + fn default() -> Self { + Self::Disabled + } +} + +pub struct PrefixReader { + prefix: Vec, + position: usize, + inner: R, +} + +impl PrefixReader { + fn new(prefix: Vec, inner: R) -> Self { + Self { + prefix, + position: 0, + inner, + } + } +} + +impl AsyncRead for PrefixReader { + fn poll_read( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + buf: &mut tokio::io::ReadBuf<'_>, + ) -> std::task::Poll> { + let this = self.get_mut(); + if this.position < this.prefix.len() && buf.remaining() != 0 { + let count = (this.prefix.len() - this.position).min(buf.remaining()); + buf.put_slice(&this.prefix[this.position..this.position + count]); + this.position += count; + return std::task::Poll::Ready(Ok(())); + } + std::pin::Pin::new(&mut this.inner).poll_read(cx, buf) + } +} + +/// Encrypted carries the RC4 state and the full handshake; Plaintext is the +/// fast path. The size gap is inherent to holding cipher state, not a defect. +#[allow(clippy::large_enum_variant)] +pub enum IncomingOutcome { + Encrypted { + read: Rc4Reader, + write: Rc4Writer, + handshake_bytes: Vec, + info_hash: [u8; 20], + }, + Plaintext { + read: PrefixReader, + write: W, + }, +} + +/// Encrypted carries the RC4-wrapped streams; PlaintextPeer carries none. +#[allow(clippy::large_enum_variant)] +pub enum OutgoingOutcome { + /// MSE handshake completed; use the RC4-wrapped streams. + Encrypted(Rc4Reader, Rc4Writer), + /// The peer answered with a plaintext BitTorrent handshake within the + /// sniff window. Abort MSE and redial plaintext on a fresh connection. + PlaintextPeer, +} + +fn sha1(parts: &[&[u8]]) -> [u8; 20] { + let mut hash = Sha1::new(); + for part in parts { + hash.update(part); + } + hash.finish() +} + +fn xor20(a: &[u8; 20], b: &[u8; 20]) -> [u8; 20] { + let mut result = [0u8; 20]; + for i in 0..20 { + result[i] = a[i] ^ b[i]; + } + result +} + +fn derive_keys(secret: &[u8], skey: &[u8], outgoing: bool) -> (Rc4, Rc4, [u8; 20]) { + let (encrypt_key, decrypt_key) = if outgoing { + ( + sha1(&[b"keyA", secret, skey]), + sha1(&[b"keyB", secret, skey]), + ) + } else { + ( + sha1(&[b"keyB", secret, skey]), + sha1(&[b"keyA", secret, skey]), + ) + }; + let mut encrypt = Rc4::new_from_slice(&encrypt_key).expect("20-byte RC4 key"); + let mut decrypt = Rc4::new_from_slice(&decrypt_key).expect("20-byte RC4 key"); + // MSE drop-1024: discard the first 1024 keystream bytes. + let mut drop = [0u8; 1024]; + encrypt.apply_keystream(&mut drop); + decrypt.apply_keystream(&mut drop); + (encrypt, decrypt, decrypt_key) +} + +async fn read_scan_for_needle( + read: &mut R, + needle: &[u8], + max_pad: usize, +) -> Result { + let mut window = Vec::with_capacity(max_pad + needle.len()); + let mut byte = [0u8; 1]; + loop { + read.read_exact(&mut byte) + .await + .context("disconnected while scanning MSE handshake")?; + window.push(byte[0]); + if window.ends_with(needle) { + return Ok(window.len() - needle.len()); + } + if window.len() >= max_pad + needle.len() { + bail!("MSE pattern not found within {max_pad} pad bytes"); + } + } +} + +async fn read_encrypted( + read: &mut R, + decrypt: &mut Rc4, + bytes: &mut [u8], +) -> Result<()> { + read.read_exact(bytes).await?; + decrypt.apply_keystream(bytes); + Ok(()) +} + +fn random_pad(max: usize) -> Vec { + let length = rand::rng().random_range(0..=max); + let mut pad = vec![0u8; length]; + rand::rng().fill_bytes(&mut pad); + pad +} + +/// Probe the responder's first 20 bytes within a short window to detect a +/// plaintext peer (which sends its BT handshake immediately on connect). +/// +/// Returns `Ok(Some(prefix))` with the 20 bytes read when they start with the +/// BT protocol prefix (peer is plaintext). Returns `Ok(None)` with the bytes +/// read so far when they do not (they are the leading bytes of the MSE +/// responder's DH public key), or when the window elapsed without a full 20 +/// bytes (treat as a slow MSE responder). The partial `sniffed` bytes must be +/// preserved as the prefix of the DH public key. +async fn sniff_plaintext(read: &mut R) -> Result<(Option>, Vec)> { + let mut sniffed = Vec::with_capacity(BT_PROTOCOL_PREFIX.len()); + let probe = async { + let mut byte = [0u8; 1]; + while sniffed.len() < BT_PROTOCOL_PREFIX.len() { + read.read_exact(&mut byte).await?; + sniffed.push(byte[0]); + } + Ok::<_, std::io::Error>(()) + }; + match tokio::time::timeout(PLAINTEXT_SNIFF_TIMEOUT, probe).await { + Ok(Ok(())) => { + let is_plaintext = sniffed == BT_PROTOCOL_PREFIX; + Ok((is_plaintext.then(|| sniffed.clone()), sniffed)) + } + Ok(Err(e)) => Err(e.into()), + Err(_elapsed) => { + // Window elapsed mid-probe: whatever we got is the leading bytes of + // the responder's public key (or a slow plaintext peer). Continue MSE. + Ok((None, sniffed)) + } + } +} + +/// Initiate MSE on a connected stream. IA must be the complete 68-byte +/// BitTorrent handshake and is consumed as part of the MSE exchange. +/// +/// Before committing to MSE, briefly probes the responder for a plaintext BT +/// handshake (see [`sniff_plaintext`]); if detected, returns +/// [`OutgoingOutcome::PlaintextPeer`] so the caller can redial plaintext +/// instead of waiting out the full MSE read timeout. +pub async fn outgoing( + mut read: R, + mut write: W, + info_hash: &[u8; 20], + initial_payload: &[u8; BT_HANDSHAKE_LEN], +) -> Result> +where + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, +{ + let dh = Dh768::generate(&mut rand::rng()); + write.write_all(&dh.public_key_bytes()).await?; + write.write_all(&random_pad(MAX_PAD)).await?; + trace!("sent MSE Ya + PadA, probing responder"); + + // Sniff for a plaintext responder before committing to MSE. + let (plaintext, sniffed) = sniff_plaintext(&mut read).await?; + if plaintext.is_some() { + debug!( + "peer answered with plaintext BT handshake within sniff window, falling back to plaintext redial" + ); + return Ok(OutgoingOutcome::PlaintextPeer); + } + + let mut server_public = [0u8; 96]; + server_public[..sniffed.len()].copy_from_slice(&sniffed); + read.read_exact(&mut server_public[sniffed.len()..]) + .await + .context("disconnected waiting for MSE responder public key")?; + let secret = dh + .shared_secret(&server_public) + .ok_or_else(|| anyhow::anyhow!("MSE degenerate remote DH key"))?; + let (mut encrypt, decrypt_base, decrypt_key) = derive_keys(&secret, info_hash, true); + + write.write_all(&sha1(&[b"req1", &secret])).await?; + let skey_hash = sha1(&[b"req2", info_hash]); + let req3 = sha1(&[b"req3", &secret]); + write.write_all(&xor20(&skey_hash, &req3)).await?; + + let pad_c = random_pad(MAX_PAD); + let mut encrypted = + Vec::with_capacity(VC_LEN + 4 + 2 + pad_c.len() + 2 + initial_payload.len()); + encrypted.extend_from_slice(&[0u8; VC_LEN]); + encrypted.extend_from_slice(&CRYPTO_RC4.to_be_bytes()); + encrypted.extend_from_slice(&u16::try_from(pad_c.len()).expect("pad_c <= MAX_PAD").to_be_bytes()); + encrypted.extend_from_slice(&pad_c); + encrypted.extend_from_slice(&u16::try_from(BT_HANDSHAKE_LEN).expect("handshake <= 65535").to_be_bytes()); + encrypted.extend_from_slice(initial_payload); + encrypt.apply_keystream(&mut encrypted); + write.write_all(&encrypted).await?; + + // PadB is plaintext of unknown length followed by the encrypted VC. Rather + // than probing each offset with a cloned decrypt state (libtorrent does the + // same), compute the encrypted-VC pattern once and search the raw bytes for + // it (matching libtorrent `read_pe_syncvc` / transmission `read_vc`). The + // VC is the first 8 keystream bytes of the decrypt stream, so after the + // search `decrypt` continues from exactly the right position. + let vc_pattern = { + let mut pattern = [0u8; VC_LEN]; + let mut probe = Rc4::new_from_slice(&decrypt_key).expect("20-byte RC4 key"); + let mut drop = [0u8; 1024]; + probe.apply_keystream(&mut drop); + probe.apply_keystream(&mut pattern); + pattern + }; + let mut raw = Vec::with_capacity(MAX_PAD + VC_LEN); + while raw.len() < MAX_PAD + VC_LEN { + let mut byte = [0u8; 1]; + read.read_exact(&mut byte).await?; + raw.push(byte[0]); + if raw.len() >= VC_LEN && raw[raw.len() - VC_LEN..] == vc_pattern { + break; + } + } + if raw.len() < VC_LEN || raw[raw.len() - VC_LEN..] != vc_pattern { + warn!("MSE verification constant not found within PadB, aborting handshake"); + bail!("MSE verification constant not found within PadB"); + } + let mut decrypt = decrypt_base; + // The VC is the first 8 keystream bytes of the decrypt stream. Decrypt the + // actual received VC ciphertext (raw's tail) and confirm it is the all-zero + // verification constant, which also advances the stream to the + // crypto-select / pad-length fields that follow. + let mut vc = raw[raw.len() - VC_LEN..].to_vec(); + decrypt.apply_keystream(&mut vc); + if vc != [0u8; VC_LEN] { + warn!("MSE invalid verification constant after PadB, aborting handshake"); + bail!("MSE invalid verification constant"); + } + + let mut select = [0u8; 4]; + read_encrypted(&mut read, &mut decrypt, &mut select).await?; + if u32::from_be_bytes(select) != CRYPTO_RC4 { + warn!("MSE responder did not select RC4, aborting handshake"); + bail!("MSE responder did not select RC4"); + } + let mut pad_length = [0u8; 2]; + read_encrypted(&mut read, &mut decrypt, &mut pad_length).await?; + let pad_length = u16::from_be_bytes(pad_length) as usize; + if pad_length > MAX_PAD { + bail!("MSE PadD exceeds {MAX_PAD} bytes"); + } + let mut pad_d = vec![0u8; pad_length]; + read_encrypted(&mut read, &mut decrypt, &mut pad_d).await?; + + trace!("MSE outgoing handshake complete, RC4 established"); + Ok(OutgoingOutcome::Encrypted( + Rc4Reader::new(read, decrypt), + Rc4Writer::new(write, encrypt), + )) +} + +/// Accept either a complete plaintext BitTorrent handshake or MSE. A +/// nonmatching partial plaintext prefix is retained as the beginning of YA. +pub async fn incoming( + mut read: R, + mut write: W, + lookup: F, +) -> Result> +where + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, + F: Fn(&[u8; 20]) -> Option<[u8; 20]>, +{ + let mut prefix = Vec::with_capacity(BT_PROTOCOL_PREFIX.len()); + while prefix.len() < BT_PROTOCOL_PREFIX.len() { + let mut byte = [0u8; 1]; + read.read_exact(&mut byte).await?; + prefix.push(byte[0]); + if prefix != BT_PROTOCOL_PREFIX[..prefix.len()] { + break; + } + } + if prefix.len() == BT_PROTOCOL_PREFIX.len() { + debug!("peer sent plaintext BT handshake, using plaintext path"); + return Ok(IncomingOutcome::Plaintext { + read: PrefixReader::new(prefix, read), + write, + }); + } + + debug!("peer did not send plaintext prefix, attempting MSE handshake"); + + let mut client_public = [0u8; 96]; + client_public[..prefix.len()].copy_from_slice(&prefix); + read.read_exact(&mut client_public[prefix.len()..]).await?; + + let dh = Dh768::generate(&mut rand::rng()); + let secret = dh + .shared_secret(&client_public) + .ok_or_else(|| anyhow::anyhow!("MSE degenerate remote DH key"))?; + + // The responder sends YB + PadB immediately, before waiting for req1. + write.write_all(&dh.public_key_bytes()).await?; + write.write_all(&random_pad(MAX_PAD)).await?; + + let req1 = sha1(&[b"req1", &secret]); + read_scan_for_needle(&mut read, &req1, MAX_PAD).await?; + let mut obfuscated_skey = [0u8; 20]; + read.read_exact(&mut obfuscated_skey).await?; + let skey_hash = xor20(&obfuscated_skey, &sha1(&[b"req3", &secret])); + let info_hash = + lookup(&skey_hash).ok_or_else(|| anyhow::anyhow!("MSE unknown info hash in SKEY"))?; + let (mut encrypt, mut decrypt, _decrypt_key) = derive_keys(&secret, &info_hash, false); + + let mut vc = [0u8; VC_LEN]; + read_encrypted(&mut read, &mut decrypt, &mut vc).await?; + if vc != [0u8; VC_LEN] { + warn!("MSE invalid verification constant, aborting handshake"); + bail!("MSE invalid verification constant"); + } + let mut provide = [0u8; 4]; + read_encrypted(&mut read, &mut decrypt, &mut provide).await?; + if u32::from_be_bytes(provide) & CRYPTO_RC4 == 0 { + bail!("MSE peer does not offer RC4"); + } + let mut pad_length = [0u8; 2]; + read_encrypted(&mut read, &mut decrypt, &mut pad_length).await?; + let pad_length = u16::from_be_bytes(pad_length) as usize; + if pad_length > MAX_PAD { + bail!("MSE PadC exceeds {MAX_PAD} bytes"); + } + let mut pad_c = vec![0u8; pad_length]; + read_encrypted(&mut read, &mut decrypt, &mut pad_c).await?; + + let mut ia_length = [0u8; 2]; + read_encrypted(&mut read, &mut decrypt, &mut ia_length).await?; + let ia_length = u16::from_be_bytes(ia_length) as usize; + if ia_length > BT_HANDSHAKE_LEN { + bail!("MSE IA length exceeds {BT_HANDSHAKE_LEN} bytes"); + } + let mut handshake_bytes = vec![0u8; ia_length]; + read_encrypted(&mut read, &mut decrypt, &mut handshake_bytes).await?; + + // Respond before waiting for the rest of the BT handshake. An initiator + // with IA=0 may not send that data until PE4 selects the cipher. + let pad_d = random_pad(MAX_PAD); + let mut response = Vec::with_capacity(VC_LEN + 4 + 2 + pad_d.len()); + response.extend_from_slice(&[0u8; VC_LEN]); + response.extend_from_slice(&CRYPTO_RC4.to_be_bytes()); + response.extend_from_slice(&u16::try_from(pad_d.len()).expect("pad_d <= MAX_PAD").to_be_bytes()); + response.extend_from_slice(&pad_d); + encrypt.apply_keystream(&mut response); + write.write_all(&response).await?; + + let mut remaining = vec![0u8; BT_HANDSHAKE_LEN - ia_length]; + read_encrypted(&mut read, &mut decrypt, &mut remaining).await?; + handshake_bytes.extend_from_slice(&remaining); + + trace!("MSE incoming handshake complete, RC4 established"); + Ok(IncomingOutcome::Encrypted { + read: Rc4Reader::new(read, decrypt), + write: Rc4Writer::new(write, encrypt), + handshake_bytes, + info_hash, + }) +} + +#[cfg(test)] +mod tests; diff --git a/crates/librqbit/src/mse/stream.rs b/crates/librqbit/src/mse/stream.rs new file mode 100644 index 000000000..4166c66b7 --- /dev/null +++ b/crates/librqbit/src/mse/stream.rs @@ -0,0 +1,268 @@ +//! Transparent RC4 stream wrappers for post-handshake traffic. + +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use rc4::{Rc4, StreamCipher}; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + +pub struct Rc4Reader { + inner: R, + rc4: Rc4, +} + +impl Rc4Reader { + pub fn new(inner: R, rc4: Rc4) -> Self { + Self { inner, rc4 } + } +} + +impl AsyncRead for Rc4Reader { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + let before = buf.filled().len(); + match Pin::new(&mut this.inner).poll_read(cx, buf) { + Poll::Ready(Ok(())) => { + this.rc4.apply_keystream(&mut buf.filled_mut()[before..]); + Poll::Ready(Ok(())) + } + other => other, + } + } +} + +pub struct Rc4Writer { + inner: W, + rc4: Rc4, + /// Ciphertext already produced (whose keystream has been consumed) but not + /// yet flushed to the underlying stream. Set when a write is Pending or + /// partially accepted; the caller must not be asked to re-encrypt. + pending: Vec, +} + +impl Rc4Writer { + pub fn new(inner: W, rc4: Rc4) -> Self { + Self { + inner, + rc4, + pending: Vec::new(), + } + } + + /// Flush any pending ciphertext to the underlying writer, consuming it. + /// Returns `Poll::Ready(Ok(()))` once everything is flushed, `Poll::Pending` + /// when the underlying writer is not writable. + fn flush_pending(&mut self, cx: &mut Context<'_>) -> Poll> + where + W: AsyncWrite + Unpin, + { + let this = self; + loop { + if this.pending.is_empty() { + return Poll::Ready(Ok(())); + } + match Pin::new(&mut this.inner).poll_write(cx, &this.pending) { + Poll::Ready(Ok(n)) => { + if n == 0 { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::WriteZero, + "failed to write RC4 stream", + ))); + } + // Consume exactly the accepted prefix; the remaining + // ciphertext stays pending for the next poll. + this.pending.drain(..n); + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + } + } +} + +impl AsyncWrite for Rc4Writer { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + data: &[u8], + ) -> Poll> { + let this = self.get_mut(); + if data.is_empty() { + return Poll::Ready(Ok(0)); + } + + // If a previous poll left ciphertext pending (this is a retry after a + // Pending or partial flush), flush it and report the logical write as + // complete without re-encrypting the caller's data. + if !this.pending.is_empty() { + match this.flush_pending(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Ready(Ok(())) => return Poll::Ready(Ok(data.len())), + } + } + + // Encrypt the whole buffer once, buffer the ciphertext, then flush it. + let mut encrypted = data.to_vec(); + this.rc4.apply_keystream(&mut encrypted); + this.pending = encrypted; + match this.flush_pending(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + Poll::Ready(Ok(())) => Poll::Ready(Ok(data.len())), + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + match this.flush_pending(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + Poll::Ready(Ok(())) => Pin::new(&mut this.inner).poll_flush(cx), + } + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + match this.flush_pending(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + Poll::Ready(Ok(())) => Pin::new(&mut this.inner).poll_shutdown(cx), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rc4::KeyInit; + use std::collections::VecDeque; + use std::sync::{Arc, Mutex}; + use tokio::io::AsyncWriteExt; + + enum Action { + Pending, + Limit(usize), + Error, + } + + struct FaultWriter { + actions: VecDeque, + bytes: Arc>>, + } + + impl AsyncWrite for FaultWriter { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + data: &[u8], + ) -> Poll> { + match self.actions.pop_front() { + Some(Action::Pending) => { + cx.waker().wake_by_ref(); + Poll::Pending + } + Some(Action::Error) => Poll::Ready(Err(io::Error::other( + "injected write failure", + ))), + Some(Action::Limit(limit)) => { + let count = limit.min(data.len()); + if let Ok(mut bytes) = self.bytes.lock() { + bytes.extend_from_slice(&data[..count]); + } + Poll::Ready(Ok(count)) + } + None => { + if let Ok(mut bytes) = self.bytes.lock() { + bytes.extend_from_slice(data); + } + Poll::Ready(Ok(data.len())) + } + } + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + fn expected(key: &[u8], plaintext: &[u8]) -> Vec { + let mut bytes = plaintext.to_vec(); + rc4::Rc4::new_from_slice(key) + .unwrap() + .apply_keystream(&mut bytes); + bytes + } + + #[tokio::test] + async fn pending_and_short_writes_preserve_state() -> io::Result<()> { + let sink = Arc::new(Mutex::new(Vec::new())); + let inner = FaultWriter { + actions: VecDeque::from([Action::Pending, Action::Limit(3)]), + bytes: sink.clone(), + }; + let key = b"writer-state"; + let mut writer = Rc4Writer::new(inner, rc4::Rc4::new_from_slice(key).unwrap()); + writer.write_all(b"abcdefgh").await?; + let actual = sink + .lock() + .map_err(|_| io::Error::other("poisoned test lock"))? + .clone(); + assert_eq!(actual, expected(key, b"abcdefgh")); + Ok(()) + } + + #[tokio::test] + async fn write_error_propagates() -> io::Result<()> { + let sink = Arc::new(Mutex::new(Vec::new())); + let inner = FaultWriter { + actions: VecDeque::from([Action::Error]), + bytes: sink.clone(), + }; + let key = b"writer-error"; + let mut writer = Rc4Writer::new(inner, rc4::Rc4::new_from_slice(key).unwrap()); + // An underlying write error propagates to the caller; nothing is + // accepted by the sink. (Unlike the self-contained rc4 impl we can't + // guarantee the RC4 keystream is un-advanced after a failed write, and + // in practice a write error tears down the connection anyway.) + assert!(writer.write(b"discarded").await.is_err()); + assert_eq!( + sink.lock() + .map_err(|_| io::Error::other("poisoned test lock"))? + .len(), + 0 + ); + Ok(()) + } + #[tokio::test] + async fn repeated_pending_does_not_reencrypt() -> io::Result<()> { + // The underlying writer returns Pending several times before accepting + // data. Rc4Writer must not re-encrypt the caller's plaintext on retry + // (the ciphertext is buffered once), and the final bytes must match + // the single-pass RC4 encryption. + let sink = Arc::new(Mutex::new(Vec::new())); + let inner = FaultWriter { + actions: VecDeque::from([Action::Pending, Action::Pending, Action::Pending]), + bytes: sink.clone(), + }; + let key = b"pending-retry"; + let mut writer = Rc4Writer::new(inner, rc4::Rc4::new_from_slice(key).unwrap()); + writer.write_all(b"persistent data").await?; + let actual = sink + .lock() + .map_err(|_| io::Error::other("poisoned test lock"))? + .clone(); + assert_eq!(actual, expected(key, b"persistent data")); + Ok(()) + } + +} \ No newline at end of file diff --git a/crates/librqbit/src/mse/tests.rs b/crates/librqbit/src/mse/tests.rs new file mode 100644 index 000000000..b55561e61 --- /dev/null +++ b/crates/librqbit/src/mse/tests.rs @@ -0,0 +1,298 @@ +use rc4::KeyInit; +use super::*; +use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex}; + +#[test] +fn default_mse_mode_is_disabled() { + // The default must stay Disabled so merging the feature is a zero-behavior + // change; users opt in via SessionOptions::mse_mode. + assert_eq!(MseMode::default(), MseMode::Disabled); +} + +fn handshake(info_hash: [u8; 20], peer_id: [u8; 20]) -> [u8; 68] { + let mut bytes = [0u8; 68]; + bytes[..20].copy_from_slice(BT_PROTOCOL_PREFIX); + bytes[28..48].copy_from_slice(&info_hash); + bytes[48..].copy_from_slice(&peer_id); + bytes +} + +#[tokio::test] +async fn incoming_accepts_zero_length_ia_before_full_handshake() -> Result<()> { + let info_hash = [0x42; 20]; + let expected_skey_hash = sha1(&[b"req2", &info_hash]); + let expected_handshake = handshake(info_hash, *b"-RQ0001-012345678901"); + let (client, server) = duplex(4096); + let (mut client_read, mut client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + + let responder = tokio::spawn(async move { + incoming(server_read, server_write, move |skey_hash| { + (*skey_hash == expected_skey_hash).then_some(info_hash) + }) + .await + }); + + let initiator_dh = Dh768::from_secret([0x37; 20]); + client_write + .write_all(&initiator_dh.public_key_bytes()) + .await?; + let mut responder_public = [0u8; 96]; + client_read.read_exact(&mut responder_public).await?; + let secret = initiator_dh + .shared_secret(&responder_public) + .context("responder returned an invalid DH key")?; + let (mut encrypt, mut decrypt, decrypt_key) = derive_keys(&secret, &info_hash, true); + + client_write.write_all(&sha1(&[b"req1", &secret])).await?; + let req3 = sha1(&[b"req3", &secret]); + client_write + .write_all(&xor20(&expected_skey_hash, &req3)) + .await?; + + let mut pe3 = Vec::new(); + pe3.extend_from_slice(&[0u8; VC_LEN]); + pe3.extend_from_slice(&CRYPTO_RC4.to_be_bytes()); + pe3.extend_from_slice(&0u16.to_be_bytes()); + pe3.extend_from_slice(&0u16.to_be_bytes()); + encrypt.apply_keystream(&mut pe3); + client_write.write_all(&pe3).await?; + + let mut encrypted_vc = [0u8; VC_LEN]; + // The initiator's decrypt stream and the responder's VC encrypt stream are + // the same keystream; rebuild an independent instance (rc4 crate has no + // Clone) to compute the encrypted-VC pattern for the needle scan. + let mut vc_probe = ::rc4::Rc4::new_from_slice(&decrypt_key).expect("20-byte RC4 key"); + let mut drop = [0u8; 1024]; + vc_probe.apply_keystream(&mut drop); + vc_probe.apply_keystream(&mut encrypted_vc); + read_scan_for_needle(&mut client_read, &encrypted_vc, MAX_PAD).await?; + decrypt.apply_keystream(&mut encrypted_vc); + assert_eq!(encrypted_vc, [0u8; VC_LEN]); + + let mut crypto_select = [0u8; 4]; + client_read.read_exact(&mut crypto_select).await?; + decrypt.apply_keystream(&mut crypto_select); + assert_eq!(u32::from_be_bytes(crypto_select), CRYPTO_RC4); + + let mut pad_d_length = [0u8; 2]; + client_read.read_exact(&mut pad_d_length).await?; + decrypt.apply_keystream(&mut pad_d_length); + let mut pad_d = vec![0u8; u16::from_be_bytes(pad_d_length) as usize]; + client_read.read_exact(&mut pad_d).await?; + decrypt.apply_keystream(&mut pad_d); + + let mut encrypted_handshake = expected_handshake; + encrypt.apply_keystream(&mut encrypted_handshake); + client_write.write_all(&encrypted_handshake).await?; + let payload = b"post-handshake payload"; + let mut encrypted_payload = *payload; + encrypt.apply_keystream(&mut encrypted_payload); + client_write.write_all(&encrypted_payload).await?; + + let outcome = responder.await??; + match outcome { + IncomingOutcome::Encrypted { + mut read, + handshake_bytes, + info_hash: resolved_info_hash, + .. + } => { + assert_eq!(resolved_info_hash, info_hash); + assert_eq!(handshake_bytes, expected_handshake); + let mut received_payload = [0u8; 22]; + read.read_exact(&mut received_payload).await?; + assert_eq!(&received_payload, payload); + } + IncomingOutcome::Plaintext { .. } => bail!("expected encrypted outcome"), + } + Ok(()) +} + +#[tokio::test] +async fn fragmented_plaintext_prefix_is_replayed() -> Result<()> { + let info_hash = [0x23; 20]; + let bytes = handshake(info_hash, [0x45; 20]); + let (client, server) = duplex(256); + let (server_read, server_write) = tokio::io::split(server); + let sender = async move { + let mut client = client; + for byte in bytes { + client.write_all(&[byte]).await?; + tokio::task::yield_now().await; + } + Ok::<_, std::io::Error>(()) + }; + let receiver = async move { + let outcome = incoming(server_read, server_write, |_| None).await?; + let mut read = match outcome { + IncomingOutcome::Plaintext { read, .. } => read, + IncomingOutcome::Encrypted { .. } => bail!("unexpected encrypted outcome"), + }; + let mut replayed = [0u8; 68]; + read.read_exact(&mut replayed).await?; + assert_eq!(replayed, bytes); + Ok::<_, anyhow::Error>(()) + }; + let (sent, received) = tokio::join!(sender, receiver); + sent?; + received?; + Ok(()) +} +#[tokio::test] +async fn duplex_handshake_preserves_payload() -> Result<()> { + let info_hash = [0x42; 20]; + let initial = handshake(info_hash, [0x11; 20]); + let (client, server) = duplex(8192); + let (client_read, client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + + let initiator = outgoing(client_read, client_write, &info_hash, &initial); + let responder = incoming(server_read, server_write, |candidate| { + (candidate == &sha1(&[b"req2", &info_hash])).then_some(info_hash) + }); + let (initiator_result, responder_result) = tokio::join!(initiator, responder); + let (mut client_read, mut client_write) = match initiator_result? { + OutgoingOutcome::Encrypted(r, w) => (r, w), + OutgoingOutcome::PlaintextPeer => bail!("unexpected plaintext peer"), + }; + let outcome = responder_result?; + let (mut server_read, mut server_write, received) = match outcome { + IncomingOutcome::Encrypted { + read, + write, + handshake_bytes, + .. + } => (read, write, handshake_bytes), + IncomingOutcome::Plaintext { .. } => bail!("unexpected plaintext outcome"), + }; + assert_eq!(received, initial); + + client_write.write_all(b"client payload").await?; + let mut client_payload = [0u8; 14]; + server_read.read_exact(&mut client_payload).await?; + assert_eq!(&client_payload, b"client payload"); + + server_write.write_all(b"server payload").await?; + let mut server_payload = [0u8; 14]; + client_read.read_exact(&mut server_payload).await?; + assert_eq!(&server_payload, b"server payload"); + Ok(()) +} + +#[tokio::test] +async fn plaintext_first_response_triggers_immediate_fallback() -> Result<()> { + // A plaintext peer answers our Ya + PadA with its 68-byte BT handshake + // immediately. `outgoing` must detect the `\x13BitTorrent protocol` + // prefix and return `PlaintextPeer` well within the 10s read timeout + // (2s sniff window is the ceiling here). + let info_hash = [0x42; 20]; + let (client, server) = duplex(8192); + let (client_read, client_write) = tokio::io::split(client); + let (mut server_read, mut server_write) = tokio::io::split(server); + let responder = async move { + // Read and discard Ya + PadA. + let mut discard = [0u8; 96]; + server_read.read_exact(&mut discard).await?; + // Reply with a plaintext BT handshake. + server_write + .write_all(&handshake(info_hash, [0x55; 20])) + .await?; + Ok::<_, std::io::Error>(()) + }; + + let initial = handshake(info_hash, [0x11; 20]); + let initiator = async { + let started = std::time::Instant::now(); + let outcome = outgoing(client_read, client_write, &info_hash, &initial).await?; + let elapsed = started.elapsed(); + match outcome { + OutgoingOutcome::PlaintextPeer => { + assert!( + elapsed < PLAINTEXT_SNIFF_TIMEOUT + Duration::from_millis(500), + "plaintext fallback took {elapsed:?}, expected <= {PLAINTEXT_SNIFF_TIMEOUT:?}" + ); + Ok::<_, anyhow::Error>(()) + } + OutgoingOutcome::Encrypted(..) => bail!("expected plaintext peer fallback"), + } + }; + + let (init, resp) = tokio::join!(initiator, responder); + init?; + resp?; + Ok(()) +} + +#[tokio::test] +async fn mse_responder_still_works_after_sniff() -> Result<()> { + // The sniff reads the first 20 bytes; an MSE responder's public key + // must still arrive intact when it does not match the BT prefix. + let info_hash = [0x42; 20]; + let initial = handshake(info_hash, [0x11; 20]); + let (client, server) = duplex(8192); + let (client_read, client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + + let initiator = outgoing(client_read, client_write, &info_hash, &initial); + let responder = incoming(server_read, server_write, |candidate| { + (candidate == &sha1(&[b"req2", &info_hash])).then_some(info_hash) + }); + let (initiator_result, responder_result) = tokio::join!(initiator, responder); + let (mut client_read, mut client_write) = match initiator_result? { + OutgoingOutcome::Encrypted(r, w) => (r, w), + OutgoingOutcome::PlaintextPeer => bail!("unexpected plaintext peer"), + }; + let (mut server_read, mut server_write, received) = match responder_result? { + IncomingOutcome::Encrypted { + read, + write, + handshake_bytes, + .. + } => (read, write, handshake_bytes), + IncomingOutcome::Plaintext { .. } => bail!("unexpected plaintext outcome"), + }; + assert_eq!(received, initial); + client_write.write_all(b"ping").await?; + let mut buf = [0u8; 4]; + server_read.read_exact(&mut buf).await?; + assert_eq!(&buf, b"ping"); + server_write.write_all(b"pong").await?; + client_read.read_exact(&mut buf).await?; + assert_eq!(&buf, b"pong"); + Ok(()) +} + +#[tokio::test] +async fn padb_without_vc_pattern_aborts_handshake() -> Result<()> { + // If the responder never sends the encrypted VC within MAX_PAD bytes, the + // PadB pattern-search must fail the handshake rather than hang or accept + // a corrupted stream. This exercises the search bound + the formal VC + // verification added for pattern misdetection. + let info_hash = [0x42; 20]; + let (client, server) = duplex(4096); + let (client_read, client_write) = tokio::io::split(client); + let (mut server_read, mut server_write) = tokio::io::split(server); + + let responder = async move { + // Read and discard Ya + PadA. + let mut ya = [0u8; 96]; + server_read.read_exact(&mut ya).await?; + // Send a fixed public key (0x01..) and a large PadB with no VC pattern. + let yb = [0x01u8; 96]; + server_write.write_all(&yb).await?; + let padb = vec![0x11u8; MAX_PAD + VC_LEN]; + server_write.write_all(&padb).await?; + Ok::<_, std::io::Error>(()) + }; + + let initial = handshake(info_hash, [0x11; 20]); + let initiator = outgoing(client_read, client_write, &info_hash, &initial); + let (initiator_result, responder_result) = tokio::join!(initiator, responder); + responder_result?; + assert!( + initiator_result.is_err(), + "outgoing should fail when VC is not found within PadB" + ); + Ok(()) +} diff --git a/crates/librqbit/src/peer_connection.rs b/crates/librqbit/src/peer_connection.rs index 6bfabe86e..4331cc3c6 100644 --- a/crates/librqbit/src/peer_connection.rs +++ b/crates/librqbit/src/peer_connection.rs @@ -4,7 +4,9 @@ use std::{ time::{Duration, Instant}, }; -use crate::{Error, Result, session::CheckedIncomingConnection, stream_connect::ConnectionKind}; +use crate::{ + Error, Result, mse::MseMode, session::CheckedIncomingConnection, stream_connect::ConnectionKind, +}; use buffers::{ByteBuf, ByteBufOwned}; use futures::TryFutureExt; use librqbit_core::{ @@ -77,6 +79,10 @@ pub struct PeerConnectionOptions { #[serde_as(as = "Option")] pub keep_alive_interval: Option, + + /// How MSE (Message Stream Encryption) is applied to this peer connection. + /// Defaults to [`MseMode::Disabled`]. + pub mse_mode: MseMode, } pub(crate) struct PeerConnection { @@ -205,7 +211,6 @@ impl PeerConnection { outgoing_chan: tokio::sync::mpsc::UnboundedReceiver, have_broadcast: tokio::sync::broadcast::Receiver, ) -> Result<()> { - use tokio::io::AsyncWriteExt; let rwtimeout = self .options .read_write_timeout @@ -217,27 +222,22 @@ impl PeerConnection { .unwrap_or_else(|| Duration::from_secs(10)); let now = Instant::now(); - let (ckind, mut read, mut write) = with_timeout( - "connecting", - connect_timeout, - self.connector.connect(self.addr), - ) - .await?; - - async move { - self.handler.on_connected(now.elapsed()); - - let mut write_buf = Box::new([0u8; MAX_MSG_LEN]); - let handshake = Handshake::new(self.info_hash, self.peer_id); - let hsz = handshake.serialize_unchecked_len(&mut *write_buf); - with_timeout( - "writing", + let conn = self + .connector + .connect_with_handshake( + self.addr, + &self.info_hash.0, + &self.peer_id.0, + connect_timeout, rwtimeout, - write - .write_all(&write_buf[..hsz]) - .map_err(Error::WriteHandshake), + self.options.mse_mode, ) .await?; + let (ckind, mut read, write) = (conn.kind, conn.read, conn.write); + + async move { + self.handler.on_connected(now.elapsed()); + let write_buf = Box::new([0u8; MAX_MSG_LEN]); let mut read_buf = ReadBuf::new(); let h = read_buf.read_handshake(&mut read, rwtimeout).await?; @@ -513,3 +513,145 @@ impl PeerConnection { } } } +#[cfg(test)] +mod mse_fallback_tests { + use super::*; + use crate::vectored_traits::AsyncReadVectoredIntoCompat; + use tokio::io::{AsyncReadExt, duplex}; + + #[tokio::test] + async fn fresh_redial_fallback_uses_a_new_stream() -> anyhow::Result<()> { + let (first_client, mut first_peer) = duplex(4096); + let (second_client, mut second_peer) = duplex(4096); + let (first_read, first_write) = tokio::io::split(first_client); + let (second_read, second_write) = tokio::io::split(second_client); + let connector = StreamConnector::with_test_connections(vec![ + ( + Box::new(first_read.into_vectored_compat()), + Box::new(first_write), + ), + ( + Box::new(second_read.into_vectored_compat()), + Box::new(second_write), + ), + ]); + let peer = async move { + // First connection: read Ya (96 bytes) then drop, so MSE fails. + let mut first_attempt = [0u8; 96]; + first_peer.read_exact(&mut first_attempt).await?; + drop(first_peer); + // Second connection: expect the plaintext BT handshake that + // connect_with_handshake wrote on the redial. + let mut plaintext = [0u8; 68]; + second_peer.read_exact(&mut plaintext).await?; + assert_eq!(&plaintext[..20], b"\x13BitTorrent protocol"); + assert_eq!(&plaintext[28..48], &[0x42; 20]); + assert_eq!(&plaintext[48..], &[0x11; 20]); + Ok::<_, std::io::Error>(()) + }; + let client = async { + let conn = connector + .connect_with_handshake( + "127.0.0.1:1".parse()?, + &[0x42; 20], + &[0x11; 20], + Duration::from_secs(1), + Duration::from_secs(1), + MseMode::Enabled, + ) + .await?; + assert!(!conn.mse_applied, "MSE should have failed and fallen back"); + // connect_with_handshake already wrote the plaintext handshake on + // the redial; drain so the peer's read completes. + drop(conn.read); + assert_eq!(connector.remaining_test_connections()?, 0); + Ok::<_, anyhow::Error>(()) + }; + let (client_result, peer_result) = tokio::join!(client, peer); + client_result?; + peer_result?; + Ok(()) + } + + #[tokio::test] + async fn disabled_skips_mse_and_uses_single_connection() -> anyhow::Result<()> { + let (client, mut peer) = duplex(4096); + let (read, write) = tokio::io::split(client); + let connector = StreamConnector::with_test_connections(vec![( + Box::new(read.into_vectored_compat()), + Box::new(write), + )]); + let peer = async move { + // Disabled MSE: the peer must receive the plaintext handshake + // directly (68 bytes), never Ya + PadA (96+ bytes first). + let mut plaintext = [0u8; 68]; + peer.read_exact(&mut plaintext).await?; + assert_eq!(&plaintext[..20], b"\x13BitTorrent protocol"); + assert_eq!(&plaintext[28..48], &[0x42; 20]); + assert_eq!(&plaintext[48..], &[0x11; 20]); + Ok::<_, std::io::Error>(()) + }; + let client = async { + let conn = connector + .connect_with_handshake( + "127.0.0.1:1".parse()?, + &[0x42; 20], + &[0x11; 20], + Duration::from_secs(1), + Duration::from_secs(1), + MseMode::Disabled, + ) + .await?; + assert!(!conn.mse_applied, "MSE must not be attempted in Disabled mode"); + // connect_with_handshake wrote the plaintext handshake already. + drop(conn.read); + // Only one connection consumed: mse::outgoing was never invoked. + assert_eq!(connector.remaining_test_connections()?, 0); + Ok::<_, anyhow::Error>(()) + }; + let (client_result, peer_result) = tokio::join!(client, peer); + client_result?; + peer_result?; + Ok(()) + } + + #[tokio::test] + async fn forced_mse_failure_returns_error_without_redial() -> anyhow::Result<()> { + let (client, mut peer) = duplex(4096); + let (read, write) = tokio::io::split(client); + let connector = StreamConnector::with_test_connections(vec![( + Box::new(read.into_vectored_compat()), + Box::new(write), + )]); + let peer = async move { + // Read Ya (96 bytes) then drop, so MSE fails. + let mut ya = [0u8; 96]; + peer.read_exact(&mut ya).await?; + drop(peer); + Ok::<_, std::io::Error>(()) + }; + let client = async { + let result = connector + .connect_with_handshake( + "127.0.0.1:1".parse()?, + &[0x42; 20], + &[0x11; 20], + Duration::from_secs(1), + Duration::from_secs(1), + MseMode::Forced, + ) + .await; + assert!( + matches!(result, Err(Error::MseForced)), + "Forced mode must error on MSE failure" + ); + // No redial in Forced mode: exactly one connection was consumed. + assert_eq!(connector.remaining_test_connections()?, 0); + Ok::<_, anyhow::Error>(()) + }; + let (client_result, peer_result) = tokio::join!(client, peer); + client_result?; + peer_result?; + Ok(()) + } +} diff --git a/crates/librqbit/src/session.rs b/crates/librqbit/src/session.rs index 75071b9c5..92ef0bd9f 100644 --- a/crates/librqbit/src/session.rs +++ b/crates/librqbit/src/session.rs @@ -23,6 +23,7 @@ use crate::{ limits::{Limits, LimitsConfig}, listen::{Accept, ListenerOptions}, merge_streams::merge_streams, + mse::MseMode, peer_connection::PeerConnectionOptions, read_buf::ReadBuf, session_persistence::{SessionPersistenceStore, json::JsonSessionPersistenceStore}, @@ -32,7 +33,8 @@ use crate::{ BoxStorageFactory, StorageFactoryExt, TorrentStorage, filesystem::FilesystemStorageFactory, }, stream_connect::{ - ConnectionKind, ConnectionOptions, SocksProxyConfig, StreamConnector, StreamConnectorArgs, + ConnectionKind, ConnectionOptions, IncomingHandshake, SocksProxyConfig, StreamConnector, + StreamConnectorArgs, }, torrent_state::{ ManagedTorrentHandle, ManagedTorrentLocked, ManagedTorrentOptions, ManagedTorrentState, @@ -67,6 +69,7 @@ use librqbit_utp::BindDevice; use parking_lot::RwLock; use peer_binary_protocol::Handshake; use serde::{Deserialize, Serialize}; +use sha1w::{ISha1, Sha1}; use tokio::sync::Notify; use tokio_util::sync::{CancellationToken, DropGuard}; use tracing::{Instrument, debug, debug_span, error, info, trace, warn}; @@ -151,6 +154,7 @@ pub struct Session { pub ipv4_only: bool, pub peer_limit: Option, client_name_and_version: String, + mse_mode: MseMode, } async fn torrent_from_url( @@ -479,6 +483,10 @@ pub struct SessionOptions { /// Override the client name and version used in User-Agent headers and /// peer extended handshakes. Defaults to "rqbit X.Y.Z". pub client_name_and_version: Option, + + /// How MSE (Message Stream Encryption) is applied to peer connections. + /// Defaults to [`MseMode::Disabled`]. + pub mse_mode: MseMode, } impl Default for SessionOptions { @@ -507,6 +515,7 @@ impl Default for SessionOptions { disable_local_service_discovery: false, ipv4_only: false, client_name_and_version: None, + mse_mode: MseMode::default(), } } } @@ -632,11 +641,15 @@ impl Session { } else { None }; - let peer_opts = opts + let mut peer_opts = opts .connect .as_ref() .and_then(|p| p.peer_opts) .unwrap_or_default(); + // The session-level MSE mode is authoritative for outgoing + // connections (including magnet metadata reads), regardless of any + // per-connection PeerConnectionOptions. + peer_opts.mse_mode = opts.mse_mode; async fn persistence_factory( opts: &SessionOptions, @@ -806,6 +819,7 @@ impl Session { disable_trackers: opts.disable_trackers, peer_limit: opts.peer_limit, client_name_and_version, + mse_mode: opts.mse_mode, #[cfg(feature = "disable-upload")] _disable_upload: opts.disable_upload, @@ -905,7 +919,7 @@ impl Session { self: Arc, addr: SocketAddr, kind: ConnectionKind, - mut reader: BoxAsyncReadVectored, + reader: BoxAsyncReadVectored, writer: BoxAsyncWrite, ) -> anyhow::Result<(Arc, CheckedIncomingConnection)> { let rwtimeout = self @@ -929,12 +943,56 @@ impl Session { bail!("Incoming ip {incoming_ip} is not in allowlist"); } - let mut read_buf = ReadBuf::new(); - let h = read_buf - .read_handshake(&mut reader, rwtimeout) - .await - .context("error reading handshake")?; - trace!("received handshake from {addr}: {:?}", h); + // Snapshot the (SKEY hash -> info_hash) mapping for every known torrent + // so the MSE acceptor can resolve the peer's obfuscated info hash. + let torrent_keys: Vec<([u8; 20], [u8; 20])> = self + .db + .read() + .torrents + .values() + .map(|torrent| { + let info_hash = torrent.info_hash().0; + let mut h = Sha1::new(); + h.update(b"req2"); + h.update(&info_hash); + (h.finish(), info_hash) + }) + .collect(); + + let lookup = |skey_hash: &[u8; 20]| -> Option<[u8; 20]> { + torrent_keys + .iter() + .find(|(k, _)| k == skey_hash) + .map(|(_, info_hash)| *info_hash) + }; + + let incoming = crate::stream_connect::accept_with_handshake( + addr, + reader, + writer, + rwtimeout, + self.mse_mode, + lookup, + ) + .await?; + + self.finish_incoming_connection(addr, kind, incoming).await + } + + async fn finish_incoming_connection( + self: Arc, + addr: SocketAddr, + kind: ConnectionKind, + incoming: IncomingHandshake, + ) -> anyhow::Result<(Arc, CheckedIncomingConnection)> { + let IncomingHandshake { + handshake: h, + read: reader, + write: writer, + read_buf, + encrypted, + } = incoming; + trace!("received handshake from {addr}: {:?} (encrypted={encrypted})", h); if h.peer_id == self.peer_id { bail!("seems like we are connecting to ourselves, ignoring"); @@ -1037,6 +1095,7 @@ impl Session { keep_alive_interval: other .keep_alive_interval .or(self.peer_opts.keep_alive_interval), + mse_mode: self.mse_mode, } } @@ -1255,8 +1314,10 @@ impl Session { let peer_rx = make_peer_rx().context( "no known way to resolve peers (no DHT, no trackers, no initial_peers)", )?; + let mut magnet_peer_opts = opts.peer_opts.unwrap_or_default(); + magnet_peer_opts.mse_mode = self.mse_mode; let resolved_magnet = self - .resolve_magnet(info_hash, peer_rx, &trackers, opts.peer_opts) + .resolve_magnet(info_hash, peer_rx, &trackers, Some(magnet_peer_opts)) .await?; // Add back seen_peers into the peer stream, as we consumed some peers @@ -1350,6 +1411,7 @@ impl Session { force_tracker_interval: opts.force_tracker_interval, peer_connect_timeout: peer_opts.connect_timeout, peer_read_write_timeout: peer_opts.read_write_timeout, + mse_mode: self.mse_mode, allow_overwrite: opts.overwrite, output_folder, ratelimits: opts.ratelimits, diff --git a/crates/librqbit/src/stream_connect.rs b/crates/librqbit/src/stream_connect.rs index ab9b0e9b1..3bb986fa4 100644 --- a/crates/librqbit/src/stream_connect.rs +++ b/crates/librqbit/src/stream_connect.rs @@ -4,13 +4,17 @@ use anyhow::{Context, bail}; use librqbit_dualstack_sockets::ConnectOpts; use librqbit_utp::{BindDevice, UtpSocketUdp}; use serde::Serialize; -use tracing::debug; use crate::{ - Error, PeerConnectionOptions, Result, + Error, PeerConnectionOptions, Result, mse::{MseMode, OutgoingOutcome}, + peer_connection::with_timeout, + read_buf::ReadBuf, type_aliases::{BoxAsyncReadVectored, BoxAsyncWrite}, vectored_traits::AsyncReadVectoredIntoCompat, }; +use peer_binary_protocol::Handshake; +use tokio::io::AsyncWriteExt; +use tracing::{debug, warn}; #[derive(Debug, Clone, Copy, Serialize)] pub enum ConnectionKind { @@ -51,6 +55,29 @@ impl Default for ConnectionOptions { } } +/// Result of connecting and (optionally) performing an MSE handshake. +pub struct OutgoingHandshake { + pub kind: ConnectionKind, + pub read: BoxAsyncReadVectored, + pub write: BoxAsyncWrite, + /// True when the MSE handshake succeeded and consumed the BT handshake as + /// its initial payload (IA); the caller must not write a plaintext + /// handshake in that case. + #[allow(dead_code)] // read by the MSE fallback tests + pub mse_applied: bool, +} + +/// Result of accepting a connection and reading its (plaintext or MSE) BT +/// handshake. +pub(crate) struct IncomingHandshake { + pub handshake: Handshake, + pub read: BoxAsyncReadVectored, + pub write: BoxAsyncWrite, + pub read_buf: ReadBuf, + /// True when the handshake arrived over an MSE-encrypted stream. + pub encrypted: bool, +} + #[derive(Debug, Clone)] pub(crate) struct SocksProxyConfig { pub host: String, @@ -125,7 +152,6 @@ gen_stats!(ConnectStatsAtomic ConnectStatsSnapshot, [], [ utp PerFamilyAtomic PerFamilySnapshot ]); -#[derive(Debug)] pub(crate) struct StreamConnector { proxy_config: Option, enable_tcp: bool, @@ -133,9 +159,43 @@ pub(crate) struct StreamConnector { utp_socket: Option>, stats: ConnectStatsAtomic, ipv4_only: bool, + // Pre-built connections consumed by `connect()` in test builds, in order. + #[cfg(test)] + test_connections: + parking_lot::Mutex>, +} + +impl std::fmt::Debug for StreamConnector { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("StreamConnector") + .field("proxy_config", &self.proxy_config) + .field("enable_tcp", &self.enable_tcp) + .field("bind_device", &self.bind_device) + .field("utp_socket", &self.utp_socket) + .field("ipv4_only", &self.ipv4_only) + .finish() + } } impl StreamConnector { + #[cfg(test)] + pub fn with_test_connections(conns: Vec<(BoxAsyncReadVectored, BoxAsyncWrite)>) -> Self { + Self { + proxy_config: None, + enable_tcp: true, + utp_socket: None, + bind_device: None, + stats: Default::default(), + ipv4_only: false, + test_connections: parking_lot::Mutex::new(conns.into()), + } + } + + #[cfg(test)] + pub fn remaining_test_connections(&self) -> anyhow::Result { + Ok(self.test_connections.lock().len()) + } + pub async fn new(config: StreamConnectorArgs) -> anyhow::Result { #[allow(clippy::single_match)] match ( @@ -158,6 +218,8 @@ impl StreamConnector { bind_device: config.bind_device, stats: Default::default(), ipv4_only: config.ipv4_only, + #[cfg(test)] + test_connections: Default::default(), }) } @@ -211,6 +273,11 @@ impl StreamConnector { &self, addr: SocketAddr, ) -> Result<(ConnectionKind, BoxAsyncReadVectored, BoxAsyncWrite)> { + #[cfg(test)] + if let Some((r, w)) = self.test_connections.lock().pop_front() { + return Ok((ConnectionKind::Tcp, r, w)); + } + if addr.port() == 0 { return Err(Error::Anyhow(anyhow::anyhow!( "invalid peer address (port 0): {}", @@ -323,4 +390,214 @@ impl StreamConnector { }; } } + + /// Connect to `addr` and complete the BitTorrent handshake, optionally + /// wrapping the stream in MSE (Message Stream Encryption) per `mse_mode`. + /// + /// Serializes the 68-byte BT handshake once; when MSE succeeds it is + /// consumed as the MSE initial payload (IA), otherwise it is written + /// plaintext. On MSE failure or a plaintext peer, redials plaintext on a + /// fresh connection (never reuses the MSE-polluted stream). `Forced` mode + /// errors instead of downgrading. The returned streams are always + /// handshake-complete: the caller must not write a BT handshake again. + pub async fn connect_with_handshake( + &self, + addr: SocketAddr, + info_hash: &[u8; 20], + peer_id: &[u8; 20], + connect_timeout: Duration, + rwtimeout: Duration, + mse_mode: MseMode, + ) -> Result { + let (ckind, read, write) = + with_timeout("connecting", connect_timeout, self.connect(addr)).await?; + + let handshake = Handshake::new( + librqbit_core::hash_id::Id20::new(*info_hash), + librqbit_core::hash_id::Id20::new(*peer_id), + ); + let mut write_buf = [0u8; 68]; + let hsz = handshake.serialize_unchecked_len(&mut write_buf); + + if mse_mode != MseMode::Disabled { + let mse_outcome = match tokio::time::timeout( + rwtimeout, + crate::mse::outgoing(read, write, info_hash, &write_buf), + ) + .await + { + Ok(Ok(outcome)) => outcome, + Ok(Err(e)) => { + if mse_mode == MseMode::Forced { + warn!(?addr, "MSE forced but handshake failed, dropping peer: {e:#}"); + return Err(Error::MseForced); + } + debug!(?addr, "MSE handshake failed, redialing plaintext: {e:#}"); + return self + .redial_plaintext(addr, connect_timeout, rwtimeout, &write_buf, hsz) + .await; + } + Err(_elapsed) => { + if mse_mode == MseMode::Forced { + warn!(?addr, "MSE forced but handshake timed out, dropping peer"); + return Err(Error::MseForced); + } + debug!(?addr, "MSE handshake timed out, redialing plaintext"); + return self + .redial_plaintext(addr, connect_timeout, rwtimeout, &write_buf, hsz) + .await; + } + }; + + match mse_outcome { + OutgoingOutcome::Encrypted(r, w) => { + debug!(?addr, "MSE handshake succeeded, RC4 established"); + return Ok(OutgoingHandshake { + kind: ckind, + read: Box::new(r.into_vectored_compat()), + write: Box::new(w), + mse_applied: true, + }); + } + OutgoingOutcome::PlaintextPeer => { + if mse_mode == MseMode::Forced { + warn!(?addr, "MSE forced but peer answered plaintext, dropping peer"); + return Err(Error::MseForced); + } + debug!(?addr, "peer answered plaintext, redialing plaintext"); + return self + .redial_plaintext(addr, connect_timeout, rwtimeout, &write_buf, hsz) + .await; + } + } + } + + Self::finish_plaintext_handshake(ckind, read, write, &write_buf, hsz, rwtimeout).await + } + + /// Redial a fresh plaintext connection after MSE failed or the peer + /// answered plaintext, writing the serialized BT handshake. + async fn redial_plaintext( + &self, + addr: SocketAddr, + connect_timeout: Duration, + rwtimeout: Duration, + write_buf: &[u8; 68], + hsz: usize, + ) -> Result { + let (nk, nr, nw) = + with_timeout("connecting", connect_timeout, self.connect(addr)).await?; + Self::finish_plaintext_handshake(nk, nr, nw, write_buf, hsz, rwtimeout).await + } + + /// Write the serialized BT handshake on a plaintext stream and wrap it. + async fn finish_plaintext_handshake( + ckind: ConnectionKind, + read: BoxAsyncReadVectored, + write: BoxAsyncWrite, + write_buf: &[u8; 68], + hsz: usize, + rwtimeout: Duration, + ) -> Result { + use futures::TryFutureExt; + let mut write = write; + with_timeout( + "writing", + rwtimeout, + write.write_all(&write_buf[..hsz]).map_err(Error::WriteHandshake), + ) + .await?; + Ok(OutgoingHandshake { kind: ckind, read, write, mse_applied: false }) + } +} + +/// Accept an incoming connection and read its BT handshake, optionally +/// attempting the MSE (Message Stream Encryption) handshake per `mse_mode`. +/// +/// With MSE enabled, probes the peer for a plaintext BT prefix; if it is a +/// plaintext peer it is accepted unless `mse_mode == Forced`, otherwise the +/// MSE acceptor runs (resolving the obfuscated info hash via `lookup`). The +/// returned streams are ready for post-handshake traffic. +pub(crate) async fn accept_with_handshake( + addr: SocketAddr, + reader: BoxAsyncReadVectored, + writer: BoxAsyncWrite, + rwtimeout: Duration, + mse_mode: MseMode, + lookup: F, +) -> anyhow::Result +where + F: Fn(&[u8; 20]) -> Option<[u8; 20]>, +{ + use crate::mse::IncomingOutcome; + use tokio::io::AsyncReadExt; + + if mse_mode == MseMode::Disabled { + let mut read_buf = ReadBuf::new(); + let mut reader = reader; + let h = read_buf + .read_handshake(&mut reader, rwtimeout) + .await + .context("error reading handshake")?; + return Ok(IncomingHandshake { + handshake: h, + read: reader, + write: writer, + read_buf, + encrypted: false, + }); + } + + let incoming = crate::mse::incoming(reader, writer, lookup); + let incoming = tokio::time::timeout(rwtimeout, incoming) + .await + .context("MSE incoming handshake timed out")??; + + match incoming { + IncomingOutcome::Encrypted { + read, + write, + handshake_bytes, + info_hash, + } => { + let (h, _size) = Handshake::deserialize(&handshake_bytes[..]) + .map_err(|e| anyhow::anyhow!("error deserializing MSE handshake: {e:?}"))?; + if h.info_hash.0 != info_hash { + bail!("MSE handshake info hash does not match SKEY"); + } + Ok(IncomingHandshake { + handshake: h, + read: Box::new(read.into_vectored_compat()), + write: Box::new(write), + read_buf: ReadBuf::new(), + encrypted: true, + }) + } + IncomingOutcome::Plaintext { read, write } => { + if mse_mode == MseMode::Forced { + warn!(?addr, "MSE forced, rejecting plaintext connection"); + bail!("MSE is forced, rejecting plaintext connection from {addr}"); + } + // 9.0's `ReadBuf::read_handshake` performs a single `read()` + // before deserializing; a `PrefixReader` replaying the 20 + // consumed prefix bytes would satisfy it with only those bytes. + // Read the full 68-byte handshake ourselves (read_exact loops), + // mirroring the Encrypted branch. + let mut handshake_bytes = [0u8; 68]; + let mut read = read; + read.read_exact(&mut handshake_bytes) + .await + .context("error reading fragmented plaintext handshake")?; + let (h, _size) = Handshake::deserialize(&handshake_bytes[..]).map_err(|e| { + anyhow::anyhow!("error deserializing plaintext handshake: {e:?}") + })?; + Ok(IncomingHandshake { + handshake: h, + read: Box::new(read.into_vectored_compat()), + write: Box::new(write), + read_buf: ReadBuf::new(), + encrypted: false, + }) + } + } } diff --git a/crates/librqbit/src/torrent_state/live/mod.rs b/crates/librqbit/src/torrent_state/live/mod.rs index b4146d8a8..af70c5734 100644 --- a/crates/librqbit/src/torrent_state/live/mod.rs +++ b/crates/librqbit/src/torrent_state/live/mod.rs @@ -487,6 +487,7 @@ impl TorrentStateLive { let options = PeerConnectionOptions { connect_timeout: self.shared.options.peer_connect_timeout, read_write_timeout: self.shared.options.peer_read_write_timeout, + mse_mode: self.shared.options.mse_mode, ..Default::default() }; let peer_connection = PeerConnection::new( @@ -552,6 +553,7 @@ impl TorrentStateLive { let options = PeerConnectionOptions { connect_timeout: state.shared.options.peer_connect_timeout, read_write_timeout: state.shared.options.peer_read_write_timeout, + mse_mode: state.shared.options.mse_mode, ..Default::default() }; let peer_connection = PeerConnection::new( diff --git a/crates/librqbit/src/torrent_state/mod.rs b/crates/librqbit/src/torrent_state/mod.rs index cfe483cdd..a29630bfd 100644 --- a/crates/librqbit/src/torrent_state/mod.rs +++ b/crates/librqbit/src/torrent_state/mod.rs @@ -112,6 +112,7 @@ pub(crate) struct ManagedTorrentOptions { pub force_tracker_interval: Option, pub peer_connect_timeout: Option, pub peer_read_write_timeout: Option, + pub mse_mode: crate::mse::MseMode, pub allow_overwrite: bool, pub output_folder: PathBuf, pub ratelimits: LimitsConfig, diff --git a/crates/rqbit/src/main.rs b/crates/rqbit/src/main.rs index a171274f2..39d40eb21 100644 --- a/crates/rqbit/src/main.rs +++ b/crates/rqbit/src/main.rs @@ -691,6 +691,7 @@ async fn async_main(mut opts: Opts, cancel: CancellationToken) -> anyhow::Result runtime_worker_threads: Some(opts.max_blocking_threads as usize), ipv4_only: opts.ipv4_only, client_name_and_version: None, + mse_mode: Default::default(), }; #[allow(clippy::needless_update)]