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<conn_id, session_name> 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 <noreply@anthropic.com>

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::<Rekey>(&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" {

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<std::path::PathBuf>) -> 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<String, Session>,
     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<IpAddr, TokenBucket>,
+}
+
+#[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<u32, ()> {
+        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<u64>,
     stream_send_credit: HashMap<u64, usize>,
     stream_pending_data: HashMap<u64, VecDeque<Vec<u8>>>,
+    /// 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<PtyOutput>,
     ) -> 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<u8>,
 ) -> 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(());
     }
-    if mode_allows_terminal_updates(&client.mode) {
+    // 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.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;
+        let client = locked
-        for session in locked.sessions.values_mut() {
+            .client_mut(&packet.header.conn_id)
-            if let Some(client) = session.clients.get_mut(&packet.header.conn_id) {
+            .ok_or_else(|| anyhow!("unknown client"))?;
         if !client.replay.accept(packet.header.seq) {
             return Ok(());
         }
+        // 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;
-                found = Some(client.send_seq);
+        client.send_seq
-                break;
-            }
-        }
-        found.ok_or_else(|| anyhow!("unknown client"))?
     };
     let out = protocol::encode_encrypted(
         PacketKind::Pong,
@@ -1120,28 +1332,29 @@ async fn handle_ping(
 
 async fn handle_detach(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packet) -> Result<()> {
     let mut locked = state.lock().expect("server state poisoned");
-    for session in locked.sessions.values_mut() {
+    locked.remove_client_everywhere(&packet.header.conn_id);
-        session.clients.remove(&packet.header.conn_id);
-    }
     Ok(())
 }
 
 fn remove_client(state: &Arc<Mutex<ServerState>>, client_id: [u8; 16]) {
     let mut locked = state.lock().expect("server state poisoned");
-    for session in locked.sessions.values_mut() {
+    locked.remove_client_everywhere(&client_id);
-        session.clients.remove(&client_id);
-    }
 }
 
-async fn handle_ack(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packet) -> Result<()> {
+async fn handle_ack(
-    let (key, _) = find_client_key(state, &packet.header.conn_id)?;
+    state: &Arc<Mutex<ServerState>>,
+    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) = locked.client_mut(&packet.header.conn_id) {
-        if let Some(client) = session.clients.get_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
@@ -1151,8 +1364,32 @@ async fn handle_ack(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packet)
         {
             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<Mutex<ServerState>>,
+    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<Mutex<ServerState>>,
+    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,8 +1890,7 @@ 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) = locked.client_mut(&client_id) {
-            if let Some(client) = session.clients.get_mut(&client_id) {
             if !client.opened_streams.contains(&stream_id)
                 || client
                     .stream_send_credit
@@ -1684,8 +1923,6 @@ async fn queue_or_send_stream_data_to_client(
                 &body,
             )?;
             send = Some((client.endpoint, packet));
-                break;
-            }
         }
         send
     };
@@ -1705,8 +1942,7 @@ 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) = locked.client_mut(&client_id) {
-                if let Some(client) = session.clients.get_mut(&client_id) {
                 let Some(pending) = client.stream_pending_data.get_mut(&stream_id) else {
                     return Ok(());
                 };
@@ -1739,8 +1975,6 @@ async fn flush_stream_pending_data_to_client(
                     &body,
                 )?;
                 send = Some((client.endpoint, packet));
-                    break;
-                }
             }
             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) = locked.client_mut(&client_id) {
-            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;
-            }
         }
     }
     let body = protocol::to_body(&StreamClose { stream_id })?;
@@ -1809,8 +2040,7 @@ 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) = locked.client_mut(&client_id) {
-            if let Some(client) = session.clients.get_mut(&client_id) {
             client.send_seq += 1;
             let packet = protocol::encode_encrypted(
                 kind,
@@ -1822,8 +2052,6 @@ async fn send_stream_packet_to_client(
                 &body,
             )?;
             found = Some((client.endpoint, packet));
-                break;
-            }
         }
         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<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+) -> 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<Mutex<ServerState>>) {
     // 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<Mutex<ServerState>>) {
         let timeout = Duration::from_secs(locked.config.client_timeout_secs.max(1));
         let now = Instant::now();
         let prewarm: HashSet<String> = 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
+            session.clients.retain(|client_id, client| {
-                .clients
+                let keep = now.duration_since(client.last_seen) <= timeout;
-                .retain(|_, client| 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<Mutex<ServerState>>) {
         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<String> = 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 {
+    locked
-        if let Some(client) = session.clients.get(client_id) {
+        .lookup_client(client_id)
-            return Ok((client.session_key, name.clone()));
+        .ok_or_else(|| anyhow!("unknown client"))
 }
-    }
+
-    Err(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<Mutex<ServerState>>,
+    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())

src/config.rs

@@ -28,6 +28,14 @@ pub struct ServerConfig {
     pub authorized_keys: Vec<String>,
     #[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()
 }

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!(
+    check_native_protocol_version(client.protocol_version, "client")?;
-        client.protocol_version == NATIVE_PROTOCOL_VERSION,
+    check_native_protocol_version(server.protocol_version, "server")?;
-        "unsupported native client protocol {}",
-        client.protocol_version
-    );
-    anyhow::ensure!(
-        server.protocol_version == NATIVE_PROTOCOL_VERSION,
-        "unsupported native server protocol {}",
-        server.protocol_version
-    );
     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<SigningKey> {
     load_ed25519_identity_with_passphrase(path, None)
 }

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<u8> for PacketKind {
@@ -71,6 +79,8 @@ impl TryFrom<u8> 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<u8> {
+    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,

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(&current_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,
+        &current_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, &current_key, SERVER_TO_CLIENT).unwrap();
+                let rekey: Rekey = protocol::from_body(&plain).unwrap();
+                let new_key = derive_rekey_session_key(
+                    &current_key,
+                    &rekey.rekey_material,
+                    &current_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,
+                    &current_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,
+                    &current_key,
+                    &protocol::to_body(&input).unwrap(),
+                );
+                send_seq += 1;
+            }
+            PacketKind::Frame | PacketKind::ResumeOk => {
+                if let Ok(plain) = protocol::decrypt_body(&packet, &current_key, SERVER_TO_CLIENT) {
+                    if let Ok(frame) = protocol::from_body::<Frame>(&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
+    );
+}

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(&current_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(&current_key, &material, &current_id, 1).unwrap();
+    let client_view = derive_rekey_session_key(&current_key, &material, &current_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(&current_key, &material, &current_id, 2).unwrap();
+    assert_ne!(server_view, next_epoch);
+    let other_material = crypto::random_32();
+    let other = derive_rekey_session_key(&current_key, &other_material, &current_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);