Merge track A: VERSION discipline, rekey, migration, rate limiting, O(1) index

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,
 };
@@ -2153,6 +2153,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! {
@@ -2193,7 +2196,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,
@@ -2233,6 +2245,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" {