Security hardening, concurrency fixes, and expanded test coverage

This commit introduces a comprehensive set of improvements to enhance the security, reliability, and configurability of the proxy server, specifically targeting adversarial resilience and high-load concurrency. Security & Cryptography: - Zeroize MTProto cryptographic key material (`dec_key`, `enc_key`) immediately after use to prevent memory leakage on early returns. - Move TLS handshake replay tracking after full policy/ALPN validation to prevent cache poisoning by unauthenticated probes. - Add `proxy_protocol_trusted_cidrs` configuration to restrict PROXY protocol headers to trusted networks, rejecting spoofed IPs. Adversarial Resilience & DoS Mitigation: - Implement "Tiny Frame Debt" tracking in the middle-relay to prevent CPU exhaustion from malicious 0-byte or 1-byte frame floods. - Add `mask_relay_max_bytes` to strictly bound unauthenticated fallback connections, preventing the proxy from being abused as an open relay. - Add a 5ms prefetch window (`mask_classifier_prefetch_timeout_ms`) to correctly assemble and classify fragmented HTTP/1.1 and HTTP/2 probes (e.g., `PRI * HTTP/2.0`) before routing them to masking heuristics. - Prevent recursive masking loops (FD exhaustion) by verifying the mask target is not the proxy's own listener via local interface enumeration. Concurrency & Reliability: - Eliminate executor waker storms during quota lock contention by replacing the spin-waker task with inline `Sleep` and exponential backoff. - Roll back user quota reservations (`rollback_me2c_quota_reservation`) if a network write fails, preventing Head-of-Line (HoL) blocking from permanently burning data quotas. - Recover gracefully from idle-registry `Mutex` poisoning instead of panicking, ensuring isolated thread failures do not break the proxy. - Fix `auth_probe_scan_start_offset` modulo logic to ensure bounds safety. Testing: - Add extensive adversarial, timing, fuzzing, and invariant test suites for both the client and handshake modules.
2026-06-19 02:00:08 +07:00 · 2026-03-22 23:06:26 +04:00
parent 6fc188f0c4
commit 91be148b72
65 changed files with 7473 additions and 210 deletions
@@ -0,0 +1,354 @@
+use super::*;
+use std::net::TcpListener as StdTcpListener;
+use std::net::SocketAddr;
+use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex};
+use tokio::net::TcpListener;
+use tokio::time::{Duration, Instant, timeout};
+
+fn closed_local_port() -> u16 {
+    let listener = StdTcpListener::bind("127.0.0.1:0").unwrap();
+    let port = listener.local_addr().unwrap().port();
+    drop(listener);
+    port
+}
+
+#[test]
+fn self_target_detection_matches_literal_ipv4_listener() {
+    let local: SocketAddr = "198.51.100.40:443".parse().unwrap();
+    assert!(is_mask_target_local_listener(
+        "198.51.100.40",
+        443,
+        local,
+        None,
+    ));
+}
+
+#[test]
+fn self_target_detection_matches_bracketed_ipv6_listener() {
+    let local: SocketAddr = "[2001:db8::44]:8443".parse().unwrap();
+    assert!(is_mask_target_local_listener(
+        "[2001:db8::44]",
+        8443,
+        local,
+        None,
+    ));
+}
+
+#[test]
+fn self_target_detection_keeps_same_ip_different_port_forwardable() {
+    let local: SocketAddr = "203.0.113.44:443".parse().unwrap();
+    assert!(!is_mask_target_local_listener(
+        "203.0.113.44",
+        8443,
+        local,
+        None,
+    ));
+}
+
+#[test]
+fn self_target_detection_normalizes_ipv4_mapped_ipv6_literal() {
+    let local: SocketAddr = "127.0.0.1:443".parse().unwrap();
+    assert!(is_mask_target_local_listener(
+        "::ffff:127.0.0.1",
+        443,
+        local,
+        None,
+    ));
+}
+
+#[test]
+fn self_target_detection_unspecified_bind_blocks_loopback_target() {
+    let local: SocketAddr = "0.0.0.0:443".parse().unwrap();
+    assert!(is_mask_target_local_listener(
+        "127.0.0.1",
+        443,
+        local,
+        None,
+    ));
+}
+
+#[test]
+fn self_target_detection_unspecified_bind_keeps_remote_target_forwardable() {
+    let local: SocketAddr = "0.0.0.0:443".parse().unwrap();
+    let remote: SocketAddr = "198.51.100.44:443".parse().unwrap();
+    assert!(!is_mask_target_local_listener(
+        "mask.example",
+        443,
+        local,
+        Some(remote),
+    ));
+}
+
+#[tokio::test]
+async fn self_target_fallback_refuses_recursive_loopback_connect() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let local_addr = listener.local_addr().unwrap();
+    let accept_task = tokio::spawn(async move {
+        timeout(Duration::from_millis(120), listener.accept())
+            .await
+            .is_ok()
+    });
+
+    let mut config = ProxyConfig::default();
+    config.general.beobachten = false;
+    config.censorship.mask = true;
+    config.censorship.mask_unix_sock = None;
+    config.censorship.mask_host = Some(local_addr.ip().to_string());
+    config.censorship.mask_port = local_addr.port();
+    config.censorship.mask_proxy_protocol = 0;
+
+    let peer: SocketAddr = "203.0.113.90:55090".parse().unwrap();
+    let beobachten = BeobachtenStore::new();
+
+    handle_bad_client(
+        tokio::io::empty(),
+        tokio::io::sink(),
+        b"GET /",
+        peer,
+        local_addr,
+        &config,
+        &beobachten,
+    )
+    .await;
+
+    let accepted = accept_task.await.unwrap();
+    assert!(
+        !accepted,
+        "self-target masking must fail closed without connecting to local listener"
+    );
+}
+
+#[tokio::test]
+async fn same_ip_different_port_still_forwards_to_mask_backend() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let backend_addr = listener.local_addr().unwrap();
+    let probe = b"GET /".to_vec();
+    let accept_task = tokio::spawn({
+        let expected = probe.clone();
+        async move {
+            let (mut stream, _) = listener.accept().await.unwrap();
+            let mut got = vec![0u8; expected.len()];
+            stream.read_exact(&mut got).await.unwrap();
+            assert_eq!(got, expected);
+        }
+    });
+
+    let mut config = ProxyConfig::default();
+    config.general.beobachten = false;
+    config.censorship.mask = true;
+    config.censorship.mask_unix_sock = None;
+    config.censorship.mask_host = Some("127.0.0.1".to_string());
+    config.censorship.mask_port = backend_addr.port();
+    config.censorship.mask_proxy_protocol = 0;
+
+    let peer: SocketAddr = "203.0.113.91:55091".parse().unwrap();
+    let local_addr: SocketAddr = "127.0.0.1:443".parse().unwrap();
+    let beobachten = BeobachtenStore::new();
+
+    handle_bad_client(
+        tokio::io::empty(),
+        tokio::io::sink(),
+        &probe,
+        peer,
+        local_addr,
+        &config,
+        &beobachten,
+    )
+    .await;
+
+    timeout(Duration::from_secs(2), accept_task)
+        .await
+        .unwrap()
+        .unwrap();
+}
+
+#[test]
+fn detect_client_type_http_boundary_get_and_post() {
+    assert_eq!(detect_client_type(b"GET "), "HTTP");
+    assert_eq!(detect_client_type(b"GET /"), "HTTP");
+
+    assert_eq!(detect_client_type(b"POST"), "HTTP");
+    assert_eq!(detect_client_type(b"POST "), "HTTP");
+    assert_eq!(detect_client_type(b"POSTX"), "HTTP");
+}
+
+#[test]
+fn detect_client_type_tls_and_length_boundaries() {
+    assert_eq!(detect_client_type(b"\x16\x03\x01"), "port-scanner");
+    assert_eq!(detect_client_type(b"\x16\x03\x01\x00"), "TLS-scanner");
+
+    assert_eq!(detect_client_type(b"123456789"), "port-scanner");
+    assert_eq!(detect_client_type(b"1234567890"), "unknown");
+}
+
+#[test]
+fn build_mask_proxy_header_v1_cross_family_falls_back_to_unknown() {
+    let peer: SocketAddr = "192.168.1.5:12345".parse().unwrap();
+    let local: SocketAddr = "[2001:db8::1]:443".parse().unwrap();
+    let header = build_mask_proxy_header(1, peer, local).unwrap();
+    assert_eq!(header, b"PROXY UNKNOWN\r\n");
+}
+
+#[test]
+fn next_mask_shape_bucket_checked_mul_overflow_fails_closed() {
+    let floor = usize::MAX / 2 + 1;
+    let cap = usize::MAX;
+    let total = floor + 1;
+    assert_eq!(next_mask_shape_bucket(total, floor, cap), total);
+}
+
+#[tokio::test]
+async fn self_target_reject_path_keeps_timing_budget() {
+    let mut config = ProxyConfig::default();
+    config.general.beobachten = false;
+    config.censorship.mask = true;
+    config.censorship.mask_unix_sock = None;
+    config.censorship.mask_host = Some("127.0.0.1".to_string());
+    config.censorship.mask_port = 443;
+
+    let local_addr: SocketAddr = "127.0.0.1:443".parse().unwrap();
+    let peer: SocketAddr = "203.0.113.92:55092".parse().unwrap();
+    let beobachten = BeobachtenStore::new();
+
+    let (client, server) = duplex(1024);
+    drop(client);
+
+    let started = Instant::now();
+    handle_bad_client(
+        server,
+        tokio::io::sink(),
+        b"GET / HTTP/1.1\r\n",
+        peer,
+        local_addr,
+        &config,
+        &beobachten,
+    )
+    .await;
+
+    let elapsed = started.elapsed();
+    assert!(
+        elapsed >= Duration::from_millis(40) && elapsed < Duration::from_millis(250),
+        "self-target reject path must keep coarse timing budget without stalling"
+    );
+}
+
+#[tokio::test]
+async fn relay_path_idle_timeout_eviction_remains_effective() {
+    let (client_read, mut client_write) = duplex(1024);
+    let (mask_read, mask_write) = duplex(1024);
+
+    tokio::spawn(async move {
+        tokio::time::sleep(Duration::from_millis(10)).await;
+        client_write.write_all(b"a").await.unwrap();
+        tokio::time::sleep(Duration::from_millis(180)).await;
+        let _ = client_write.write_all(b"b").await;
+    });
+
+    let started = Instant::now();
+    relay_to_mask(
+        client_read,
+        tokio::io::sink(),
+        mask_read,
+        mask_write,
+        b"init",
+        false,
+        0,
+        0,
+        false,
+        0,
+        false,
+        5 * 1024 * 1024,
+    )
+    .await;
+
+    let elapsed = started.elapsed();
+    assert!(
+        elapsed >= Duration::from_millis(90) && elapsed < Duration::from_millis(180),
+        "idle-timeout eviction must occur before late trickle write"
+    );
+}
+
+#[tokio::test]
+async fn offline_mask_target_refusal_respects_timing_normalization_budget() {
+    let mut config = ProxyConfig::default();
+    config.general.beobachten = false;
+    config.censorship.mask = true;
+    config.censorship.mask_unix_sock = None;
+    config.censorship.mask_host = Some("127.0.0.1".to_string());
+    config.censorship.mask_port = closed_local_port();
+    config.censorship.mask_timing_normalization_enabled = true;
+    config.censorship.mask_timing_normalization_floor_ms = 120;
+    config.censorship.mask_timing_normalization_ceiling_ms = 120;
+
+    let peer: SocketAddr = "203.0.113.93:55093".parse().unwrap();
+    let local_addr: SocketAddr = "127.0.0.1:443".parse().unwrap();
+    let beobachten = BeobachtenStore::new();
+
+    let (mut client, server) = duplex(1024);
+    let started = Instant::now();
+    let task = tokio::spawn(async move {
+        handle_bad_client(
+            server,
+            tokio::io::sink(),
+            b"GET / HTTP/1.1\r\n\r\n",
+            peer,
+            local_addr,
+            &config,
+            &beobachten,
+        )
+        .await;
+    });
+
+    client.shutdown().await.unwrap();
+    timeout(Duration::from_secs(2), task).await.unwrap().unwrap();
+    let elapsed = started.elapsed();
+
+    assert!(
+        elapsed >= Duration::from_millis(110) && elapsed < Duration::from_millis(220),
+        "offline-refusal path must honor normalization budget without unbounded drift"
+    );
+}
+
+#[tokio::test]
+async fn offline_mask_target_refusal_with_idle_client_is_bounded_by_consume_timeout() {
+    let mut config = ProxyConfig::default();
+    config.general.beobachten = false;
+    config.censorship.mask = true;
+    config.censorship.mask_unix_sock = None;
+    config.censorship.mask_host = Some("127.0.0.1".to_string());
+    config.censorship.mask_port = closed_local_port();
+    config.censorship.mask_timing_normalization_enabled = false;
+
+    let peer: SocketAddr = "203.0.113.94:55094".parse().unwrap();
+    let local_addr: SocketAddr = "127.0.0.1:443".parse().unwrap();
+    let beobachten = BeobachtenStore::new();
+
+    let (mut client, server) = duplex(1024);
+    let started = Instant::now();
+    let task = tokio::spawn(async move {
+        handle_bad_client(
+            server,
+            tokio::io::sink(),
+            b"GET / HTTP/1.1\r\n\r\n",
+            peer,
+            local_addr,
+            &config,
+            &beobachten,
+        )
+        .await;
+    });
+
+    tokio::time::sleep(Duration::from_millis(120)).await;
+    client
+        .write_all(b"still-open-before-timeout")
+        .await
+        .expect("connection should still be open before consume timeout expires");
+
+    timeout(Duration::from_secs(2), task).await.unwrap().unwrap();
+    let elapsed = started.elapsed();
+
+    assert!(
+        elapsed >= Duration::from_millis(190) && elapsed < Duration::from_millis(350),
+        "offline-refusal path must not retain idle client indefinitely"
+    );
+}