use bytes::Bytes; use rand::Rng; use thiserror::Error; use std::collections::{BTreeMap, VecDeque}; use std::time::{Duration, Instant}; /// Upper bound on a single frame's retransmit timer, after exponential backoff /// is applied to the adaptive RTO. Past this the session is dead from the /// user's point of view, and waiting longer only delays recovery. const MAX_EFFECTIVE_RTO: Duration = Duration::from_secs(8); use crate::congestion::CongestionController; use crate::crypto::{NoiseRole, NoiseSession, SessionCipher}; use crate::framing::{AdaptivePadder, FrameHeader, FrameKind, FramedPacket, PaddingStrategy}; #[derive(Debug, Error)] pub enum ProtocolError { #[error("state error: {0}")] State(String), #[error("crypto error: {0}")] Crypto(String), #[error("framing error: {0}")] Framing(String), } #[derive(Debug, Clone)] pub struct ProtocolConfig { pub role: NoiseRole, pub psk: [u8; 32], pub session_id: u32, pub handshake_payload: Vec, pub max_padding: usize, pub padding_strategy: PaddingStrategy, pub obfuscation_key: [u8; 8], pub max_reorder: u64, pub max_reorder_buffer: usize, pub ack_delay_ms: u64, pub rto_ms: u64, pub max_retries: u8, pub max_sent_history: usize, /// Key-derived handshake padding range (Kerckhoffs's principle). /// Different access keys produce different handshake packet sizes. pub handshake_pad_min: usize, pub handshake_pad_max: usize, pub mtu: usize, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum OstpState { Init, Handshaking, Established, Closing, Closed, } pub enum OstpEvent { Start, Inbound(Bytes), Outbound(u16, Bytes), // stream_id, payload Close, Tick, } pub enum ProtocolAction { SendDatagram(Bytes), // Fully formed datagram to send globally DeliverApp(u16, Bytes), // stream_id, payload HandshakePayload(Bytes, Option), // Passed from client's handshake, Optional response to send Multiple(Vec), Noop, } pub struct ProtocolMachine { role: NoiseRole, state: OstpState, noise: NoiseSession, send_cipher: Option, recv_cipher: Option, send_nonce: u64, expected_recv_nonce: u64, reorder_buffer: BTreeMap, sent_history: VecDeque, session_id: u32, handshake_payload: Vec, padder: AdaptivePadder, obfuscation_key: [u8; 8], max_reorder: u64, max_reorder_buffer: usize, ack_delay: Duration, rto: Duration, max_retries: u8, max_sent_history: usize, ack_pending: bool, last_ack_sent: Instant, /// Rate-limit: prevents sending a NACK more than once per 30ms to avoid storms last_nack_sent: Instant, /// Tracks when expected_recv_nonce last advanced. Used for gap recovery: /// if the receiver is stuck waiting for a lost frame that the sender already /// evicted from sent_history, this timer detects the deadlock and skips /// the gap to restore liveness. last_recv_advance: Instant, /// Congestion controller (BBR-inspired adaptive window) cc: CongestionController, /// Key-derived handshake padding range handshake_pad_min: usize, handshake_pad_max: usize, _mtu: usize, } // ── Gap recovery (see `ProtocolMachine::recover_stalled_gap`) ──────────────── // How long the receive sequence may sit stuck behind a missing frame, with // later frames already buffered, before that frame is declared unrecoverable // and skipped. Derived from the live RTO so it scales with the path instead of // guessing, then clamped: the floor keeps a fast link from discarding a frame // that is merely late, the ceiling bounds how long a stall can be visible to // the user before the tunnel unblocks itself. const GAP_RECOVERY_RTO_MULTIPLIER: u32 = 8; const GAP_RECOVERY_MIN: Duration = Duration::from_secs(2); const GAP_RECOVERY_MAX: Duration = Duration::from_secs(10); #[derive(Debug, Clone)] struct SentFrame { nonce: u64, bytes: Bytes, last_sent: Instant, retries: u8, is_retransmittable: bool, } impl ProtocolMachine { pub fn new(config: ProtocolConfig) -> Result { let noise = NoiseSession::new( config.role, &config.psk, )?; Ok(Self { role: config.role, state: OstpState::Init, noise, send_cipher: None, recv_cipher: None, send_nonce: 0, expected_recv_nonce: 0, reorder_buffer: BTreeMap::new(), sent_history: VecDeque::with_capacity(config.max_sent_history.max(1)), session_id: config.session_id, handshake_payload: config.handshake_payload, padder: AdaptivePadder::new(config.mtu, config.max_padding, config.padding_strategy), obfuscation_key: config.obfuscation_key, max_reorder: config.max_reorder.max(1), max_reorder_buffer: config.max_reorder_buffer.max(1), ack_delay: Duration::from_millis(config.ack_delay_ms.max(1)), rto: Duration::from_millis(config.rto_ms.max(1)), max_retries: config.max_retries.max(1), max_sent_history: config.max_sent_history.max(1), ack_pending: false, last_ack_sent: Instant::now(), last_nack_sent: Instant::now() - Duration::from_secs(1), last_recv_advance: Instant::now(), cc: CongestionController::new(config.mtu as u64), handshake_pad_min: config.handshake_pad_min.max(8), handshake_pad_max: config.handshake_pad_max.max(config.handshake_pad_min + 16), _mtu: config.mtu, }) } pub fn in_flight_count(&self) -> usize { // COUNT ONLY retransmittable Data frames — control frames (Ack/Nack) must not // contribute to this counter or they will trigger false backpressure. self.sent_history.iter().filter(|f| f.is_retransmittable).count() } /// Sum of retry counters across in-flight frames. Test-only: lets a test /// assert the core retransmit invariant (a retry is only ever charged to a /// frame that was actually put on the wire) without needing to advance the /// clock through several seconds of exponential backoff. #[cfg(test)] fn total_retries(&self) -> usize { self.sent_history .iter() .filter(|f| f.is_retransmittable) .map(|f| f.retries as usize) .sum() } pub fn cwnd_packets(&self) -> usize { self.cc.cwnd_packets() as usize } /// Whether the pacing bucket currently allows releasing another packet. /// /// The congestion window bounds how much may be UNACKNOWLEDGED; it says /// nothing about how fast that window is emptied onto the wire. Sending a /// whole window back-to-back is what drives a deep buffer into standing /// queue, so admission is gated on both. pub fn can_pace_packet(&self) -> bool { self.cc.can_pace_packet() } pub fn on_send(&mut self, bytes: u64) { self.cc.on_send(bytes); } pub fn state(&self) -> OstpState { self.state } pub fn on_event(&mut self, event: OstpEvent) -> Result { match (self.state, event) { (OstpState::Init, OstpEvent::Start) => { match self.role { NoiseRole::Initiator => { self.state = OstpState::Handshaking; let mut out = vec![0_u8; 1024]; let n = self.noise.write_handshake(&self.handshake_payload, &mut out)?; out.truncate(n); self.wrap_datagram_handshake(&out) .map(ProtocolAction::SendDatagram) } NoiseRole::Responder => { self.state = OstpState::Handshaking; Ok(ProtocolAction::Noop) } } } (OstpState::Init, OstpEvent::Inbound(raw)) => { self.state = OstpState::Handshaking; self.handle_inbound(raw) } (OstpState::Handshaking, OstpEvent::Inbound(raw)) => { self.handle_inbound(raw) } (OstpState::Handshaking, OstpEvent::Start) => Ok(ProtocolAction::Noop), (OstpState::Established, OstpEvent::Outbound(stream_id, app_data)) => { self.build_tracked_datagram(stream_id, FrameKind::Data, app_data) .map(ProtocolAction::SendDatagram) } (OstpState::Established, OstpEvent::Inbound(raw)) => { self.handle_inbound(raw) } (OstpState::Established, OstpEvent::Close) => { self.state = OstpState::Closing; self.build_tracked_datagram(0, FrameKind::Close, Bytes::new()) .map(ProtocolAction::SendDatagram) } (OstpState::Closing, OstpEvent::Inbound(raw)) => { // The remote may still have data or ACKs in transit when we initiated // Close. Stay in Closing and process them; handle_inbound transitions to // Closed only when it actually receives the peer's Close frame — the old // code force-closed after a single inbound packet, losing in-flight data. // (Ported from 0.3.x 47d44fa.) self.handle_inbound(raw) } (OstpState::Established, OstpEvent::Tick) => self.handle_tick(), // Retransmit our Close frame (and drain pending) while waiting for teardown. (OstpState::Closing, OstpEvent::Tick) => self.handle_tick(), (OstpState::Closed, _) => Ok(ProtocolAction::Noop), (_, OstpEvent::Close) => { self.state = OstpState::Closed; Ok(ProtocolAction::Noop) } _ => Ok(ProtocolAction::Noop), } } fn handle_inbound(&mut self, raw: Bytes) -> Result { let mut raw_vec = raw.to_vec(); let is_handshake = self.state == OstpState::Handshaking || self.state == OstpState::Init; crate::crypto::deobfuscate_packet_inplace(&mut raw_vec, &self.obfuscation_key, is_handshake); if raw_vec.len() < 4 { return Err(ProtocolError::Framing("datagram too short".to_string())); } let session_id = u32::from_be_bytes([raw_vec[0], raw_vec[1], raw_vec[2], raw_vec[3]]); if session_id != self.session_id { // Per-packet, attacker-triggerable event: keep at debug and don't // dump internal session ids (log-flood + info-leak surface). tracing::debug!("session id mismatch (is_handshake={})", is_handshake); return Err(ProtocolError::State("session id mismatch".to_string())); } if self.state == OstpState::Handshaking { self.handle_handshake_inbound(&raw_vec) } else if self.state == OstpState::Established { self.handle_data_inbound(&raw_vec) } else { Ok(ProtocolAction::Noop) } } fn handle_handshake_inbound(&mut self, raw_vec: &[u8]) -> Result { // Wire format: [session_id:4][noise_len:2][noise_payload:N][random_padding:*] // Extract noise_len to pass exactly the right bytes to snow if raw_vec.len() < 6 { return Err(ProtocolError::Framing("handshake too short for length prefix".to_string())); } let noise_len = u16::from_be_bytes([raw_vec[4], raw_vec[5]]) as usize; if raw_vec.len() < 6 + noise_len { return Err(ProtocolError::Framing(format!( "handshake truncated: expected {} noise bytes, got {}", noise_len, raw_vec.len() - 6 ))); } let mut read_out = vec![0_u8; 1024]; let n = self.noise.read_handshake(&raw_vec[6..6 + noise_len], &mut read_out).map_err(|e| { ProtocolError::Crypto(format!("noise-read: {:?} (raw_len={}, noise_len={})", e, raw_vec.len(), noise_len)) })?; read_out.truncate(n); let response = match self.role { NoiseRole::Responder => { let mut write_out = vec![0_u8; 1024]; let out_n = self.noise.write_handshake(&self.handshake_payload, &mut write_out)?; write_out.truncate(out_n); Some(self.wrap_datagram_handshake(&write_out)?) } NoiseRole::Initiator => None, }; // Transport keys come from Noise's Split() over the final chaining key, // so they depend on the ephemeral `ee` DH secret and give the session // forward secrecy. (Previously these were derived from the handshake // hash, which never absorbs the DH result — see raw_split's SECURITY // note. That is the wire-breaking change gated by PROTOCOL_VERSION.) let (send_key, recv_key) = self.noise.raw_split(self.role)?; self.send_cipher = Some(SessionCipher::new(&send_key)); self.recv_cipher = Some(SessionCipher::new(&recv_key)); self.state = OstpState::Established; let extracted_payload = read_out[..n].to_vec(); Ok(ProtocolAction::HandshakePayload(Bytes::from(extracted_payload), response)) } /// Restores liveness when the receive sequence is stuck behind a frame that /// can never arrive. /// /// Delivery is gated on `expected_recv_nonce`, so a single missing frame /// holds back every later frame. That is correct *while the sender can still /// retransmit* — but the sender drops a frame from `sent_history` once it /// exceeds `max_retries + 2` attempts (see the zombie eviction in /// `handle_tick`). After that the frame is gone for good and the two sides /// deadlock: the receiver buffers forever and NACKs a nonce nobody can /// resend. /// /// That deadlock is invisible to the keepalive watchdog, which is why it /// presented as a hard freeze rather than a reconnect: retransmits, ACKs and /// NACKs keep flowing, so the client's `last_valid_recv` keeps refreshing and /// its stall detector never fires. The RTT readout freezes at its last value /// for the same reason — Pong rides in a Data frame stuck behind the gap. /// /// So: once we have been stuck long enough that retransmission has provably /// given up, skip to the lowest buffered nonce and drain. This drops the /// missing frame's payload (one RelayMessage — a chunk of one stream), which /// is a real cost, but the alternative is a permanently dead tunnel. fn recover_stalled_gap(&mut self) -> Vec { let mut recovered = Vec::new(); if self.reorder_buffer.is_empty() { return recovered; } // Wait out the sender's full retransmit budget before giving up, so a // frame that is merely late is never discarded. The sender backs off // exponentially, so key this off the live RTO estimate rather than a // flat constant, with a floor that keeps low-RTT links from skipping // too eagerly and a ceiling that bounds the visible freeze. let timeout = self .cc .rto() .saturating_mul(GAP_RECOVERY_RTO_MULTIPLIER) .clamp(GAP_RECOVERY_MIN, GAP_RECOVERY_MAX); if self.last_recv_advance.elapsed() < timeout { return recovered; } let Some(&resume_at) = self.reorder_buffer.keys().next() else { return recovered; }; let skipped = resume_at.saturating_sub(self.expected_recv_nonce); tracing::warn!( "Gap recovery: no progress for {:?}; skipping {} unrecoverable frame(s) \ (nonce {} -> {}) to unblock the session", self.last_recv_advance.elapsed(), skipped, self.expected_recv_nonce, resume_at ); self.expected_recv_nonce = resume_at; while let Some(buffered) = self.reorder_buffer.remove(&self.expected_recv_nonce) { recovered.push(buffered); match self.expected_recv_nonce.checked_add(1) { Some(next) => self.expected_recv_nonce = next, // u64 nonce space exhausted: stop draining rather than wrap. // The session is finished either way; the caller's next decrypt // will fail and tear it down. None => break, } } self.last_recv_advance = Instant::now(); // The peer must learn the sequence moved on, or it will keep // retransmitting into the void. self.ack_pending = true; recovered } fn handle_data_inbound(&mut self, raw_vec: &[u8]) -> Result { // Check for a stalled gap before classifying this frame, so the rest of // the function sees an already-advanced `expected_recv_nonce`. Runs here // rather than on Tick because both tick handlers discard DeliverApp // actions, and because inbound frames keep arriving throughout the stall // (retransmits/ACKs/NACKs/keepalives) — so this path is reliably reached. let recovered = self.recover_stalled_gap(); let result = self.handle_data_inbound_frame(raw_vec)?; if recovered.is_empty() { return Ok(result); } // Recovered payloads are older than anything this frame produces, so // they go first to preserve delivery order. let mut all = recovered; match result { ProtocolAction::Noop => {} ProtocolAction::Multiple(list) => all.extend(list), single => all.push(single), } Ok(if all.len() == 1 { all.pop().unwrap() } else { ProtocolAction::Multiple(all) }) } fn handle_data_inbound_frame(&mut self, raw_vec: &[u8]) -> Result { if raw_vec.len() < 12 { return Err(ProtocolError::Framing("data datagram too short".to_string())); } let nonce = u64::from_be_bytes(raw_vec[4..12].try_into().map_err(|_| ProtocolError::Framing("data datagram too short for nonce".into()))?); if nonce < self.expected_recv_nonce { // Duplicate — the ACK we sent was likely lost or delayed. tracing::debug!("Duplicate frame nonce={} (expected {}), forcing ACK", nonce, self.expected_recv_nonce); if let Some(ack_frame) = self.force_build_ack()? { return Ok(ProtocolAction::SendDatagram(ack_frame)); } return Ok(ProtocolAction::Noop); } if nonce > self.expected_recv_nonce + self.max_reorder { tracing::debug!("Frame nonce={} exceeds max reorder window (expected={}, max_gap={}), sending NACK", nonce, self.expected_recv_nonce, self.max_reorder ); if let Ok(nack_frame) = self.build_control_datagram( 0, FrameKind::Nack, Bytes::copy_from_slice(&self.expected_recv_nonce.to_be_bytes()), ) { return Ok(ProtocolAction::SendDatagram(nack_frame)); } return Ok(ProtocolAction::Noop); } let ciphertext = &raw_vec[12..]; let cipher = self.recv_cipher.as_ref().ok_or_else(|| { ProtocolError::State("missing recv cipher".to_string()) })?; let session_id_bytes = self.session_id.to_be_bytes(); let plaintext = cipher.decrypt(nonce, ciphertext, &session_id_bytes)?; let packet = FramedPacket::decode_zero_copy(Bytes::from(plaintext))?; let mut outbound_actions = Vec::new(); // Fast path processing for Nacks: act immediately, bypass sequence queue if packet.header.kind == FrameKind::Nack && packet.payload.len() >= 8 { let req_nonce = u64::from_be_bytes(packet.payload[..8].try_into().map_err(|_| ProtocolError::Framing("nack payload too short".into()))?); if let Some(cached_frame) = self.lookup_sent_frame(req_nonce) { tracing::debug!("NACK received: retransmitting nonce={}", req_nonce); self.cc.on_loss(cached_frame.len() as u64); outbound_actions.push(ProtocolAction::SendDatagram(cached_frame)); } else { tracing::debug!("NACK received: nonce={} not found in sent_history (evicted)", req_nonce); // Estimate ~1200 bytes lost for evicted frames self.cc.on_loss(1200); } } if packet.header.kind == FrameKind::Ack { let ranges = parse_ack_ranges(&packet.payload)?; self.drop_acked_frames(&ranges); } let action = match packet.header.kind { FrameKind::Data => { ProtocolAction::DeliverApp(packet.header.stream_id, packet.payload) } FrameKind::Close => { tracing::debug!("Received Close frame, terminating session"); self.state = OstpState::Closed; ProtocolAction::Noop } FrameKind::KeepAlive => ProtocolAction::Noop, _ => ProtocolAction::Noop, }; let mut app_actions = Vec::new(); if matches!(packet.header.kind, FrameKind::Data | FrameKind::Close | FrameKind::KeepAlive) { self.ack_pending = true; } if nonce == self.expected_recv_nonce { app_actions.push(action); self.expected_recv_nonce = self.expected_recv_nonce.checked_add(1).ok_or_else(|| { ProtocolError::Crypto("recv nonce sequence exhausted".to_string()) })?; self.last_recv_advance = Instant::now(); // Drain continuous queue while let Some(buffered_action) = self.reorder_buffer.remove(&self.expected_recv_nonce) { app_actions.push(buffered_action); self.expected_recv_nonce = self.expected_recv_nonce.checked_add(1).ok_or_else(|| { ProtocolError::Crypto("recv nonce sequence exhausted".to_string()) })?; } self.last_recv_advance = Instant::now(); } else { // Gap detected if nonce >= self.expected_recv_nonce { if self.reorder_buffer.len() < self.max_reorder_buffer { self.reorder_buffer.insert(nonce, action); } else { tracing::warn!("Reorder buffer still full after gap recovery, dropping frame nonce={}", nonce); } } else { tracing::debug!("Frame nonce={} arrived too late after gap recovery, dropping", nonce); } // Rate-limited NACK: send at most once per (rto/2) to prevent retransmit storms. // Using rto/2 means we send a NACK before the sender's timer fires, prompting // fast retransmit without flooding. Floor at 10ms to handle very low-RTT links. let nack_cooldown = (self.cc.rto() / 2).max(Duration::from_millis(10)); if self.last_nack_sent.elapsed() >= nack_cooldown { self.last_nack_sent = Instant::now(); let nack_payload = self.expected_recv_nonce.to_be_bytes(); if let Ok(nack_frame) = self.build_control_datagram(0, FrameKind::Nack, Bytes::copy_from_slice(&nack_payload)) { outbound_actions.push(ProtocolAction::SendDatagram(nack_frame)); } } } if let Some(ack_frame) = self.build_ack_if_due()? { outbound_actions.push(ProtocolAction::SendDatagram(ack_frame)); } // Collate both types of output (application payloads and wire actions like Nacks/Retransmissions) let mut all_actions = Vec::new(); all_actions.extend(outbound_actions); all_actions.extend(app_actions); if all_actions.is_empty() { Ok(ProtocolAction::Noop) } else if all_actions.len() == 1 { Ok(all_actions.pop().unwrap()) } else { Ok(ProtocolAction::Multiple(all_actions)) } } fn wrap_datagram_handshake(&self, noise_payload: &[u8]) -> Result { // Anti-DPI: add random padding after the Noise payload to prevent // size fingerprinting. The padding range is derived from the access key // (Kerckhoffs's principle), so different keys produce different size // distributions — no universal filter can be built from the binary alone. // // Wire format: [session_id:4][noise_len:2][noise_payload:N][random_padding] let pad_len: usize = rand::thread_rng().gen_range(self.handshake_pad_min..=self.handshake_pad_max); let mut pad = vec![0u8; pad_len]; rand::thread_rng().fill(&mut pad[..]); let noise_len = noise_payload.len() as u16; let mut out = Vec::with_capacity(4 + 2 + noise_payload.len() + pad_len); out.extend_from_slice(&self.session_id.to_be_bytes()); out.extend_from_slice(&noise_len.to_be_bytes()); out.extend_from_slice(noise_payload); out.extend_from_slice(&pad); crate::crypto::obfuscate_packet_inplace(&mut out, &self.obfuscation_key, true); Ok(Bytes::from(out)) } fn build_tracked_datagram(&mut self, stream_id: u16, kind: FrameKind, payload: Bytes) -> Result { self.build_datagram(stream_id, kind, payload, true) } fn build_control_datagram(&mut self, stream_id: u16, kind: FrameKind, payload: Bytes) -> Result { self.build_datagram(stream_id, kind, payload, false) } fn build_datagram(&mut self, stream_id: u16, kind: FrameKind, payload: Bytes, is_retransmittable: bool) -> Result { let padding = self.padder.build_padding(payload.len()); let header = FrameHeader { version: 1, kind, stream_id, payload_len: payload.len() as u32, pad_len: padding.len() as u16, }; let packet = FramedPacket { header, payload, padding: Bytes::from(padding), }; let plaintext = packet.encode(); let cipher = self.send_cipher.as_ref().ok_or_else(|| { ProtocolError::State("missing send cipher".to_string()) })?; let nonce = self.send_nonce; self.send_nonce = self.send_nonce.checked_add(1).ok_or_else(|| { ProtocolError::Crypto("send nonce sequence exhausted".to_string()) })?; let session_id_bytes = self.session_id.to_be_bytes(); let ciphertext = cipher.encrypt(nonce, plaintext.as_ref(), &session_id_bytes)?; let mut out = Vec::with_capacity(4 + 8 + ciphertext.len()); out.extend_from_slice(&session_id_bytes); out.extend_from_slice(&nonce.to_be_bytes()); out.extend_from_slice(&ciphertext); crate::crypto::obfuscate_packet_inplace(&mut out, &self.obfuscation_key, false); let final_bytes = Bytes::from(out); self.push_sent_frame(nonce, final_bytes.clone(), is_retransmittable); Ok(final_bytes) } pub fn set_session_keys(&mut self, session_id: u32, obfuscation_key: [u8; 8]) { self.session_id = session_id; self.obfuscation_key = obfuscation_key; } fn handle_tick(&mut self) -> Result { let mut actions = Vec::new(); // ── Pending ACK flush ───────────────────────────────────────── if let Some(ack_frame) = self.build_ack_if_due()? { actions.push(ProtocolAction::SendDatagram(ack_frame)); } let now = Instant::now(); // Use the adaptive RTO from the congestion controller (RFC 6298 SRTT + 4*RTTVAR). // Falls back to rto_initial before the first ACK is received. let base_rto_ms = self.cc.rto().max(self.rto).as_millis().max(1) as u64; // ── Zombie frame eviction ──────────────────────────────────── // Evict frames that exceeded max_retries + 2 grace retries. let grace = self.max_retries.saturating_add(2); let before = self.sent_history.len(); self.sent_history.retain(|f| !f.is_retransmittable || f.retries <= grace); let evicted = before - self.sent_history.len(); if evicted > 0 { tracing::debug!("Evicted {} zombie frames from sent_history (remaining={})", evicted, self.sent_history.len()); } // ── Retransmit expired frames ──────────────────────────────── // Limit retransmits per tick to prevent bandwidth saturation // Backoff starts from retry #0 (immediately effective): // effective_rto = base_rto * 2^retries, capped at 2^6 = 64× let mut retransmit_budget: usize = self.cc.retransmit_budget(); for frame in self.sent_history.iter_mut() { if !frame.is_retransmittable { continue; } // Out of budget for this tick — stop scanning rather than walking the // rest of the queue. sent_history is in send order, so everything we // skip is strictly newer than what we already handled; deferring it to // the next tick preserves oldest-first retransmit priority. if retransmit_budget == 0 { break; } // Exponential backoff, but bounded in absolute terms. base_rto is // itself adaptive and can reach RTO_MAX (16s) on a congested path; // multiplying that by the 64x backoff cap yields a frame that sits // unretransmitted for ~17 MINUTES, long past the point where the // session is simply dead to the user. Cap the product so backoff // stays a backoff rather than an outage. let backoff_factor = 1u64 << (frame.retries as u64).min(6); let effective_rto = Duration::from_millis(base_rto_ms.saturating_mul(backoff_factor)) .min(MAX_EFFECTIVE_RTO); if now.duration_since(frame.last_sent) >= effective_rto { // Only burn the retry counter and reset the RTO timer when the // frame is ACTUALLY put on the wire. Doing it unconditionally // meant that whenever the per-tick budget ran out — which is // exactly when loss is heavy and retransmits matter most — // frames accumulated "phantom retries" they never actually got, // and the zombie eviction above then silently dropped them after // `grace` such rounds. The peer never received that data and // never would: that stream stalls forever while the session // itself stays healthy, which is precisely the reported "tunnel // frozen at 0 b/s but the session still up" symptom. frame.last_sent = now; frame.retries = frame.retries.saturating_add(1); actions.push(ProtocolAction::SendDatagram(frame.bytes.clone())); retransmit_budget -= 1; } } if actions.is_empty() { Ok(ProtocolAction::Noop) } else if actions.len() == 1 { Ok(actions.pop().unwrap()) } else { Ok(ProtocolAction::Multiple(actions)) } } fn build_ack_if_due(&mut self) -> Result, ProtocolError> { if !self.ack_pending { return Ok(None); } let now = Instant::now(); if now.duration_since(self.last_ack_sent) < self.ack_delay { return Ok(None); } let payload = self.build_ack_payload(); if payload.is_empty() { self.ack_pending = false; return Ok(None); } let frame = self.build_control_datagram(0, FrameKind::Ack, payload)?; self.ack_pending = false; self.last_ack_sent = now; Ok(Some(frame)) } fn force_build_ack(&mut self) -> Result, ProtocolError> { let payload = self.build_ack_payload(); if payload.is_empty() { self.ack_pending = false; return Ok(None); } let frame = self.build_control_datagram(0, FrameKind::Ack, payload)?; self.ack_pending = false; self.last_ack_sent = Instant::now(); Ok(Some(frame)) } fn build_ack_payload(&self) -> Bytes { const MAX_RANGES: usize = 8; let mut ranges = Vec::new(); if self.expected_recv_nonce > 0 { ranges.push((0_u64, self.expected_recv_nonce - 1)); } let mut current_start: Option = None; let mut last = 0_u64; for &nonce in self.reorder_buffer.keys() { if current_start.is_none() { current_start = Some(nonce); last = nonce; } else if nonce == last + 1 { last = nonce; } else { ranges.push((current_start.unwrap(), last)); current_start = Some(nonce); last = nonce; } } if let Some(start) = current_start { ranges.push((start, last)); } if ranges.is_empty() { return Bytes::new(); } if ranges.len() > MAX_RANGES { // Always preserve the cumulative range (index 0) so the sender knows // all frames up to expected_recv_nonce are received. Truncate SACK ranges. let mut trimmed = vec![ranges[0]]; let tail_start = ranges.len().saturating_sub(MAX_RANGES - 1); trimmed.extend_from_slice(&ranges[tail_start..]); ranges = trimmed; } let mut out = Vec::with_capacity(1 + ranges.len() * 16); out.push(ranges.len() as u8); for (start, end) in ranges { out.extend_from_slice(&start.to_be_bytes()); out.extend_from_slice(&end.to_be_bytes()); } Bytes::from(out) } fn lookup_sent_frame(&mut self, nonce: u64) -> Option { if let Some(frame) = self.sent_history.iter_mut().rev().find(|f| f.nonce == nonce) { frame.last_sent = Instant::now(); frame.retries = frame.retries.saturating_add(1); return Some(frame.bytes.clone()); } None } fn push_sent_frame(&mut self, nonce: u64, bytes: Bytes, is_retransmittable: bool) { if is_retransmittable { self.cc.on_send(bytes.len() as u64); } self.sent_history.push_back(SentFrame { nonce, bytes, last_sent: Instant::now(), retries: 0, is_retransmittable, }); if self.sent_history.len() > self.max_sent_history { let overflow = self.sent_history.len() - self.max_sent_history; tracing::debug!("sent_history overflow: evicting {} oldest frames (cap={})", overflow, self.max_sent_history ); while self.sent_history.len() > self.max_sent_history { self.sent_history.pop_front(); } } } fn drop_acked_frames(&mut self, ranges: &[(u64, u64)]) { let now = Instant::now(); let mut acked_bytes = 0u64; let mut min_rtt: Option = None; for frame in self.sent_history.iter() { if nonce_in_ranges(frame.nonce, ranges) { acked_bytes += frame.bytes.len() as u64; // Karn's algorithm: never take an RTT sample from a frame that // was retransmitted. `last_sent` is bumped on every retransmit, // so an ACK for the ORIGINAL transmission would be measured // against the retransmit time, yielding a spuriously small RTT // that drags SRTT/RTO down and triggers more spurious // retransmits. Only unambiguous (never-retried) frames qualify. if frame.retries == 0 { let rtt = now.duration_since(frame.last_sent); min_rtt = Some(min_rtt.map_or(rtt, |m| m.min(rtt))); } } } self.sent_history.retain(|frame| !nonce_in_ranges(frame.nonce, ranges)); // Notify congestion controller. Feed an RTT sample only when we had at // least one unambiguous ACK; otherwise update the window without // polluting the RTT estimator. if acked_bytes > 0 { match min_rtt { Some(rtt) => self.cc.on_ack(acked_bytes, rtt), None => self.cc.on_ack_no_rtt(acked_bytes), } } } } fn parse_ack_ranges(payload: &[u8]) -> Result, ProtocolError> { if payload.is_empty() { return Ok(Vec::new()); } let count = payload[0] as usize; let expected = 1 + count * 16; if payload.len() < expected { return Err(ProtocolError::Framing("ack payload truncated".to_string())); } let mut ranges = Vec::with_capacity(count); let mut idx = 1; for _ in 0..count { let start = u64::from_be_bytes(payload[idx..idx + 8].try_into().map_err(|_| ProtocolError::Framing("ack range start invalid".into()))?); let end = u64::from_be_bytes(payload[idx + 8..idx + 16].try_into().map_err(|_| ProtocolError::Framing("ack range end invalid".into()))?); ranges.push((start, end)); idx += 16; } Ok(ranges) } fn nonce_in_ranges(nonce: u64, ranges: &[(u64, u64)]) -> bool { ranges.iter().any(|(start, end)| nonce >= *start && nonce <= *end) } #[cfg(test)] mod tests { use super::*; use crate::framing::PaddingStrategy; fn test_psk() -> [u8; 32] { let mut psk = [0u8; 32]; psk[0] = 0xAB; psk[15] = 0xCD; psk[31] = 0xEF; psk } fn make_config(role: NoiseRole) -> ProtocolConfig { ProtocolConfig { role, psk: test_psk(), session_id: 1, handshake_payload: vec![], max_padding: 64, padding_strategy: PaddingStrategy::Adaptive, obfuscation_key: [0u8; 8], max_reorder: 128, max_reorder_buffer: 256, ack_delay_ms: 5, rto_ms: 100, max_retries: 4, max_sent_history: 1024, handshake_pad_min: 8, handshake_pad_max: 32, mtu: 1400, } } /// Full handshake: Initiator -> Responder -> Initiator -> Established fn do_handshake() -> (ProtocolMachine, ProtocolMachine) { let mut client = ProtocolMachine::new(make_config(NoiseRole::Initiator)).unwrap(); let mut server = ProtocolMachine::new(make_config(NoiseRole::Responder)).unwrap(); // Client sends handshake message 1 let action = client.on_event(OstpEvent::Start).unwrap(); let msg1 = match action { ProtocolAction::SendDatagram(d) => d, _ => panic!("expected SendDatagram from client Start"), }; assert_eq!(client.state(), OstpState::Handshaking); // Server receives msg1 and responds let action = server.on_event(OstpEvent::Start).unwrap(); assert!(matches!(action, ProtocolAction::Noop)); let action = server.on_event(OstpEvent::Inbound(msg1)).unwrap(); let msg2 = match action { ProtocolAction::Multiple(actions) => { actions.into_iter().find_map(|a| match a { ProtocolAction::SendDatagram(d) => Some(d), _ => None, }).expect("server should send datagram in handshake response") } ProtocolAction::SendDatagram(d) => d, ProtocolAction::HandshakePayload(_, Some(d)) => d, other => panic!("unexpected server response: {:?}", std::mem::discriminant(&other)), }; // Client receives msg2 -> Established let action = client.on_event(OstpEvent::Inbound(msg2)).unwrap(); match action { ProtocolAction::HandshakePayload(_, _) => {} ProtocolAction::Multiple(_) => {} _ => {} } // Both should be Established assert_eq!(client.state(), OstpState::Established); assert_eq!(server.state(), OstpState::Established); (client, server) } #[test] fn test_full_handshake() { let (client, server) = do_handshake(); assert_eq!(client.state(), OstpState::Established); assert_eq!(server.state(), OstpState::Established); } #[test] fn test_data_exchange_client_to_server() { let (mut client, mut server) = do_handshake(); // Client sends data let payload = Bytes::from_static(b"hello from client"); let action = client.on_event(OstpEvent::Outbound(1, payload.clone())).unwrap(); let datagram = match action { ProtocolAction::SendDatagram(d) => d, _ => panic!("expected SendDatagram"), }; // Server receives and decrypts let action = server.on_event(OstpEvent::Inbound(datagram)).unwrap(); match action { ProtocolAction::DeliverApp(stream_id, data) => { assert_eq!(stream_id, 1); assert_eq!(data.as_ref(), b"hello from client"); } ProtocolAction::Multiple(actions) => { let found = actions.iter().any(|a| matches!(a, ProtocolAction::DeliverApp(1, d) if d.as_ref() == b"hello from client" )); assert!(found, "expected DeliverApp in Multiple"); } _ => panic!("expected DeliverApp or Multiple"), } } #[test] fn test_data_exchange_server_to_client() { let (mut client, mut server) = do_handshake(); // Server sends data let payload = Bytes::from_static(b"hello from server"); let action = server.on_event(OstpEvent::Outbound(2, payload.clone())).unwrap(); let datagram = match action { ProtocolAction::SendDatagram(d) => d, _ => panic!("expected SendDatagram"), }; // Client receives let action = client.on_event(OstpEvent::Inbound(datagram)).unwrap(); match action { ProtocolAction::DeliverApp(stream_id, data) => { assert_eq!(stream_id, 2); assert_eq!(data.as_ref(), b"hello from server"); } ProtocolAction::Multiple(actions) => { let found = actions.iter().any(|a| matches!(a, ProtocolAction::DeliverApp(2, d) if d.as_ref() == b"hello from server" )); assert!(found, "expected DeliverApp in Multiple"); } _ => panic!("expected DeliverApp or Multiple"), } } #[test] fn test_close_sequence() { let (mut client, mut server) = do_handshake(); // Client sends Close let action = client.on_event(OstpEvent::Close).unwrap(); let close_datagram = match action { ProtocolAction::SendDatagram(d) => d, _ => panic!("expected SendDatagram for Close"), }; assert_eq!(client.state(), OstpState::Closing); // Server receives Close let _action = server.on_event(OstpEvent::Inbound(close_datagram)).unwrap(); assert_eq!(server.state(), OstpState::Closed); } #[test] fn test_wrong_psk_handshake_fails() { let mut client = ProtocolMachine::new(make_config(NoiseRole::Initiator)).unwrap(); let mut bad_psk_config = make_config(NoiseRole::Responder); bad_psk_config.psk = [0xFF; 32]; // Different PSK let mut server = ProtocolMachine::new(bad_psk_config).unwrap(); let action = client.on_event(OstpEvent::Start).unwrap(); let msg1 = match action { ProtocolAction::SendDatagram(d) => d, _ => panic!("expected SendDatagram"), }; let _ = server.on_event(OstpEvent::Start).unwrap(); // Server should fail to process handshake with wrong PSK let result = server.on_event(OstpEvent::Inbound(msg1)); // Either an error or the server stays in Handshaking (never reaches Established) assert!(result.is_err() || server.state() != OstpState::Established); } #[test] fn test_congestion_controller_after_handshake() { let (client, _server) = do_handshake(); // CC should be in SlowStart after handshake let budget = client.cc.retransmit_budget(); assert!(budget >= 2, "initial retransmit budget should be >= 2, got {}", budget); } #[test] fn test_multiple_data_frames() { let (mut client, mut server) = do_handshake(); // Send 10 frames for i in 0..10u8 { let payload = Bytes::from(vec![i; 100]); let action = client.on_event(OstpEvent::Outbound(1, payload)).unwrap(); let datagram = match action { ProtocolAction::SendDatagram(d) => d, _ => panic!("expected SendDatagram for frame {}", i), }; let action = server.on_event(OstpEvent::Inbound(datagram)).unwrap(); match action { ProtocolAction::DeliverApp(_, data) => { assert_eq!(data.len(), 100); assert_eq!(data[0], i); } ProtocolAction::Multiple(actions) => { let found = actions.iter().any(|a| matches!(a, ProtocolAction::DeliverApp(_, d) if d.len() == 100 && d[0] == i )); assert!(found, "frame {} not found in Multiple", i); } _ => panic!("unexpected action for frame {}", i), } } // Verify in-flight state assert!(client.in_flight_count() > 0, "should have in-flight frames"); } #[test] fn test_tick_no_crash() { let (mut client, mut server) = do_handshake(); // Tick should not crash on either side let _ = client.on_event(OstpEvent::Tick).unwrap(); let _ = server.on_event(OstpEvent::Tick).unwrap(); } /// A retry may only be charged to a frame that was actually retransmitted. /// /// The retransmit loop is budget-limited per tick. It used to bump /// `retries` and reset `last_sent` for every due frame regardless of /// whether the budget allowed it to actually send — so under heavy loss /// (exactly when the budget runs out) frames racked up retries they never /// received, and the zombie eviction dropped them after `max_retries + 2` /// such rounds. That data was never delivered and never would be: the /// stream stalls permanently while the session itself stays up. #[test] fn test_retransmit_budget_charges_retries_only_for_frames_actually_sent() { let (mut client, _server) = do_handshake(); // Queue far more in-flight frames than a single tick's budget allows. const FRAMES: usize = 40; for i in 0..FRAMES { let payload = Bytes::from(vec![i as u8; 200]); client.on_event(OstpEvent::Outbound(1, payload)).unwrap(); } assert_eq!(client.in_flight_count(), FRAMES); assert_eq!(client.total_retries(), 0, "nothing retransmitted yet"); // Let every frame's RTO lapse so that on the next tick all FRAMES frames // are due at once and the per-tick budget is guaranteed to run out. The // effective RTO here is max(cc.rto(), config rto_ms) = 100ms at retries=0. std::thread::sleep(Duration::from_millis(150)); let sent = count_datagrams(&client.on_event(OstpEvent::Tick).unwrap()); assert!(sent > 0, "expected some retransmits after the RTO lapsed"); assert!( sent < FRAMES, "budget should have capped this tick below the {FRAMES} due frames, got {sent}" ); assert_eq!( client.total_retries(), sent, "charged {} retries but only put {} frames on the wire — the \ difference is phantom retries that will silently evict live data", client.total_retries(), sent ); assert_eq!( client.in_flight_count(), FRAMES, "nothing was acked, so no frame may be evicted yet" ); } /// Count how many datagrams an action tree actually puts on the wire. fn count_datagrams(action: &ProtocolAction) -> usize { match action { ProtocolAction::SendDatagram(_) => 1, ProtocolAction::Multiple(list) => list.iter().map(count_datagrams).sum(), _ => 0, } } /// Count how many application payloads an action tree actually delivers. fn delivered_payloads(action: &ProtocolAction) -> Vec { match action { ProtocolAction::DeliverApp(_, data) => vec![data.clone()], ProtocolAction::Multiple(list) => list.iter().flat_map(delivered_payloads).collect(), _ => Vec::new(), } } /// Build `count` data frames on `client`, returning them without delivering /// any — lets a test choose which ones to "lose" in transit. fn make_data_frames(client: &mut ProtocolMachine, count: u8) -> Vec { (0..count) .map(|i| { let payload = Bytes::from(vec![i; 32]); match client.on_event(OstpEvent::Outbound(1, payload)).unwrap() { ProtocolAction::SendDatagram(d) => d, _ => panic!("expected SendDatagram for frame {i}"), } }) .collect() } /// The freeze this fixes: a frame is lost, the sender eventually stops /// retransmitting it, and the receiver — which gates delivery on /// `expected_recv_nonce` — waits for it forever. Every later frame piles up /// undelivered while the transport itself stays healthy, so nothing upstream /// notices. Recovery must eventually skip the hole and release the backlog. #[test] fn test_gap_recovery_releases_permanently_stalled_frames() { let (mut client, mut server) = do_handshake(); let frames = make_data_frames(&mut client, 4); // Frame 0 arrives in order and is delivered straight through. let action = server.on_event(OstpEvent::Inbound(frames[0].clone())).unwrap(); assert_eq!(delivered_payloads(&action).len(), 1, "in-order frame should deliver"); // Frame 1 is lost. 2 and 3 arrive but must be held back — delivering them // now would reorder the stream. for idx in [2usize, 3] { let action = server.on_event(OstpEvent::Inbound(frames[idx].clone())).unwrap(); assert!( delivered_payloads(&action).is_empty(), "frame {idx} must stay buffered behind the missing frame" ); } // Stand in for "the sender exhausted its retries and dropped frame 1": // the sequence has not advanced for longer than the recovery timeout. server.last_recv_advance = Instant::now() - GAP_RECOVERY_MAX - Duration::from_secs(1); // The next inbound frame (a retransmitted duplicate, which is exactly what // a real stalled session keeps receiving) must unblock the backlog. let action = server.on_event(OstpEvent::Inbound(frames[0].clone())).unwrap(); let delivered = delivered_payloads(&action); assert_eq!( delivered.len(), 2, "both buffered frames must be released once the gap is declared unrecoverable" ); // ...and in order: frame 2 before frame 3. assert_eq!(delivered[0][0], 2); assert_eq!(delivered[1][0], 3); } /// Recovery must not be trigger-happy: a frame that is merely late still has /// to be waited for, or we would discard data the sender is about to resend. #[test] fn test_gap_recovery_does_not_fire_before_timeout() { let (mut client, mut server) = do_handshake(); let frames = make_data_frames(&mut client, 3); server.on_event(OstpEvent::Inbound(frames[0].clone())).unwrap(); let action = server.on_event(OstpEvent::Inbound(frames[2].clone())).unwrap(); assert!(delivered_payloads(&action).is_empty()); // Well inside the timeout — the gap must still be respected. let action = server.on_event(OstpEvent::Inbound(frames[0].clone())).unwrap(); assert!( delivered_payloads(&action).is_empty(), "must keep waiting while retransmission is still plausible" ); // And once the genuinely-late frame shows up, normal in-order delivery // resumes with nothing dropped. let action = server.on_event(OstpEvent::Inbound(frames[1].clone())).unwrap(); let delivered = delivered_payloads(&action); assert_eq!(delivered.len(), 2, "late frame plus the buffered one"); assert_eq!(delivered[0][0], 1); assert_eq!(delivered[1][0], 2); } }