From f9c1973c1364e1414431ecf0754c8ddf96b4ade6 Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Sun, 14 Jun 2026 10:45:44 -0400 Subject: [PATCH] Harden native protocol: version discipline, rekey, migration, rate limit, O(1) client index Track A protocol hardening for native-v1. 1. Protocol VERSION discipline (protocol.rs, native.rs, dosh-server.rs, dosh-client.rs): bump protocol::VERSION to 2 and NATIVE_PROTOCOL_VERSION to 2. Foreign wire-version datagrams now get a clear named "protocol version mismatch - upgrade dosh" reject from the server instead of a silent timeout (peek_foreign_wire_version + VERSION_MISMATCH_REASON). The native handshake's plaintext protocol_version is also checked server-side before any crypto via check_native_protocol_version, returning a typed ProtocolVersionMismatch. 2. Transport rekey (spec section 11/9): server rotates a client's traffic key after rekey_after_packets OR rekey_after_secs (config knobs). Rotated keys are derived independently of the handshake keys from the current key plus fresh server CSPRNG material shipped confidentially in an AEAD Rekey packet (derive_rekey_session_key). The previous epoch's key is retained briefly so in-flight pre-rekey packets still decrypt (matched by session_key_id), and stale-epoch packets are ignored, not fatal. Client handles Rekey/RekeyAck. 3. Connection migration (spec section 11): every authenticated, replay-accepted packet now migrates client.endpoint to the new source address (input, resize, ping, ack, resume, stream*, rekey-ack), not just resume. Ping now verifies its AEAD tag before acting so migration cannot be spoofed. 4. Native auth rate limiting: per-source-IP token bucket (native_auth_rate_limit_per_minute) enforced in handle_native_client_hello BEFORE any X25519/Ed25519 work; over-limit hellos get a non-crypto reject. 5. Speed: replaced O(sessions x clients) per-packet linear scans with an O(1) HashMap index in ServerState, kept in sync on every client insert/remove (attach handlers, detach, remove_client, cleanup reap, session-exit drain). New config keys (ServerConfig, with defaults): rekey_after_packets = 100000, rekey_after_secs = 3600. Tests: version-mismatch reject (no hang), rate-limit flood reject, rekey round-trip end-to-end + key-derivation unit tests, source-address migration, client-index sync + cleanup purge. fmt and full test suite green. Co-Authored-By: Claude Opus 4.8 --- src/bin/dosh-client.rs | 58 ++- src/bin/dosh-server.rs | 792 +++++++++++++++++++++++++++++-------- src/config.rs | 22 ++ src/native.rs | 90 ++++- src/protocol.rs | 46 ++- tests/integration_smoke.rs | 311 +++++++++++++++ tests/protocol_auth.rs | 133 +++++++ 7 files changed, 1283 insertions(+), 169 deletions(-) diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index e247746..770cac8 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -17,7 +17,7 @@ use dosh::native::{ use dosh::protocol::{ self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input, NativeAuthCheckOkBody, NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody, - NativeUserAuthBody, PacketKind, Resize, ResumeRequest, SERVER_TO_CLIENT, StreamClose, + NativeUserAuthBody, PacketKind, Rekey, Resize, ResumeRequest, SERVER_TO_CLIENT, StreamClose, StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust, TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope, }; @@ -2145,6 +2145,9 @@ async fn run_terminal( None }; + // Retain the previous epoch's key briefly after a rekey so any in-flight + // pre-rekey frame still decrypts instead of triggering a needless reconnect. + let mut previous_session_key: Option<[u8; 32]> = None; let mut recv_buf = vec![0u8; 65535]; loop { tokio::select! { @@ -2185,7 +2188,16 @@ async fn run_terminal( } match packet.header.kind { PacketKind::Frame | PacketKind::ResumeOk => { - let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else { + let decrypted = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) + .or_else(|err| { + // Fall back to the previous epoch key for in-flight + // pre-rekey frames before declaring the link stale. + match previous_session_key { + Some(prev) => protocol::decrypt_body(&packet, &prev, SERVER_TO_CLIENT), + None => Err(err), + } + }); + let Ok(plain) = decrypted else { if let Some(frame) = reconnect( &socket, &mut cred, @@ -2225,6 +2237,48 @@ async fn run_terminal( PacketKind::Pong => { last_packet_at = Instant::now(); } + PacketKind::Rekey => { + // Server-initiated transport rekey (spec §11). The Rekey is + // sealed under the current key; once decrypted we derive the + // next key from the shipped fresh material + current key, + // switch atomically (keeping the old key for grace), and + // confirm with a RekeyAck under the NEW key. + let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else { + continue; + }; + let Ok(rekey) = protocol::from_body::(&plain) else { + continue; + }; + let previous_key = cred.session_key; + let previous_key_id = cred.session_key_id; + let Ok(new_key) = dosh::native::derive_rekey_session_key( + &previous_key, + &rekey.rekey_material, + &previous_key_id, + rekey.epoch, + ) else { + continue; + }; + // Only accept if our derivation matches the server's id. + if protocol::session_key_id(&new_key) != rekey.new_session_key_id { + continue; + } + previous_session_key = Some(previous_key); + cred.session_key = new_key; + cred.session_key_id = rekey.new_session_key_id; + last_packet_at = Instant::now(); + send_seq += 1; + let ack = protocol::encode_encrypted( + PacketKind::RekeyAck, + cred.client_id, + send_seq, + cred.last_rendered_seq, + &cred.session_key, + CLIENT_TO_SERVER, + b"", + )?; + socket.send_to(&ack, addr).await?; + } PacketKind::AttachReject => { let reject: AttachReject = protocol::from_body(&packet.body)?; if reject.reason == "unknown client" { diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index 4fffdc2..bb185e7 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -20,7 +20,7 @@ use dosh::protocol::{ }; use dosh::pty::{PtyHandle, PtyOutput, spawn_pty_session}; use std::collections::{HashMap, HashSet, VecDeque}; -use std::net::SocketAddr; +use std::net::{IpAddr, SocketAddr}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -143,6 +143,9 @@ async fn serve(config_path: Option) -> Result<()> { if let Err(err) = retransmit_pending(&retransmit_state, &retransmit_socket).await { eprintln!("retransmit error: {err:#}"); } + if let Err(err) = maybe_rekey_clients(&retransmit_state, &retransmit_socket).await { + eprintln!("rekey error: {err:#}"); + } } }); @@ -171,6 +174,82 @@ struct ServerState { sessions: HashMap, pending_native: HashMap<[u8; 16], PendingNativeAuth>, next_server_stream_id: u64, + /// O(1) reverse index from a live client's `conn_id` to its session name, so + /// per-packet handlers don't linear-scan every session's client map. Kept in + /// lockstep with `Session::clients` on every insert and remove. + client_index: HashMap<[u8; 16], String>, + /// Per-source-IP token bucket for native ClientHello flood protection, + /// checked before any X25519/Ed25519 work. + native_auth_limiter: NativeAuthRateLimiter, +} + +/// Per-source-IP token bucket throttling native auth ClientHellos. +/// +/// Threat model item: a malicious unauthenticated client flooding auth attempts. +/// We charge one token per ClientHello *before* doing any X25519/Ed25519 crypto, +/// so a flood cannot burn server CPU. The bucket refills continuously at +/// `per_minute / 60` tokens per second up to a burst of `per_minute`, matching +/// the `native_auth_rate_limit_per_minute` value advertised in `ServerHello`. +struct NativeAuthRateLimiter { + per_minute: u32, + buckets: HashMap, +} + +#[derive(Clone, Copy)] +struct TokenBucket { + tokens: f64, + last_refill: Instant, +} + +impl NativeAuthRateLimiter { + fn new(per_minute: u32) -> Self { + Self { + per_minute, + buckets: HashMap::new(), + } + } + + /// Try to spend one token for `ip` at `now`. Returns `Ok(remaining)` whole + /// tokens left when allowed, or `Err(())` when the bucket is empty. A + /// `per_minute` of 0 disables native auth entirely (no tokens ever). + fn check(&mut self, ip: IpAddr, now: Instant) -> std::result::Result { + if self.per_minute == 0 { + return Err(()); + } + let capacity = self.per_minute as f64; + let refill_per_sec = capacity / 60.0; + let bucket = self.buckets.entry(ip).or_insert(TokenBucket { + tokens: capacity, + last_refill: now, + }); + let elapsed = now + .saturating_duration_since(bucket.last_refill) + .as_secs_f64(); + bucket.tokens = (bucket.tokens + elapsed * refill_per_sec).min(capacity); + bucket.last_refill = now; + if bucket.tokens < 1.0 { + return Err(()); + } + bucket.tokens -= 1.0; + Ok(bucket.tokens as u32) + } + + /// Drop buckets that have fully refilled, so the map can't grow unbounded + /// under a spoofed-source-IP flood. Called from the periodic cleanup task. + fn evict_full(&mut self, now: Instant) { + if self.per_minute == 0 { + self.buckets.clear(); + return; + } + let capacity = self.per_minute as f64; + let refill_per_sec = capacity / 60.0; + self.buckets.retain(|_, bucket| { + let elapsed = now + .saturating_duration_since(bucket.last_refill) + .as_secs_f64(); + (bucket.tokens + elapsed * refill_per_sec) < capacity + }); + } } struct Session { @@ -202,6 +281,15 @@ struct ClientState { opened_streams: HashSet, stream_send_credit: HashMap, stream_pending_data: HashMap>>, + /// Current transport key epoch (0 = original handshake key). Bumped on rekey. + epoch: u64, + /// When the current epoch began, for the wall-clock rekey trigger. + epoch_started: Instant, + /// Packets seen/sent in the current epoch, for the packet-count trigger. + epoch_packets: u64, + /// The previous epoch's key, retained briefly so in-flight pre-rekey packets + /// still decrypt (matched via `session_key_id`) instead of dropping fatally. + previous_session_key: Option<[u8; 32]>, } #[derive(Clone)] @@ -226,6 +314,7 @@ impl ServerState { secret: [u8; 32], pty_tx: mpsc::UnboundedSender, ) -> Self { + let per_minute = config.native_auth_rate_limit_per_minute; Self { config, secret, @@ -233,6 +322,8 @@ impl ServerState { sessions: HashMap::new(), pending_native: HashMap::new(), next_server_stream_id: 1u64 << 63, + client_index: HashMap::new(), + native_auth_limiter: NativeAuthRateLimiter::new(per_minute), } } @@ -283,6 +374,72 @@ impl ServerState { ); Ok(()) } + + /// Insert a client into `session_name` and record it in the O(1) index. The + /// session must already exist (callers `ensure_session` first). + fn insert_client(&mut self, session_name: &str, client_id: [u8; 16], client: ClientState) { + if let Some(session) = self.sessions.get_mut(session_name) { + session.clients.insert(client_id, client); + self.client_index + .insert(client_id, session_name.to_string()); + } + } + + /// Remove a client from whichever session holds it, keeping the index in + /// sync. Returns true when a client was actually removed. + fn remove_client_everywhere(&mut self, client_id: &[u8; 16]) -> bool { + if let Some(session_name) = self.client_index.remove(client_id) { + if let Some(session) = self.sessions.get_mut(&session_name) { + return session.clients.remove(client_id).is_some(); + } + } + false + } + + /// O(1) resolve of a live client's session key and session name via the + /// index, replacing the former O(sessions) linear scan. + fn lookup_client(&self, client_id: &[u8; 16]) -> Option<([u8; 32], String)> { + let session_name = self.client_index.get(client_id)?; + let session = self.sessions.get(session_name)?; + let client = session.clients.get(client_id)?; + Some((client.session_key, session_name.clone())) + } + + /// Resolve the decryption key for an inbound packet, honoring rekey grace. + /// + /// During the brief window after a rekey, in-flight packets sealed under the + /// previous epoch's key must still decrypt instead of being dropped as fatal + /// (spec: "stale encrypted packets ... ignored, not fatal"). We pick whichever + /// of the current or retained-previous key matches the packet's + /// `session_key_id`, falling back to the current key when neither matches (so + /// a genuinely stale/forged packet still fails the AEAD check, not silently). + fn lookup_client_key_for( + &self, + client_id: &[u8; 16], + session_key_id: &[u8; 16], + ) -> Option<([u8; 32], String)> { + let session_name = self.client_index.get(client_id)?; + let session = self.sessions.get(session_name)?; + let client = session.clients.get(client_id)?; + if protocol::session_key_id(&client.session_key) == *session_key_id { + return Some((client.session_key, session_name.clone())); + } + if let Some(previous) = client.previous_session_key + && protocol::session_key_id(&previous) == *session_key_id + { + return Some((previous, session_name.clone())); + } + Some((client.session_key, session_name.clone())) + } + + /// O(1) mutable access to a live client via the conn_id index. + fn client_mut(&mut self, client_id: &[u8; 16]) -> Option<&mut ClientState> { + let session_name = self.client_index.get(client_id)?; + self.sessions + .get_mut(session_name)? + .clients + .get_mut(client_id) + } } fn mode_uses_pty(mode: &str) -> bool { @@ -383,6 +540,14 @@ async fn handle_packet( peer: SocketAddr, raw: &[u8], ) -> Result<()> { + // A peer running a different wire protocol version cannot be decoded at all, + // because the header VERSION byte is part of the AEAD AAD and the framing. + // Rather than drop it (which the peer sees as a silent timeout), answer with + // a clear, named version-mismatch reject so the user is told to upgrade. + if let Some(foreign) = protocol::peek_foreign_wire_version(raw) { + let _ = foreign; + return send_reject(socket, peer, protocol::VERSION_MISMATCH_REASON).await; + } let packet = protocol::decode(raw)?; match packet.header.kind { PacketKind::NativeClientHello => { @@ -407,6 +572,7 @@ async fn handle_packet( | PacketKind::StreamClose | PacketKind::StreamEof | PacketKind::StreamWindowAdjust + | PacketKind::RekeyAck if find_client_key(state, &packet.header.conn_id).is_err() => { send_reject_to_client(socket, peer, packet.header.conn_id, "unknown client").await @@ -414,7 +580,8 @@ async fn handle_packet( PacketKind::Input => handle_input(state, peer, &packet).await, PacketKind::Resize => handle_resize(state, peer, &packet).await, PacketKind::Ping => handle_ping(state, socket, peer, &packet).await, - PacketKind::Ack => handle_ack(state, &packet).await, + PacketKind::Ack => handle_ack(state, peer, &packet).await, + PacketKind::RekeyAck => handle_rekey_ack(state, peer, &packet).await, PacketKind::Detach => handle_detach(state, &packet).await, PacketKind::StreamOpen => handle_stream_open(state, socket, peer, &packet).await, PacketKind::StreamOpenOk => handle_stream_open_ok(state, socket, peer, &packet).await, @@ -423,7 +590,7 @@ async fn handle_packet( PacketKind::StreamWindowAdjust => { handle_stream_window_adjust(state, socket, peer, &packet).await } - PacketKind::StreamClose => handle_stream_close(state, &packet).await, + PacketKind::StreamClose => handle_stream_close(state, peer, &packet).await, _ => Ok(()), } } @@ -435,6 +602,27 @@ async fn handle_native_client_hello( body: Vec, ) -> Result<()> { let req: NativeClientHelloBody = protocol::from_body(&body)?; + // Defense in depth against a peer that shares our wire VERSION but advertises + // a different native handshake protocol_version (the spec's negotiation + // point). Reject with a named, actionable reason before any crypto. + if let Err(err) = + dosh::native::check_native_protocol_version(req.hello.protocol_version, "client") + { + return send_reject(socket, peer, &err.to_string()).await; + } + // Rate-limit native auth *before* touching the host key or any X25519/Ed25519 + // work, so a flood of ClientHellos from one source cannot burn server CPU. + // A separate, non-crypto reject lets a legitimate client back off cleanly. + let rate_limit_remaining = { + let mut locked = state.lock().expect("server state poisoned"); + locked.native_auth_limiter.check(peer.ip(), Instant::now()) + }; + let rate_limit_remaining = match rate_limit_remaining { + Ok(remaining) => remaining, + Err(()) => { + return send_reject(socket, peer, "native auth rate limit exceeded").await; + } + }; let built: Result<([u8; 16], NativeServerHello)> = { let mut locked = state.lock().expect("server state poisoned"); if !locked.config.native_auth { @@ -468,7 +656,7 @@ async fn handle_native_client_hello( chosen_aead: "chacha20poly1305".to_string(), server_key_epoch: 1, auth_challenge: crypto::random_32(), - rate_limit_remaining: Some(locked.config.native_auth_rate_limit_per_minute), + rate_limit_remaining: Some(rate_limit_remaining), host_signature: Vec::new(), }; sign_server_hello(&host_signing, &req.hello, &mut hello)?; @@ -637,7 +825,8 @@ async fn handle_native_user_auth( let client_id = crypto::random_16(); let snapshot = screen_snapshot(session.parser.screen()); let output_seq = session.output_seq; - session.clients.insert( + locked.insert_client( + &session_name, client_id, ClientState { endpoint: peer, @@ -655,6 +844,10 @@ async fn handle_native_user_auth( opened_streams: HashSet::new(), stream_send_credit: HashMap::new(), stream_pending_data: HashMap::new(), + epoch: 0, + epoch_started: Instant::now(), + epoch_packets: 0, + previous_session_key: None, }, ); Ok(( @@ -765,7 +958,9 @@ async fn handle_bootstrap_attach( let client_id = crypto::random_16(); let snapshot = screen_snapshot(session.parser.screen()); let output_seq = session.output_seq; - session.clients.insert( + let bootstrap_session = req.bootstrap.session.clone(); + locked.insert_client( + &bootstrap_session, client_id, ClientState { endpoint: peer, @@ -783,6 +978,10 @@ async fn handle_bootstrap_attach( opened_streams: HashSet::new(), stream_send_credit: HashMap::new(), stream_pending_data: HashMap::new(), + epoch: 0, + epoch_started: Instant::now(), + epoch_packets: 0, + previous_session_key: None, }, ); ( @@ -878,7 +1077,9 @@ async fn handle_ticket_attach( let client_id = crypto::random_16(); let snapshot = screen_snapshot(session.parser.screen()); let output_seq = session.output_seq; - session.clients.insert( + let ticket_session = req.session.clone(); + locked.insert_client( + &ticket_session, client_id, ClientState { endpoint: peer, @@ -896,6 +1097,10 @@ async fn handle_ticket_attach( opened_streams: HashSet::new(), stream_send_credit: HashMap::new(), stream_pending_data: HashMap::new(), + epoch: 0, + epoch_started: Instant::now(), + epoch_packets: 0, + previous_session_key: None, }, ); (client_id, output_seq, snapshot) @@ -956,7 +1161,7 @@ async fn handle_resume( peer: SocketAddr, packet: &protocol::Packet, ) -> Result<()> { - let (key, session_name) = match find_client_key(state, &packet.header.conn_id) { + let (key, session_name) = match find_client_decrypt_key(state, &packet.header) { Ok(found) => found, Err(_) => return send_reject(socket, peer, "unknown client").await, }; @@ -1018,7 +1223,7 @@ async fn handle_input( peer: SocketAddr, packet: &protocol::Packet, ) -> Result<()> { - let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; + let (key, session_name) = find_client_decrypt_key(state, &packet.header)?; let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; let input: Input = protocol::from_body(&body)?; let mut locked = state.lock().expect("server state poisoned"); @@ -1053,7 +1258,7 @@ async fn handle_resize( peer: SocketAddr, packet: &protocol::Packet, ) -> Result<()> { - let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; + let (key, session_name) = find_client_decrypt_key(state, &packet.header)?; let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; let resize: Resize = protocol::from_body(&body)?; let mut locked = state.lock().expect("server state poisoned"); @@ -1068,8 +1273,11 @@ async fn handle_resize( if !client.replay.accept(packet.header.seq) { return Ok(()); } + // Connection migration happens for any authenticated, fresh packet (spec §11), + // independent of the session mode; a view-only client roams too. + client.endpoint = peer; + client.last_seen = Instant::now(); if mode_allows_terminal_updates(&client.mode) { - client.endpoint = peer; client.cols = resize.cols; client.rows = resize.rows; apply_terminal_size( @@ -1088,22 +1296,26 @@ async fn handle_ping( peer: SocketAddr, packet: &protocol::Packet, ) -> Result<()> { - let (key, _) = find_client_key(state, &packet.header.conn_id)?; + let (key, _) = find_client_decrypt_key(state, &packet.header)?; + // Authenticate the ping (verify its AEAD tag) BEFORE acting on it. Pings carry + // an empty plaintext but still a tag, so this rejects a spoofed ping that + // could otherwise replay-poison or, worse, hijack `endpoint` migration. + let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; let seq = { let mut locked = state.lock().expect("server state poisoned"); - let mut found = None; - for session in locked.sessions.values_mut() { - if let Some(client) = session.clients.get_mut(&packet.header.conn_id) { - if !client.replay.accept(packet.header.seq) { - return Ok(()); - } - client.last_seen = Instant::now(); - client.send_seq += 1; - found = Some(client.send_seq); - break; - } + let client = locked + .client_mut(&packet.header.conn_id) + .ok_or_else(|| anyhow!("unknown client"))?; + if !client.replay.accept(packet.header.seq) { + return Ok(()); } - found.ok_or_else(|| anyhow!("unknown client"))? + // Connection migration: any authenticated, fresh packet from a new source + // address roams the session there (spec §11). Gated on replay-accept so a + // captured old packet can't redirect the stream. + client.endpoint = peer; + client.last_seen = Instant::now(); + client.send_seq += 1; + client.send_seq }; let out = protocol::encode_encrypted( PacketKind::Pong, @@ -1120,39 +1332,64 @@ async fn handle_ping( async fn handle_detach(state: &Arc>, packet: &protocol::Packet) -> Result<()> { let mut locked = state.lock().expect("server state poisoned"); - for session in locked.sessions.values_mut() { - session.clients.remove(&packet.header.conn_id); - } + locked.remove_client_everywhere(&packet.header.conn_id); Ok(()) } fn remove_client(state: &Arc>, client_id: [u8; 16]) { let mut locked = state.lock().expect("server state poisoned"); - for session in locked.sessions.values_mut() { - session.clients.remove(&client_id); - } + locked.remove_client_everywhere(&client_id); } -async fn handle_ack(state: &Arc>, packet: &protocol::Packet) -> Result<()> { - let (key, _) = find_client_key(state, &packet.header.conn_id)?; +async fn handle_ack( + state: &Arc>, + peer: SocketAddr, + packet: &protocol::Packet, +) -> Result<()> { + let (key, _) = find_client_decrypt_key(state, &packet.header)?; let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; let mut locked = state.lock().expect("server state poisoned"); - for session in locked.sessions.values_mut() { - if let Some(client) = session.clients.get_mut(&packet.header.conn_id) { - if !client.replay.accept(packet.header.seq) { - return Ok(()); - } - client.last_seen = Instant::now(); - client.last_acked = packet.header.ack; - while client - .pending - .front() - .is_some_and(|pending| pending.output_seq <= packet.header.ack) - { - client.pending.pop_front(); - } + if let Some(client) = locked.client_mut(&packet.header.conn_id) { + if !client.replay.accept(packet.header.seq) { return Ok(()); } + // Connection migration on any authenticated, fresh packet (spec §11). + client.endpoint = peer; + client.last_seen = Instant::now(); + client.last_acked = packet.header.ack; + while client + .pending + .front() + .is_some_and(|pending| pending.output_seq <= packet.header.ack) + { + client.pending.pop_front(); + } + } + Ok(()) +} + +/// The client confirms it has switched to the new epoch key. Sealed under the +/// new (now-current) key, so decrypting it authenticates the confirmation. We +/// retire the previous-epoch key immediately, ending the grace early. +async fn handle_rekey_ack( + state: &Arc>, + peer: SocketAddr, + packet: &protocol::Packet, +) -> Result<()> { + let (key, _) = find_client_decrypt_key(state, &packet.header)?; + let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; + let mut locked = state.lock().expect("server state poisoned"); + if let Some(client) = locked.client_mut(&packet.header.conn_id) { + if !client.replay.accept(packet.header.seq) { + return Ok(()); + } + client.endpoint = peer; + client.last_seen = Instant::now(); + // Only retire the old key if the ack was sealed under the *current* key, + // confirming the client really is on the new epoch. + if protocol::session_key_id(&client.session_key) == packet.header.session_key_id { + client.previous_session_key = None; + } } Ok(()) } @@ -1163,7 +1400,7 @@ async fn handle_stream_open( peer: SocketAddr, packet: &protocol::Packet, ) -> Result<()> { - let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; + let (key, session_name) = find_client_decrypt_key(state, &packet.header)?; let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; let open: StreamOpen = protocol::from_body(&body)?; let allowed = { @@ -1276,7 +1513,7 @@ async fn handle_stream_data( peer: SocketAddr, packet: &protocol::Packet, ) -> Result<()> { - let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; + let (key, session_name) = find_client_decrypt_key(state, &packet.header)?; let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; let data: StreamData = protocol::from_body(&body)?; let writer = { @@ -1314,7 +1551,7 @@ async fn handle_stream_open_ok( peer: SocketAddr, packet: &protocol::Packet, ) -> Result<()> { - let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; + let (key, session_name) = find_client_decrypt_key(state, &packet.header)?; let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; let ok: StreamOpenOk = protocol::from_body(&body)?; { @@ -1346,7 +1583,7 @@ async fn handle_stream_open_reject( peer: SocketAddr, packet: &protocol::Packet, ) -> Result<()> { - let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; + let (key, session_name) = find_client_decrypt_key(state, &packet.header)?; let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; let reject: StreamOpenReject = protocol::from_body(&body)?; let mut locked = state.lock().expect("server state poisoned"); @@ -1376,7 +1613,7 @@ async fn handle_stream_window_adjust( peer: SocketAddr, packet: &protocol::Packet, ) -> Result<()> { - let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; + let (key, session_name) = find_client_decrypt_key(state, &packet.header)?; let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; let adjust: StreamWindowAdjust = protocol::from_body(&body)?; { @@ -1406,9 +1643,10 @@ async fn handle_stream_window_adjust( async fn handle_stream_close( state: &Arc>, + peer: SocketAddr, packet: &protocol::Packet, ) -> Result<()> { - let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; + let (key, session_name) = find_client_decrypt_key(state, &packet.header)?; let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; let close: StreamClose = protocol::from_body(&body)?; let mut locked = state.lock().expect("server state poisoned"); @@ -1421,6 +1659,8 @@ async fn handle_stream_close( if !client.replay.accept(packet.header.seq) { return Ok(()); } + // Connection migration on any authenticated, fresh packet (spec §11). + client.endpoint = peer; client.last_seen = Instant::now(); client.stream_writers.remove(&close.stream_id); client.opened_streams.remove(&close.stream_id); @@ -1650,42 +1890,39 @@ async fn queue_or_send_stream_data_to_client( let send = { let mut locked = state.lock().expect("server state poisoned"); let mut send = None; - for session in locked.sessions.values_mut() { - if let Some(client) = session.clients.get_mut(&client_id) { - if !client.opened_streams.contains(&stream_id) - || client - .stream_send_credit - .get(&stream_id) - .copied() - .unwrap_or(0) - < bytes.len() - || client - .stream_pending_data - .get(&stream_id) - .is_some_and(|pending| !pending.is_empty()) - { - client - .stream_pending_data - .entry(stream_id) - .or_default() - .push_back(bytes); - return Ok(()); - } - *client.stream_send_credit.entry(stream_id).or_default() -= bytes.len(); - client.send_seq += 1; - let body = protocol::to_body(&StreamData { stream_id, bytes })?; - let packet = protocol::encode_encrypted( - PacketKind::StreamData, - client_id, - client.send_seq, - client.last_acked, - &client.session_key, - SERVER_TO_CLIENT, - &body, - )?; - send = Some((client.endpoint, packet)); - break; + if let Some(client) = locked.client_mut(&client_id) { + if !client.opened_streams.contains(&stream_id) + || client + .stream_send_credit + .get(&stream_id) + .copied() + .unwrap_or(0) + < bytes.len() + || client + .stream_pending_data + .get(&stream_id) + .is_some_and(|pending| !pending.is_empty()) + { + client + .stream_pending_data + .entry(stream_id) + .or_default() + .push_back(bytes); + return Ok(()); } + *client.stream_send_credit.entry(stream_id).or_default() -= bytes.len(); + client.send_seq += 1; + let body = protocol::to_body(&StreamData { stream_id, bytes })?; + let packet = protocol::encode_encrypted( + PacketKind::StreamData, + client_id, + client.send_seq, + client.last_acked, + &client.session_key, + SERVER_TO_CLIENT, + &body, + )?; + send = Some((client.endpoint, packet)); } send }; @@ -1705,42 +1942,39 @@ async fn flush_stream_pending_data_to_client( let send = { let mut locked = state.lock().expect("server state poisoned"); let mut send = None; - for session in locked.sessions.values_mut() { - if let Some(client) = session.clients.get_mut(&client_id) { - let Some(pending) = client.stream_pending_data.get_mut(&stream_id) else { - return Ok(()); - }; - let Some(bytes) = pending.front() else { - client.stream_pending_data.remove(&stream_id); - return Ok(()); - }; - let credit = client - .stream_send_credit - .get(&stream_id) - .copied() - .unwrap_or(0); - if credit < bytes.len() { - return Ok(()); - } - let bytes = pending.pop_front().expect("pending front exists"); - if pending.is_empty() { - client.stream_pending_data.remove(&stream_id); - } - *client.stream_send_credit.entry(stream_id).or_default() -= bytes.len(); - client.send_seq += 1; - let body = protocol::to_body(&StreamData { stream_id, bytes })?; - let packet = protocol::encode_encrypted( - PacketKind::StreamData, - client_id, - client.send_seq, - client.last_acked, - &client.session_key, - SERVER_TO_CLIENT, - &body, - )?; - send = Some((client.endpoint, packet)); - break; + if let Some(client) = locked.client_mut(&client_id) { + let Some(pending) = client.stream_pending_data.get_mut(&stream_id) else { + return Ok(()); + }; + let Some(bytes) = pending.front() else { + client.stream_pending_data.remove(&stream_id); + return Ok(()); + }; + let credit = client + .stream_send_credit + .get(&stream_id) + .copied() + .unwrap_or(0); + if credit < bytes.len() { + return Ok(()); } + let bytes = pending.pop_front().expect("pending front exists"); + if pending.is_empty() { + client.stream_pending_data.remove(&stream_id); + } + *client.stream_send_credit.entry(stream_id).or_default() -= bytes.len(); + client.send_seq += 1; + let body = protocol::to_body(&StreamData { stream_id, bytes })?; + let packet = protocol::encode_encrypted( + PacketKind::StreamData, + client_id, + client.send_seq, + client.last_acked, + &client.session_key, + SERVER_TO_CLIENT, + &body, + )?; + send = Some((client.endpoint, packet)); } send }; @@ -1785,14 +2019,11 @@ async fn send_stream_close_to_client( ) -> Result<()> { { let mut locked = state.lock().expect("server state poisoned"); - for session in locked.sessions.values_mut() { - if let Some(client) = session.clients.get_mut(&client_id) { - client.stream_writers.remove(&stream_id); - client.opened_streams.remove(&stream_id); - client.stream_send_credit.remove(&stream_id); - client.stream_pending_data.remove(&stream_id); - break; - } + if let Some(client) = locked.client_mut(&client_id) { + client.stream_writers.remove(&stream_id); + client.opened_streams.remove(&stream_id); + client.stream_send_credit.remove(&stream_id); + client.stream_pending_data.remove(&stream_id); } } let body = protocol::to_body(&StreamClose { stream_id })?; @@ -1809,21 +2040,18 @@ async fn send_stream_packet_to_client( let send = { let mut locked = state.lock().expect("server state poisoned"); let mut found = None; - for session in locked.sessions.values_mut() { - if let Some(client) = session.clients.get_mut(&client_id) { - client.send_seq += 1; - let packet = protocol::encode_encrypted( - kind, - client_id, - client.send_seq, - client.last_acked, - &client.session_key, - SERVER_TO_CLIENT, - &body, - )?; - found = Some((client.endpoint, packet)); - break; - } + if let Some(client) = locked.client_mut(&client_id) { + client.send_seq += 1; + let packet = protocol::encode_encrypted( + kind, + client_id, + client.send_seq, + client.last_acked, + &client.session_key, + SERVER_TO_CLIENT, + &body, + )?; + found = Some((client.endpoint, packet)); } found }; @@ -1860,6 +2088,8 @@ async fn broadcast_output( let mut sends = Vec::new(); for (client_id, client) in session.clients.iter_mut() { client.send_seq += 1; + // Count traffic toward the packet-count rekey trigger (spec §11). + client.epoch_packets = client.epoch_packets.saturating_add(1); // Only materialize a screen snapshot when this client has fallen too // far behind; the common case sends the raw output bytes and must not // clone the whole vt100 grid per client per packet. @@ -1919,6 +2149,7 @@ async fn broadcast_session_exit( let output_seq = session.output_seq; let mut sends = Vec::new(); for (client_id, mut client) in session.clients.drain() { + locked.client_index.remove(&client_id); client.send_seq += 1; let frame = Frame { session: session_name.clone(), @@ -1990,6 +2221,102 @@ async fn retransmit_pending( Ok(()) } +/// How long the previous epoch's key is retained after a rekey so in-flight +/// pre-rekey packets still decrypt instead of being dropped as fatal. +const REKEY_PREVIOUS_KEY_GRACE_SECS: u64 = 5; + +/// Rotate transport traffic keys for any client past its rekey threshold (spec +/// §11). Triggers on either a packet count or a wall-clock interval, configured +/// via `rekey_after_packets` / `rekey_after_secs`. The new key is derived from +/// fresh server material independent of the handshake keys +/// (see [`dosh::native::derive_rekey_session_key`]); the server installs it +/// immediately, retains the old key for a short grace, and ships the fresh +/// material to the client in a `Rekey` packet sealed under the *current* key. +async fn maybe_rekey_clients( + state: &Arc>, + socket: &Arc, +) -> Result<()> { + let sends = { + let mut locked = state.lock().expect("server state poisoned"); + let after_packets = locked.config.rekey_after_packets; + let after_secs = locked.config.rekey_after_secs; + if after_packets == 0 && after_secs == 0 { + return Ok(()); + } + let now = Instant::now(); + let grace = Duration::from_secs(REKEY_PREVIOUS_KEY_GRACE_SECS); + let mut sends = Vec::new(); + for session in locked.sessions.values_mut() { + for (client_id, client) in session.clients.iter_mut() { + // Expire the retained previous-epoch key once its grace is over. + if client.previous_session_key.is_some() + && now.duration_since(client.epoch_started) >= grace + { + client.previous_session_key = None; + } + let by_packets = after_packets != 0 && client.epoch_packets >= after_packets; + let by_time = after_secs != 0 + && now.duration_since(client.epoch_started).as_secs() >= after_secs; + if !(by_packets || by_time) { + continue; + } + // Don't start a new rekey while the previous epoch's grace is + // still active, so we never juggle three live keys at once. + if client.previous_session_key.is_some() { + continue; + } + let previous_key = client.session_key; + let previous_key_id = protocol::session_key_id(&previous_key); + let new_epoch = client.epoch.saturating_add(1); + let rekey_material = crypto::random_32(); + let new_key = match dosh::native::derive_rekey_session_key( + &previous_key, + &rekey_material, + &previous_key_id, + new_epoch, + ) { + Ok(key) => key, + Err(err) => { + eprintln!("rekey derive error: {err:#}"); + continue; + } + }; + let new_key_id = protocol::session_key_id(&new_key); + let rekey = protocol::Rekey { + epoch: new_epoch, + rekey_material, + new_session_key_id: new_key_id, + }; + let body = protocol::to_body(&rekey)?; + client.send_seq += 1; + // The Rekey itself is sealed under the CURRENT key so the client + // can decrypt it before switching. + let packet = protocol::encode_encrypted( + PacketKind::Rekey, + *client_id, + client.send_seq, + client.last_acked, + &previous_key, + SERVER_TO_CLIENT, + &body, + )?; + // Install the new key, keep the old one for the grace window. + client.previous_session_key = Some(previous_key); + client.session_key = new_key; + client.epoch = new_epoch; + client.epoch_started = now; + client.epoch_packets = 0; + sends.push((client.endpoint, packet)); + } + } + sends + }; + for (endpoint, packet) in sends { + socket.send_to(&packet, endpoint).await?; + } + Ok(()) +} + fn cleanup_disconnected_clients(state: &Arc>) { // Sessions removed here are dropped after the lock is released so their // shells are killed (PtyHandle::drop) without holding up the event loop. @@ -1998,16 +2325,26 @@ fn cleanup_disconnected_clients(state: &Arc>) { let timeout = Duration::from_secs(locked.config.client_timeout_secs.max(1)); let now = Instant::now(); let prewarm: HashSet = locked.config.prewarm_sessions.iter().cloned().collect(); + // Collect timed-out clients first so we can purge them from the O(1) + // conn_id index too, keeping it in lockstep with `Session::clients`. + let mut timed_out: Vec<[u8; 16]> = Vec::new(); for session in locked.sessions.values_mut() { - session - .clients - .retain(|_, client| now.duration_since(client.last_seen) <= timeout); + session.clients.retain(|client_id, client| { + let keep = now.duration_since(client.last_seen) <= timeout; + if !keep { + timed_out.push(*client_id); + } + keep + }); if session.clients.is_empty() { session.empty_since.get_or_insert(now); } else { session.empty_since = None; } } + for client_id in timed_out { + locked.client_index.remove(&client_id); + } // Evict half-finished native handshakes. Each unauthenticated ClientHello // inserts a PendingNativeAuth that is otherwise only removed by a matching // UserAuth, so without this a flood of hellos grows the map without bound. @@ -2015,6 +2352,9 @@ fn cleanup_disconnected_clients(state: &Arc>) { locked .pending_native .retain(|_, pending| now.duration_since(pending.created_at) < handshake_ttl); + // Drop fully-refilled rate-limit buckets so a spoofed-source flood cannot + // grow the map without bound. + locked.native_auth_limiter.evict_full(now); // Reap sessions that have been clientless past the grace period. Prewarmed // sessions stay hot on purpose, even with no clients attached. let to_reap: Vec = locked @@ -2041,12 +2381,22 @@ fn find_client_key( client_id: &[u8; 16], ) -> Result<([u8; 32], String)> { let locked = state.lock().expect("server state poisoned"); - for (name, session) in &locked.sessions { - if let Some(client) = session.clients.get(client_id) { - return Ok((client.session_key, name.clone())); - } - } - Err(anyhow!("unknown client")) + locked + .lookup_client(client_id) + .ok_or_else(|| anyhow!("unknown client")) +} + +/// Like [`find_client_key`] but selects the key (current or retained previous +/// epoch) matching the packet header's `session_key_id`, so a packet sealed +/// under the pre-rekey key during the grace window still decrypts. +fn find_client_decrypt_key( + state: &Arc>, + header: &protocol::Header, +) -> Result<([u8; 32], String)> { + let locked = state.lock().expect("server state poisoned"); + locked + .lookup_client_key_for(&header.conn_id, &header.session_key_id) + .ok_or_else(|| anyhow!("unknown client")) } fn parse_nonce(raw: &str) -> Result<[u8; 12]> { @@ -2067,6 +2417,134 @@ fn parse_size(raw: &str) -> Result<(u16, u16)> { mod tests { use super::*; + #[test] + fn native_auth_rate_limiter_throttles_per_ip_and_refills() { + let mut limiter = NativeAuthRateLimiter::new(3); + let ip_a: IpAddr = "10.0.0.1".parse().unwrap(); + let ip_b: IpAddr = "10.0.0.2".parse().unwrap(); + let t0 = Instant::now(); + + // Burst of 3 allowed, 4th rejected for ip_a. + assert!(limiter.check(ip_a, t0).is_ok()); + assert!(limiter.check(ip_a, t0).is_ok()); + assert!(limiter.check(ip_a, t0).is_ok()); + assert!(limiter.check(ip_a, t0).is_err()); + + // A different source IP has its own independent bucket. + assert!(limiter.check(ip_b, t0).is_ok()); + + // After 20s at 3/min (one token every 20s), ip_a gets exactly one back. + let t1 = t0 + Duration::from_secs(20); + assert!(limiter.check(ip_a, t1).is_ok()); + assert!(limiter.check(ip_a, t1).is_err()); + } + + #[test] + fn native_auth_rate_limiter_zero_per_minute_blocks_everything() { + let mut limiter = NativeAuthRateLimiter::new(0); + assert!( + limiter + .check("127.0.0.1".parse().unwrap(), Instant::now()) + .is_err() + ); + } + + #[test] + fn native_auth_rate_limiter_evicts_refilled_buckets() { + let mut limiter = NativeAuthRateLimiter::new(60); + let ip: IpAddr = "10.0.0.9".parse().unwrap(); + let t0 = Instant::now(); + assert!(limiter.check(ip, t0).is_ok()); + assert_eq!(limiter.buckets.len(), 1); + // 60/min = one token per second; after 60s the bucket is full again. + limiter.evict_full(t0 + Duration::from_secs(60)); + assert!(limiter.buckets.is_empty()); + } + + fn test_client_state(session_key: [u8; 32]) -> ClientState { + ClientState { + endpoint: "127.0.0.1:9".parse().unwrap(), + mode: "forward-only".to_string(), + session_key, + last_acked: 0, + replay: ReplayWindow::default(), + send_seq: 1, + cols: 80, + rows: 24, + last_seen: Instant::now(), + pending: VecDeque::new(), + allowed_forwardings: Vec::new(), + stream_writers: HashMap::new(), + opened_streams: HashSet::new(), + stream_send_credit: HashMap::new(), + stream_pending_data: HashMap::new(), + epoch: 0, + epoch_started: Instant::now(), + epoch_packets: 0, + previous_session_key: None, + } + } + + #[test] + fn client_index_stays_in_sync_with_session_clients() { + let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); + state + .ensure_session("work", 80, 24, "forward-only", &[]) + .unwrap(); + let client_id = [5u8; 16]; + let key = [6u8; 32]; + + // Inserting through the helper records the index for O(1) lookup. + state.insert_client("work", client_id, test_client_state(key)); + assert_eq!( + state.client_index.get(&client_id).map(String::as_str), + Some("work") + ); + assert_eq!( + state.lookup_client(&client_id), + Some((key, "work".to_string())) + ); + assert!(state.client_mut(&client_id).is_some()); + + // Removing purges both the session map and the index. + assert!(state.remove_client_everywhere(&client_id)); + assert!(!state.client_index.contains_key(&client_id)); + assert!(state.lookup_client(&client_id).is_none()); + assert!(state.client_mut(&client_id).is_none()); + assert!(!state.sessions["work"].clients.contains_key(&client_id)); + // Removing an absent client is a no-op. + assert!(!state.remove_client_everywhere(&client_id)); + } + + #[test] + fn cleanup_purges_timed_out_clients_from_index() { + let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let config = ServerConfig { + client_timeout_secs: 1, + ..ServerConfig::default() + }; + let mut state = ServerState::new(config, [0u8; 32], pty_tx); + state + .ensure_session("work", 80, 24, "forward-only", &[]) + .unwrap(); + let client_id = [8u8; 16]; + let mut client = test_client_state([1u8; 32]); + // Make the client look stale so cleanup reaps it. + client.last_seen = Instant::now() - Duration::from_secs(60); + state.insert_client("work", client_id, client); + + let state = Arc::new(Mutex::new(state)); + cleanup_disconnected_clients(&state); + + let locked = state.lock().unwrap(); + assert!( + !locked.client_index.contains_key(&client_id), + "timed-out client must be dropped from the index, not leaked" + ); + assert!(locked.lookup_client(&client_id).is_none()); + } + #[test] fn forward_only_session_does_not_allocate_pty() { let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); @@ -2196,6 +2674,10 @@ mod tests { opened_streams: HashSet::new(), stream_send_credit: HashMap::new(), stream_pending_data: HashMap::new(), + epoch: 0, + epoch_started: Instant::now(), + epoch_packets: 0, + previous_session_key: None, }, )]), output_seq: 0, @@ -2203,6 +2685,8 @@ mod tests { empty_since: None, }, ); + // Keep the O(1) conn_id index in sync, matching `insert_client`. + state.client_index.insert(client_id, "test".to_string()); let state = Arc::new(Mutex::new(state)); queue_or_send_stream_data_to_client(&state, &sender, client_id, 42, b"hello".to_vec()) diff --git a/src/config.rs b/src/config.rs index 0fb0f45..3f96c9e 100644 --- a/src/config.rs +++ b/src/config.rs @@ -28,6 +28,14 @@ pub struct ServerConfig { pub authorized_keys: Vec, #[serde(default = "default_native_auth_rate_limit_per_minute")] pub native_auth_rate_limit_per_minute: u32, + /// Rotate a client's transport traffic key after this many packets in the + /// current epoch (spec §11). `0` disables the packet-count trigger. + #[serde(default = "default_rekey_after_packets")] + pub rekey_after_packets: u64, + /// Rotate a client's transport traffic key after this many wall-clock seconds + /// in the current epoch (spec §11). `0` disables the time trigger. + #[serde(default = "default_rekey_after_secs")] + pub rekey_after_secs: u64, #[serde(default = "default_true")] pub allow_tcp_forwarding: bool, #[serde(default)] @@ -61,6 +69,8 @@ impl Default for ServerConfig { host_key: default_host_key(), authorized_keys: default_authorized_keys(), native_auth_rate_limit_per_minute: default_native_auth_rate_limit_per_minute(), + rekey_after_packets: default_rekey_after_packets(), + rekey_after_secs: default_rekey_after_secs(), allow_tcp_forwarding: true, allow_remote_forwarding: false, allow_remote_non_loopback_bind: false, @@ -175,6 +185,18 @@ fn default_native_auth_rate_limit_per_minute() -> u32 { 30 } +fn default_rekey_after_packets() -> u64 { + // Rotate well before any AEAD nonce-reuse concern; ChaCha20-Poly1305 with a + // per-direction monotonic seq is safe far beyond this, but a bounded epoch + // keeps forward-secrecy windows small. + 100_000 +} + +fn default_rekey_after_secs() -> u64 { + // One hour per epoch by default. + 3600 +} + fn default_auth_preference() -> String { "native,ssh".to_string() } diff --git a/src/native.rs b/src/native.rs index d3af508..765d20f 100644 --- a/src/native.rs +++ b/src/native.rs @@ -15,7 +15,7 @@ use std::str::FromStr; use x25519_dalek::{PublicKey as X25519PublicKey, StaticSecret}; pub const HOST_KEY_ALGORITHM: &str = "dosh-ed25519"; -pub const NATIVE_PROTOCOL_VERSION: u8 = 1; +pub const NATIVE_PROTOCOL_VERSION: u8 = 2; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct HostPublicKey { @@ -248,17 +248,55 @@ pub fn sign_server_hello( Ok(()) } +/// Error raised when a native handshake peer speaks a different protocol version. +/// +/// Carries the peer's advertised version so callers can render an actionable, +/// named message ("upgrade dosh") rather than letting a mismatch surface as an +/// opaque decrypt failure or a silent timeout. The `Display` text deliberately +/// embeds [`crate::protocol::VERSION_MISMATCH_REASON`] so the server's reject and +/// the client's local error read the same way. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ProtocolVersionMismatch { + pub local: u8, + pub remote: u8, + pub peer: &'static str, +} + +impl std::fmt::Display for ProtocolVersionMismatch { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{} ({} speaks native protocol v{}, this build speaks v{})", + crate::protocol::VERSION_MISMATCH_REASON, + self.peer, + self.remote, + self.local, + ) + } +} + +impl std::error::Error for ProtocolVersionMismatch {} + +/// Verify a peer-advertised native protocol version against this build. +/// +/// `peer` names whose version was wrong ("client" or "server") for the error +/// message. Returns a typed [`ProtocolVersionMismatch`] so the caller can both +/// match on it and print an actionable, upgrade-oriented message. +pub fn check_native_protocol_version(version: u8, peer: &'static str) -> Result<()> { + if version != NATIVE_PROTOCOL_VERSION { + return Err(ProtocolVersionMismatch { + local: NATIVE_PROTOCOL_VERSION, + remote: version, + peer, + } + .into()); + } + Ok(()) +} + pub fn verify_server_hello(client: &NativeClientHello, server: &NativeServerHello) -> Result<()> { - anyhow::ensure!( - client.protocol_version == NATIVE_PROTOCOL_VERSION, - "unsupported native client protocol {}", - client.protocol_version - ); - anyhow::ensure!( - server.protocol_version == NATIVE_PROTOCOL_VERSION, - "unsupported native server protocol {}", - server.protocol_version - ); + check_native_protocol_version(client.protocol_version, "client")?; + check_native_protocol_version(server.protocol_version, "server")?; let transcript = server_hello_transcript(client, server)?; verify_host_signature(&server.host_key, &transcript, &server.host_signature) } @@ -369,6 +407,36 @@ pub fn derive_native_session_key( crypto::hkdf32(shared.as_bytes(), &salt, b"dosh/native/chacha20poly1305/v1") } +/// Derive a rotated transport key for transport rekey (spec §11 / §9). +/// +/// The rotated key is derived **independently of the handshake traffic keys**: +/// the input keying material is the current epoch's key (which both peers already +/// hold) mixed with fresh, server-generated `rekey_material` (32 bytes of CSPRNG +/// output, delivered confidentially inside an AEAD `Rekey` packet sealed under +/// the current key). The new `epoch` and the previous epoch's `session_key_id` +/// salt the derivation so each epoch's key is unique. It never re-derives from +/// the handshake DH output, satisfying the spec requirement that "rotated session +/// keys must be derived independently ... from fresh randomness" and "must not +/// reuse handshake traffic keys." +/// +/// Both peers run this identically — the client never needs the server's +/// long-term secret, only the fresh `rekey_material` it receives in the `Rekey` +/// packet plus the current key it already shares. +pub fn derive_rekey_session_key( + current_key: &[u8; 32], + rekey_material: &[u8; 32], + previous_session_key_id: &[u8; 16], + epoch: u64, +) -> Result<[u8; 32]> { + let mut ikm = Vec::with_capacity(64); + ikm.extend_from_slice(current_key); + ikm.extend_from_slice(rekey_material); + let mut salt = b"dosh/native/rekey/v1".to_vec(); + salt.extend_from_slice(previous_session_key_id); + salt.extend_from_slice(&epoch.to_be_bytes()); + crypto::hkdf32(&ikm, &salt, b"dosh/native/rekey/chacha20poly1305/v1") +} + pub fn load_ed25519_identity(path: &Path) -> Result { load_ed25519_identity_with_passphrase(path, None) } diff --git a/src/protocol.rs b/src/protocol.rs index ac6ac84..f4c9add 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -5,8 +5,14 @@ use anyhow::{Context, Result, bail}; use serde::{Deserialize, Serialize}; pub const MAGIC: &[u8; 4] = b"DOSH"; -pub const VERSION: u8 = 1; +pub const VERSION: u8 = 2; pub const HEADER_LEN: usize = 58; + +/// Stable, user-facing reason string the server puts in an `AttachReject` when a +/// native handshake arrives carrying a `protocol_version` it cannot speak. The +/// client recognizes this prefix and surfaces a clear "upgrade dosh" message +/// instead of the generic transport rejection or, worse, a silent timeout. +pub const VERSION_MISMATCH_REASON: &str = "protocol version mismatch — upgrade dosh"; const HEADER_AAD_LEN: usize = HEADER_LEN - 2; pub const CLIENT_TO_SERVER: u32 = 1; pub const SERVER_TO_CLIENT: u32 = 2; @@ -39,6 +45,8 @@ pub enum PacketKind { StreamEof = 23, StreamClose = 24, NativeAuthCheckOk = 25, + Rekey = 26, + RekeyAck = 27, } impl TryFrom for PacketKind { @@ -71,6 +79,8 @@ impl TryFrom for PacketKind { 23 => Self::StreamEof, 24 => Self::StreamClose, 25 => Self::NativeAuthCheckOk, + 26 => Self::Rekey, + 27 => Self::RekeyAck, _ => bail!("unknown packet kind {value}"), }) } @@ -110,7 +120,11 @@ impl Header { bail!("bad magic"); } if input[4] != VERSION { - bail!("bad protocol version {}", input[4]); + bail!( + "{} (peer wire protocol v{}, this build speaks v{VERSION})", + VERSION_MISMATCH_REASON, + input[4] + ); } let kind = PacketKind::try_from(input[5])?; let flags = u16::from_be_bytes(input[6..8].try_into().unwrap()); @@ -133,6 +147,20 @@ impl Header { } } +/// Cheaply inspect a datagram that carries our [`MAGIC`] but whose wire +/// [`VERSION`] byte differs from this build's. Returns the peer's advertised +/// wire version when (and only when) the packet is a Dosh packet we cannot +/// otherwise decode because of a version skew, so the receiver can answer with a +/// clear, named version-mismatch reject instead of dropping it (a silent +/// timeout for the peer). Returns `None` for our own version, foreign magic, or +/// runt packets. +pub fn peek_foreign_wire_version(input: &[u8]) -> Option { + if input.len() < 5 || &input[..4] != MAGIC || input[4] == VERSION { + return None; + } + Some(input[4]) +} + #[derive(Debug, Clone)] pub struct Packet { pub header: Header, @@ -369,6 +397,20 @@ pub struct StreamClose { pub stream_id: u64, } +/// Server→client transport rekey, sealed under the *current* session key. +/// +/// Carries the fresh server-generated `rekey_material` and the new `epoch`; both +/// peers feed these into [`crate::native::derive_rekey_session_key`] to compute +/// the next traffic key. `new_session_key_id` is the id the next epoch's packets +/// will carry, so the client can recognize and switch atomically. The client +/// replies with a [`PacketKind::RekeyAck`] encrypted under the *new* key. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Rekey { + pub epoch: u64, + pub rekey_material: [u8; 32], + pub new_session_key_id: [u8; 16], +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Frame { pub session: String, diff --git a/tests/integration_smoke.rs b/tests/integration_smoke.rs index 0acf62c..54c75aa 100644 --- a/tests/integration_smoke.rs +++ b/tests/integration_smoke.rs @@ -1246,3 +1246,314 @@ fn resume_updates_udp_endpoint_for_roaming() { "expected output on resumed socket, got {text:?}" ); } + +#[test] +fn transport_rekey_round_trip_keeps_session_alive() { + use dosh::native::derive_rekey_session_key; + use dosh::protocol::Rekey; + + let dir = tempfile::tempdir().unwrap(); + let port = free_udp_port(); + let config = write_server_config(&dir, port); + // Rotate after a single packet so a rekey fires almost immediately. + let mut raw = fs::read_to_string(&config).unwrap(); + raw.push_str("rekey_after_packets = 1\n"); + fs::write(&config, raw).unwrap(); + let mut server = start_server(&dir, &config); + + let (socket, bootstrap, ok) = direct_attach(&config, port, "read-write"); + socket + .set_read_timeout(Some(Duration::from_millis(300))) + .unwrap(); + + // Track the live transport key; it rotates when a Rekey arrives. + let mut current_key = bootstrap.session_key; + let mut current_key_id = protocol::session_key_id(¤t_key); + let mut send_seq = 2u64; + + // Produce some output so the server starts sending frames (and rekeys). + let input = Input { + bytes: b"printf DOSH_REKEY_ONE\\n\n".to_vec(), + }; + send_encrypted( + &socket, + port, + PacketKind::Input, + ok.client_id, + send_seq, + 0, + ¤t_key, + &protocol::to_body(&input).unwrap(), + ); + send_seq += 1; + + // Drive the loop: decrypt frames under the live key, and when a Rekey lands, + // adopt the new epoch key and ack it. Confirm a post-rekey input still works. + let mut rekeyed = false; + let mut saw_post_rekey_output = false; + let mut buf = [0u8; 65535]; + let deadline = std::time::Instant::now() + Duration::from_secs(8); + while std::time::Instant::now() < deadline { + let Ok((n, _)) = socket.recv_from(&mut buf) else { + continue; + }; + let Ok(packet) = protocol::decode(&buf[..n]) else { + continue; + }; + match packet.header.kind { + PacketKind::Rekey => { + let plain = + protocol::decrypt_body(&packet, ¤t_key, SERVER_TO_CLIENT).unwrap(); + let rekey: Rekey = protocol::from_body(&plain).unwrap(); + let new_key = derive_rekey_session_key( + ¤t_key, + &rekey.rekey_material, + ¤t_key_id, + rekey.epoch, + ) + .unwrap(); + assert_eq!( + protocol::session_key_id(&new_key), + rekey.new_session_key_id, + "client-derived rekey key id must match the server's" + ); + current_key = new_key; + current_key_id = rekey.new_session_key_id; + // Ack under the NEW key. + send_encrypted( + &socket, + port, + PacketKind::RekeyAck, + ok.client_id, + send_seq, + 0, + ¤t_key, + b"", + ); + send_seq += 1; + rekeyed = true; + // Now send a fresh input under the new key. + let input = Input { + bytes: b"printf DOSH_REKEY_TWO\\n\n".to_vec(), + }; + send_encrypted( + &socket, + port, + PacketKind::Input, + ok.client_id, + send_seq, + 0, + ¤t_key, + &protocol::to_body(&input).unwrap(), + ); + send_seq += 1; + } + PacketKind::Frame | PacketKind::ResumeOk => { + if let Ok(plain) = protocol::decrypt_body(&packet, ¤t_key, SERVER_TO_CLIENT) { + if let Ok(frame) = protocol::from_body::(&plain) { + let text = String::from_utf8_lossy(&frame.bytes); + if rekeyed && text.contains("DOSH_REKEY_TWO") { + saw_post_rekey_output = true; + break; + } + } + } + } + _ => {} + } + } + + let _ = server.kill(); + let _ = server.wait(); + + assert!(rekeyed, "server never initiated a rekey"); + assert!( + saw_post_rekey_output, + "post-rekey input/output round trip failed under the new epoch key" + ); +} + +#[test] +fn input_from_new_source_address_migrates_connection() { + // Spec §11: the server must accept client source-address migration after ANY + // valid authenticated/encrypted packet from a new address, not just resume. + let dir = tempfile::tempdir().unwrap(); + let port = free_udp_port(); + let config = write_server_config(&dir, port); + let mut server = start_server(&dir, &config); + let (_old_socket, bootstrap, ok) = direct_attach(&config, port, "read-write"); + + // A fresh socket = a new source address (new ephemeral port). The very first + // packet from it is an ordinary encrypted Input, not a ResumeRequest. + let new_socket = UdpSocket::bind("127.0.0.1:0").unwrap(); + new_socket + .set_read_timeout(Some(Duration::from_secs(2))) + .unwrap(); + let input = Input { + bytes: b"printf DOSH_MIGRATE\\n\n".to_vec(), + }; + let packet = protocol::encode_encrypted( + PacketKind::Input, + ok.client_id, + 2, + 0, + &bootstrap.session_key, + CLIENT_TO_SERVER, + &protocol::to_body(&input).unwrap(), + ) + .unwrap(); + new_socket + .send_to(&packet, format!("127.0.0.1:{port}")) + .unwrap(); + + // Output frames for this input must now be delivered to the NEW socket, + // proving the server migrated `endpoint` off the original address. + let text = collect_frame_text(&new_socket, &bootstrap.session_key, 2000); + + let _ = server.kill(); + let _ = server.wait(); + + assert!( + text.contains("DOSH_MIGRATE"), + "expected server output on the migrated socket, got {text:?}" + ); +} + +#[test] +fn native_auth_rate_limit_rejects_flood_before_crypto() { + use dosh::native::{NATIVE_PROTOCOL_VERSION, NativeClientHello}; + use dosh::protocol::NativeClientHelloBody; + + let dir = tempfile::tempdir().unwrap(); + let port = free_udp_port(); + let config = write_server_config(&dir, port); + // Squeeze the rate limit down to 2/min so a short burst trips it. + let mut raw = fs::read_to_string(&config).unwrap(); + raw.push_str("native_auth_rate_limit_per_minute = 2\n"); + fs::write(&config, raw).unwrap(); + write_native_client_auth(&dir, &config); + let mut server = start_server(&dir, &config); + + let socket = std::net::UdpSocket::bind("127.0.0.1:0").unwrap(); + socket + .set_read_timeout(Some(Duration::from_secs(2))) + .unwrap(); + let user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string()); + let make_hello = || { + let (_, client_public) = dosh::native::generate_native_ephemeral(); + let hello = NativeClientHello { + protocol_version: NATIVE_PROTOCOL_VERSION, + client_random: crypto::random_32(), + client_ephemeral_public: client_public, + requested_host: "local".to_string(), + requested_user: user.clone(), + requested_session: "default".to_string(), + requested_mode: "read-write".to_string(), + terminal_size: (80, 24), + supported_aead: vec!["chacha20poly1305".to_string()], + supported_user_key_algorithms: vec!["ssh-ed25519".to_string()], + cached_host_key_fingerprint: None, + attach_ticket_envelope: None, + requested_env: Vec::new(), + }; + protocol::encode_plain( + PacketKind::NativeClientHello, + [0u8; 16], + 1, + 0, + &protocol::to_body(&NativeClientHelloBody { hello }).unwrap(), + ) + .unwrap() + }; + + let mut reasons = Vec::new(); + // Burst of hellos; with a 2/min budget the later ones must be rejected. + for _ in 0..6 { + socket + .send_to(&make_hello(), format!("127.0.0.1:{port}")) + .unwrap(); + let mut buf = [0u8; 65535]; + if let Ok((n, _)) = socket.recv_from(&mut buf) { + let packet = protocol::decode(&buf[..n]).unwrap(); + if packet.header.kind == PacketKind::AttachReject { + let reject: AttachReject = protocol::from_body(&packet.body).unwrap(); + reasons.push(reject.reason); + } + } + } + + let _ = server.kill(); + let _ = server.wait(); + + assert!( + reasons.iter().any(|r| r.contains("rate limit")), + "expected a rate-limit reject within the burst, got {reasons:?}" + ); +} + +#[test] +fn native_hello_with_mismatched_protocol_version_gets_named_reject_not_hang() { + use dosh::native::{NATIVE_PROTOCOL_VERSION, NativeClientHello}; + use dosh::protocol::{NativeClientHelloBody, VERSION_MISMATCH_REASON}; + + let dir = tempfile::tempdir().unwrap(); + let port = free_udp_port(); + let config = write_server_config(&dir, port); + write_native_client_auth(&dir, &config); + let mut server = start_server(&dir, &config); + + let socket = std::net::UdpSocket::bind("127.0.0.1:0").unwrap(); + socket + .set_read_timeout(Some(Duration::from_secs(2))) + .unwrap(); + + // Same wire VERSION as the server (so the datagram decodes), but an + // incompatible native handshake protocol_version. This is the spec's + // plaintext negotiation point: the server must answer with a clear, named + // reject rather than letting the client time out. + let hello = NativeClientHello { + protocol_version: NATIVE_PROTOCOL_VERSION.wrapping_add(1), + client_random: crypto::random_32(), + client_ephemeral_public: [3u8; 32], + requested_host: "local".to_string(), + requested_user: std::env::var("USER").unwrap_or_else(|_| "unknown".to_string()), + requested_session: "default".to_string(), + requested_mode: "read-write".to_string(), + terminal_size: (80, 24), + supported_aead: vec!["chacha20poly1305".to_string()], + supported_user_key_algorithms: vec!["ssh-ed25519".to_string()], + cached_host_key_fingerprint: None, + attach_ticket_envelope: None, + requested_env: Vec::new(), + }; + let packet = protocol::encode_plain( + PacketKind::NativeClientHello, + [0u8; 16], + 1, + 0, + &protocol::to_body(&NativeClientHelloBody { hello }).unwrap(), + ) + .unwrap(); + socket + .send_to(&packet, format!("127.0.0.1:{port}")) + .unwrap(); + + let mut buf = [0u8; 65535]; + // recv with the 2s read timeout above: a hang would surface as a recv error + // here instead of a reject, which is exactly what this test guards against. + let (n, _) = socket + .recv_from(&mut buf) + .expect("server must answer a version-mismatched hello, not hang"); + let packet = protocol::decode(&buf[..n]).unwrap(); + let reject: AttachReject = protocol::from_body(&packet.body).unwrap(); + + let _ = server.kill(); + let _ = server.wait(); + + assert_eq!(packet.header.kind, PacketKind::AttachReject); + assert!( + reject.reason.contains(VERSION_MISMATCH_REASON), + "expected a named protocol-version-mismatch reject, got {:?}", + reject.reason + ); +} diff --git a/tests/protocol_auth.rs b/tests/protocol_auth.rs index 43f7f9f..1166359 100644 --- a/tests/protocol_auth.rs +++ b/tests/protocol_auth.rs @@ -126,6 +126,139 @@ fn attach_ticket_is_sealed_and_verifies_scope() { ); } +#[test] +fn peek_foreign_wire_version_flags_only_version_skew() { + let key = crypto::random_32(); + let mut packet = protocol::encode_encrypted( + PacketKind::Input, + crypto::random_16(), + 1, + 0, + &key, + CLIENT_TO_SERVER, + b"hi", + ) + .unwrap(); + // A correctly framed packet for this build is not "foreign". + assert_eq!(protocol::peek_foreign_wire_version(&packet), None); + // Bumping the wire version byte makes it undecodable but recognizable. + packet[4] = protocol::VERSION.wrapping_add(7); + assert_eq!( + protocol::peek_foreign_wire_version(&packet), + Some(protocol::VERSION.wrapping_add(7)) + ); + assert!(protocol::decode(&packet).is_err()); + // Non-Dosh datagrams and runts are ignored. + assert_eq!(protocol::peek_foreign_wire_version(b"XXXX\x01"), None); + assert_eq!(protocol::peek_foreign_wire_version(b"DOS"), None); +} + +#[test] +fn rekey_key_derivation_agrees_and_is_independent_per_epoch() { + use dosh::native::derive_rekey_session_key; + + // The handshake/current key both peers already share. + let current_key = crypto::random_32(); + let current_id = protocol::session_key_id(¤t_key); + // Fresh server-generated material, delivered confidentially in the Rekey. + let material = crypto::random_32(); + + // Both peers derive identically from shared current key + shipped material. + let server_view = derive_rekey_session_key(¤t_key, &material, ¤t_id, 1).unwrap(); + let client_view = derive_rekey_session_key(¤t_key, &material, ¤t_id, 1).unwrap(); + assert_eq!(server_view, client_view); + + // The rotated key must not equal the handshake/current key. + assert_ne!(server_view, current_key); + + // A different epoch (or different material) yields a different key. + let next_epoch = derive_rekey_session_key(¤t_key, &material, ¤t_id, 2).unwrap(); + assert_ne!(server_view, next_epoch); + let other_material = crypto::random_32(); + let other = derive_rekey_session_key(¤t_key, &other_material, ¤t_id, 1).unwrap(); + assert_ne!(server_view, other); +} + +#[test] +fn rekey_round_trip_decrypts_old_and_new_epoch_packets() { + use dosh::native::derive_rekey_session_key; + + let key_epoch0 = crypto::random_32(); + let id0 = protocol::session_key_id(&key_epoch0); + let conn_id = crypto::random_16(); + + // Pre-rekey packet sealed under epoch-0 key. + let pre = protocol::encode_encrypted( + PacketKind::Frame, + conn_id, + 5, + 0, + &key_epoch0, + protocol::SERVER_TO_CLIENT, + b"before", + ) + .unwrap(); + + // Rotate to epoch 1. + let material = crypto::random_32(); + let key_epoch1 = derive_rekey_session_key(&key_epoch0, &material, &id0, 1).unwrap(); + let post = protocol::encode_encrypted( + PacketKind::Frame, + conn_id, + 6, + 0, + &key_epoch1, + protocol::SERVER_TO_CLIENT, + b"after", + ) + .unwrap(); + + let pre = protocol::decode(&pre).unwrap(); + let post = protocol::decode(&post).unwrap(); + + // Each epoch's key carries its own session_key_id; the receiver picks the + // right one and both decrypt correctly. + assert_eq!(pre.header.session_key_id, id0); + assert_eq!( + pre.header.session_key_id, + protocol::session_key_id(&key_epoch0) + ); + assert_eq!( + post.header.session_key_id, + protocol::session_key_id(&key_epoch1) + ); + assert_eq!( + protocol::decrypt_body(&pre, &key_epoch0, protocol::SERVER_TO_CLIENT).unwrap(), + b"before" + ); + assert_eq!( + protocol::decrypt_body(&post, &key_epoch1, protocol::SERVER_TO_CLIENT).unwrap(), + b"after" + ); + + // A stale-epoch packet under the wrong key is rejected via session_key_id + // BEFORE any AEAD work — ignorable, not a fatal decrypt error. + let err = protocol::decrypt_body(&post, &key_epoch0, protocol::SERVER_TO_CLIENT).unwrap_err(); + assert!(err.to_string().contains("session key id")); +} + +#[test] +fn check_native_protocol_version_names_the_mismatch() { + use dosh::native::{NATIVE_PROTOCOL_VERSION, check_native_protocol_version}; + check_native_protocol_version(NATIVE_PROTOCOL_VERSION, "server").unwrap(); + let err = check_native_protocol_version(NATIVE_PROTOCOL_VERSION.wrapping_add(1), "server") + .unwrap_err(); + let message = err.to_string(); + assert!( + message.contains(protocol::VERSION_MISMATCH_REASON), + "expected actionable upgrade message, got {message:?}" + ); + assert!( + message.contains("server"), + "should name the wrong peer: {message:?}" + ); +} + #[test] fn replay_window_rejects_duplicates_but_allows_bounded_out_of_order() { let mut replay = ReplayWindow::new(8);