Add packet session key identifiers

src/auth.rs

@@ -1,5 +1,6 @@
 use crate::config::{ServerConfig, expand_tilde};
 use crate::crypto;
+use crate::protocol;
 use anyhow::{Context, Result};
 use base64::Engine;
 use base64::engine::general_purpose::URL_SAFE_NO_PAD;
@@ -105,12 +106,7 @@ pub fn build_bootstrap(
         &client_nonce,
         format!("dosh/session/{user}/{session}/{issued_at}").as_bytes(),
     )?;
-    let session_key_id = {
-        let digest = crypto::sha256(&session_key);
-        let mut out = [0u8; 16];
-        out.copy_from_slice(&digest[..16]);
-        out
-    };
+    let session_key_id = protocol::session_key_id(&session_key);
     let attach_token = attach_token(
         secret,
         &user,

src/bin/dosh-server.rs

@@ -435,12 +435,7 @@ async fn handle_native_user_auth(
                 }
             }
             let session_key = pending.session_key;
-            let key_id = {
-                let digest = crypto::sha256(&session_key);
-                let mut out = [0u8; 16];
-                out.copy_from_slice(&digest[..16]);
-                out
-            };
+            let key_id = protocol::session_key_id(&session_key);
             let attach_ticket_psk = crypto::random_32();
             let issued_at = now_secs()?;
             let attach_ticket = build_attach_ticket(
@@ -668,12 +663,7 @@ async fn handle_ticket_attach(
     }
 
     let session_key = crypto::random_32();
-    let session_key_id = {
-        let digest = crypto::sha256(&session_key);
-        let mut out = [0u8; 16];
-        out.copy_from_slice(&digest[..16]);
-        out
-    };
+    let session_key_id = protocol::session_key_id(&session_key);
     let (client_id, output_seq, snapshot) = {
         let mut locked = state.lock().expect("server state poisoned");
         if !locked.sessions.contains_key(&req.session) {

src/protocol.rs

@@ -6,7 +6,8 @@ use serde::{Deserialize, Serialize};
 
 pub const MAGIC: &[u8; 4] = b"DOSH";
 pub const VERSION: u8 = 1;
-pub const HEADER_LEN: usize = 42;
+pub const HEADER_LEN: usize = 58;
+const HEADER_AAD_LEN: usize = HEADER_LEN - 2;
 pub const CLIENT_TO_SERVER: u32 = 1;
 pub const SERVER_TO_CLIENT: u32 = 2;
 
@@ -80,6 +81,7 @@ pub struct Header {
     pub conn_id: [u8; 16],
     pub seq: u64,
     pub ack: u64,
+    pub session_key_id: [u8; 16],
     pub body_len: u16,
 }
 
@@ -93,7 +95,8 @@ impl Header {
         out[8..24].copy_from_slice(&self.conn_id);
         out[24..32].copy_from_slice(&self.seq.to_be_bytes());
         out[32..40].copy_from_slice(&self.ack.to_be_bytes());
-        out[40..42].copy_from_slice(&self.body_len.to_be_bytes());
+        out[40..56].copy_from_slice(&self.session_key_id);
+        out[56..58].copy_from_slice(&self.body_len.to_be_bytes());
         out
     }
 
@@ -113,13 +116,16 @@ impl Header {
         conn_id.copy_from_slice(&input[8..24]);
         let seq = u64::from_be_bytes(input[24..32].try_into().unwrap());
         let ack = u64::from_be_bytes(input[32..40].try_into().unwrap());
-        let body_len = u16::from_be_bytes(input[40..42].try_into().unwrap());
+        let mut session_key_id = [0u8; 16];
+        session_key_id.copy_from_slice(&input[40..56]);
+        let body_len = u16::from_be_bytes(input[56..58].try_into().unwrap());
         Ok(Self {
             kind,
             flags,
             conn_id,
             seq,
             ack,
+            session_key_id,
             body_len,
         })
     }
@@ -147,6 +153,7 @@ pub fn encode_plain(
         conn_id,
         seq,
         ack,
+        session_key_id: [0u8; 16],
         body_len: body.len() as u16,
     };
     let mut out = Vec::with_capacity(HEADER_LEN + body.len());
@@ -171,10 +178,11 @@ pub fn encode_encrypted(
         conn_id,
         seq,
         ack,
+        session_key_id: session_key_id(key),
         body_len: 0,
     };
     let aad_without_len = header.aad();
-    let ciphertext = crypto::seal(key, &nonce, &aad_without_len[..40], plaintext)?;
+    let ciphertext = crypto::seal(key, &nonce, &aad_without_len[..HEADER_AAD_LEN], plaintext)?;
     if ciphertext.len() > u16::MAX as usize {
         bail!("packet body too large");
     }
@@ -206,7 +214,18 @@ pub fn decrypt_body(packet: &Packet, key: &[u8; 32], direction: u32) -> Result<V
     }
     let nonce = crypto::nonce_from(direction, packet.header.seq);
     let aad = packet.header.aad();
-    crypto::open(key, &nonce, &aad[..40], &packet.body)
+    let expected_key_id = session_key_id(key);
+    if packet.header.session_key_id != expected_key_id {
+        bail!("stale or wrong session key id");
+    }
+    crypto::open(key, &nonce, &aad[..HEADER_AAD_LEN], &packet.body)
+}
+
+pub fn session_key_id(key: &[u8; 32]) -> [u8; 16] {
+    let digest = crypto::sha256(key);
+    let mut out = [0u8; 16];
+    out.copy_from_slice(&digest[..16]);
+    out
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

tests/protocol_auth.rs

@@ -50,10 +50,33 @@ fn encrypted_packet_round_trips() {
     let decoded = protocol::decode(&packet).unwrap();
     assert_eq!(decoded.header.kind, PacketKind::Input);
     assert_eq!(decoded.header.conn_id, conn_id);
+    assert_eq!(
+        decoded.header.session_key_id,
+        protocol::session_key_id(&key)
+    );
     let plain = protocol::decrypt_body(&decoded, &key, CLIENT_TO_SERVER).unwrap();
     assert_eq!(plain, b"hello");
 }
 
+#[test]
+fn encrypted_packet_rejects_wrong_session_key_id_before_decrypt() {
+    let key = crypto::random_32();
+    let wrong_key = crypto::random_32();
+    let packet = protocol::encode_encrypted(
+        PacketKind::Input,
+        crypto::random_16(),
+        1,
+        0,
+        &key,
+        CLIENT_TO_SERVER,
+        b"hello",
+    )
+    .unwrap();
+    let decoded = protocol::decode(&packet).unwrap();
+    let err = protocol::decrypt_body(&decoded, &wrong_key, CLIENT_TO_SERVER).unwrap_err();
+    assert!(err.to_string().contains("session key id"));
+}
+
 #[test]
 fn attach_ticket_is_sealed_and_verifies_scope() {
     let dir = tempfile::tempdir().unwrap();