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/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!(
-        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<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,