Enhance TLS Emulator with ALPN Support and Add Adversarial Tests

- Modified `build_emulated_server_hello` to accept ALPN (Application-Layer Protocol Negotiation) as an optional parameter, allowing for the embedding of ALPN markers in the application data payload.
- Implemented logic to handle oversized ALPN values and ensure they do not interfere with the application data payload.
- Added new security tests in `emulator_security_tests.rs` to validate the behavior of the ALPN embedding, including scenarios for oversized ALPN and preference for certificate payloads over ALPN markers.
- Introduced `send_adversarial_tests.rs` to cover edge cases and potential issues in the middle proxy's send functionality, ensuring robustness against various failure modes.
- Updated `middle_proxy` module to include new test modules and ensure proper handling of writer commands during data transmission.

src/transport/middle_proxy/mod.rs

@@ -27,6 +27,8 @@ mod health_regression_tests;
 mod health_integration_tests;
 #[cfg(test)]
 mod health_adversarial_tests;
+#[cfg(test)]
+mod send_adversarial_tests;
 
 use bytes::Bytes;
 

src/transport/middle_proxy/pool.rs

@@ -692,6 +692,7 @@ impl MePool {
         }
     }
 
+    #[allow(dead_code)]
     pub(super) fn draining_active_runtime(&self) -> u64 {
         self.draining_active_runtime.load(Ordering::Relaxed)
     }

src/transport/middle_proxy/registry.rs

@@ -454,6 +454,7 @@ impl ConnRegistry {
         true
     }
 
+    #[allow(dead_code)]
     pub(super) async fn non_empty_writer_ids(&self, writer_ids: &[u64]) -> HashSet<u64> {
         let inner = self.inner.read().await;
         let mut out = HashSet::<u64>::with_capacity(writer_ids.len());

src/transport/middle_proxy/send.rs

@@ -372,17 +372,20 @@ impl MePool {
                 }
                 let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port());
                 let (payload, meta) = build_routed_payload(effective_our_addr);
-                match w.tx.try_send(WriterCommand::Data(payload.clone())) {
-                    Ok(()) => {
-                        self.stats.increment_me_writer_pick_success_try_total(pick_mode);
+                match w.tx.clone().try_reserve_owned() {
+                    Ok(permit) => {
                         if !self.registry.bind_writer(conn_id, w.id, meta).await {
                             debug!(
                                 conn_id,
                                 writer_id = w.id,
-                                "ME writer disappeared before bind commit, retrying"
+                                "ME writer disappeared before bind commit, pruning stale writer"
                             );
+                            drop(permit);
+                            self.remove_writer_and_close_clients(w.id).await;
                             continue;
                         }
+                        permit.send(WriterCommand::Data(payload.clone()));
+                        self.stats.increment_me_writer_pick_success_try_total(pick_mode);
                         if w.generation < self.current_generation() {
                             self.stats.increment_pool_stale_pick_total();
                             debug!(
@@ -422,18 +425,21 @@ impl MePool {
             self.stats.increment_me_writer_pick_blocking_fallback_total();
             let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port());
             let (payload, meta) = build_routed_payload(effective_our_addr);
-            match w.tx.send(WriterCommand::Data(payload.clone())).await {
-                Ok(()) => {
-                    self.stats
-                        .increment_me_writer_pick_success_fallback_total(pick_mode);
+            match w.tx.clone().reserve_owned().await {
+                Ok(permit) => {
                     if !self.registry.bind_writer(conn_id, w.id, meta).await {
                         debug!(
                             conn_id,
                             writer_id = w.id,
-                            "ME writer disappeared before fallback bind commit, retrying"
+                            "ME writer disappeared before fallback bind commit, pruning stale writer"
                         );
+                        drop(permit);
+                        self.remove_writer_and_close_clients(w.id).await;
                         continue;
                     }
+                    permit.send(WriterCommand::Data(payload.clone()));
+                    self.stats
+                        .increment_me_writer_pick_success_fallback_total(pick_mode);
                     if w.generation < self.current_generation() {
                         self.stats.increment_pool_stale_pick_total();
                     }

src/transport/middle_proxy/send_adversarial_tests.rs

@@ -0,0 +1,263 @@
+use std::collections::HashMap;
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering};
+use std::time::{Duration, Instant};
+
+use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
+
+use super::codec::WriterCommand;
+use super::pool::{MePool, MeWriter, WriterContour};
+use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode};
+use crate::crypto::SecureRandom;
+use crate::network::probe::NetworkDecision;
+use crate::stats::Stats;
+
+async fn make_pool() -> (Arc<MePool>, Arc<SecureRandom>) {
+    let general = GeneralConfig {
+        me_route_no_writer_mode: MeRouteNoWriterMode::AsyncRecoveryFailfast,
+        me_route_no_writer_wait_ms: 50,
+        me_writer_pick_mode: MeWriterPickMode::SortedRr,
+        me_deterministic_writer_sort: true,
+        ..GeneralConfig::default()
+    };
+
+    let rng = Arc::new(SecureRandom::new());
+    let pool = MePool::new(
+        None,
+        vec![1u8; 32],
+        None,
+        false,
+        None,
+        Vec::new(),
+        1,
+        None,
+        12,
+        1200,
+        HashMap::new(),
+        HashMap::new(),
+        None,
+        NetworkDecision::default(),
+        None,
+        rng.clone(),
+        Arc::new(Stats::default()),
+        general.me_keepalive_enabled,
+        general.me_keepalive_interval_secs,
+        general.me_keepalive_jitter_secs,
+        general.me_keepalive_payload_random,
+        general.rpc_proxy_req_every,
+        general.me_warmup_stagger_enabled,
+        general.me_warmup_step_delay_ms,
+        general.me_warmup_step_jitter_ms,
+        general.me_reconnect_max_concurrent_per_dc,
+        general.me_reconnect_backoff_base_ms,
+        general.me_reconnect_backoff_cap_ms,
+        general.me_reconnect_fast_retry_count,
+        general.me_single_endpoint_shadow_writers,
+        general.me_single_endpoint_outage_mode_enabled,
+        general.me_single_endpoint_outage_disable_quarantine,
+        general.me_single_endpoint_outage_backoff_min_ms,
+        general.me_single_endpoint_outage_backoff_max_ms,
+        general.me_single_endpoint_shadow_rotate_every_secs,
+        general.me_floor_mode,
+        general.me_adaptive_floor_idle_secs,
+        general.me_adaptive_floor_min_writers_single_endpoint,
+        general.me_adaptive_floor_min_writers_multi_endpoint,
+        general.me_adaptive_floor_recover_grace_secs,
+        general.me_adaptive_floor_writers_per_core_total,
+        general.me_adaptive_floor_cpu_cores_override,
+        general.me_adaptive_floor_max_extra_writers_single_per_core,
+        general.me_adaptive_floor_max_extra_writers_multi_per_core,
+        general.me_adaptive_floor_max_active_writers_per_core,
+        general.me_adaptive_floor_max_warm_writers_per_core,
+        general.me_adaptive_floor_max_active_writers_global,
+        general.me_adaptive_floor_max_warm_writers_global,
+        general.hardswap,
+        general.me_pool_drain_ttl_secs,
+        general.me_pool_drain_threshold,
+        general.effective_me_pool_force_close_secs(),
+        general.me_pool_min_fresh_ratio,
+        general.me_hardswap_warmup_delay_min_ms,
+        general.me_hardswap_warmup_delay_max_ms,
+        general.me_hardswap_warmup_extra_passes,
+        general.me_hardswap_warmup_pass_backoff_base_ms,
+        general.me_bind_stale_mode,
+        general.me_bind_stale_ttl_secs,
+        general.me_secret_atomic_snapshot,
+        general.me_deterministic_writer_sort,
+        general.me_writer_pick_mode,
+        general.me_writer_pick_sample_size,
+        MeSocksKdfPolicy::default(),
+        general.me_writer_cmd_channel_capacity,
+        general.me_route_channel_capacity,
+        general.me_route_backpressure_base_timeout_ms,
+        general.me_route_backpressure_high_timeout_ms,
+        general.me_route_backpressure_high_watermark_pct,
+        general.me_reader_route_data_wait_ms,
+        general.me_health_interval_ms_unhealthy,
+        general.me_health_interval_ms_healthy,
+        general.me_warn_rate_limit_ms,
+        general.me_route_no_writer_mode,
+        general.me_route_no_writer_wait_ms,
+        general.me_route_inline_recovery_attempts,
+        general.me_route_inline_recovery_wait_ms,
+    );
+
+    (pool, rng)
+}
+
+async fn insert_writer(
+    pool: &Arc<MePool>,
+    writer_id: u64,
+    writer_dc: i32,
+    addr: SocketAddr,
+    register_in_registry: bool,
+) -> mpsc::Receiver<WriterCommand> {
+    let (tx, rx) = mpsc::channel::<WriterCommand>(8);
+    let writer = MeWriter {
+        id: writer_id,
+        addr,
+        source_ip: addr.ip(),
+        writer_dc,
+        generation: pool.current_generation(),
+        contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())),
+        created_at: Instant::now(),
+        tx: tx.clone(),
+        cancel: CancellationToken::new(),
+        degraded: Arc::new(AtomicBool::new(false)),
+        rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
+        draining: Arc::new(AtomicBool::new(false)),
+        draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)),
+        drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)),
+        allow_drain_fallback: Arc::new(AtomicBool::new(false)),
+    };
+
+    pool.writers.write().await.push(writer);
+    {
+        let mut map = pool.proxy_map_v4.write().await;
+        map.entry(writer_dc)
+            .or_insert_with(Vec::new)
+            .push((addr.ip(), addr.port()));
+    }
+    pool.rebuild_endpoint_dc_map().await;
+    if register_in_registry {
+        pool.registry.register_writer(writer_id, tx).await;
+    }
+    rx
+}
+
+async fn recv_data_count(rx: &mut mpsc::Receiver<WriterCommand>, budget: Duration) -> usize {
+    let start = Instant::now();
+    let mut data_count = 0usize;
+    while Instant::now().duration_since(start) < budget {
+        let remaining = budget.saturating_sub(Instant::now().duration_since(start));
+        match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await {
+            Ok(Some(WriterCommand::Data(_))) => data_count += 1,
+            Ok(Some(WriterCommand::DataAndFlush(_))) => data_count += 1,
+            Ok(Some(WriterCommand::Close)) => {}
+            Ok(None) => break,
+            Err(_) => break,
+        }
+    }
+    data_count
+}
+
+#[tokio::test]
+async fn send_proxy_req_does_not_replay_when_first_bind_commit_fails() {
+    let (pool, _rng) = make_pool().await;
+    pool.rr.store(0, Ordering::Relaxed);
+
+    let (conn_id, _rx) = pool.registry.register().await;
+    let mut stale_rx = insert_writer(
+        &pool,
+        10,
+        2,
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 10)), 443),
+        false,
+    )
+    .await;
+    let mut live_rx = insert_writer(
+        &pool,
+        11,
+        2,
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 11)), 443),
+        true,
+    )
+    .await;
+
+    let result = pool
+        .send_proxy_req(
+            conn_id,
+            2,
+            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 30000),
+            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443),
+            b"hello",
+            0,
+            None,
+        )
+        .await;
+
+    assert!(result.is_ok());
+    assert_eq!(recv_data_count(&mut stale_rx, Duration::from_millis(50)).await, 0);
+    assert_eq!(recv_data_count(&mut live_rx, Duration::from_millis(50)).await, 1);
+
+    let bound = pool.registry.get_writer(conn_id).await;
+    assert!(bound.is_some());
+    assert_eq!(bound.expect("writer should be bound").writer_id, 11);
+}
+
+#[tokio::test]
+async fn send_proxy_req_prunes_iterative_stale_bind_failures_without_data_replay() {
+    let (pool, _rng) = make_pool().await;
+    pool.rr.store(0, Ordering::Relaxed);
+
+    let (conn_id, _rx) = pool.registry.register().await;
+
+    let mut stale_rx_1 = insert_writer(
+        &pool,
+        21,
+        2,
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 21)), 443),
+        false,
+    )
+    .await;
+    let mut stale_rx_2 = insert_writer(
+        &pool,
+        22,
+        2,
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 22)), 443),
+        false,
+    )
+    .await;
+    let mut live_rx = insert_writer(
+        &pool,
+        23,
+        2,
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 23)), 443),
+        true,
+    )
+    .await;
+
+    let result = pool
+        .send_proxy_req(
+            conn_id,
+            2,
+            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 30001),
+            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443),
+            b"storm",
+            0,
+            None,
+        )
+        .await;
+
+    assert!(result.is_ok());
+    assert_eq!(recv_data_count(&mut stale_rx_1, Duration::from_millis(50)).await, 0);
+    assert_eq!(recv_data_count(&mut stale_rx_2, Duration::from_millis(50)).await, 0);
+    assert_eq!(recv_data_count(&mut live_rx, Duration::from_millis(50)).await, 1);
+
+    let writers = pool.writers.read().await;
+    let writer_ids = writers.iter().map(|w| w.id).collect::<Vec<_>>();
+    drop(writers);
+    assert_eq!(writer_ids, vec![23]);
+}