feat: implement challenge-response mechanism for session closure

client/src/crypto.rs

@@ -16,6 +16,12 @@ pub(crate) fn derive_key(shared_secret: &[u8], label: &[u8]) -> Vec<u8> {
     hmac::sign(&key, label).as_ref().to_vec()
 }
 
+/// Compute HMAC-SHA256(key, data) and return the 32-byte tag.
+pub(crate) fn hmac_sha256(key_bytes: &[u8], data: &[u8]) -> Vec<u8> {
+    let key = hmac::Key::new(hmac::HMAC_SHA256, key_bytes);
+    hmac::sign(&key, data).as_ref().to_vec()
+}
+
 pub(crate) fn b64_encode(data: &[u8]) -> String {
     B64.encode(data)
 }

client/src/main.rs

@@ -185,9 +185,25 @@ fn main() -> Result<()> {
     );
     println!("decrypted sum={sum}");
 
-    // 6. Close
+    // 6. Close (challenge-response)
+    // 6a. Request close challenge
+    let challenge_req = types::Request::CloseChallenge {
+        session_id: session_id.clone(),
+    }
+    .to_json()?;
+    let challenge_resp = post_json(&server, &challenge_req)?;
+    let challenge_resp: types::Response = types::Response::from_json(&challenge_resp)?;
+    let challenge = match challenge_resp {
+        types::Response::CloseChallenge { challenge_b64 } => crypto::b64_decode(&challenge_b64)?,
+        types::Response::Error { error } => return Err(anyhow!("close-challenge failed: {error}")),
+        other => return Err(anyhow!("unexpected close-challenge response: {other:?}")),
+    };
+
+    // 6b. Compute response = HMAC-SHA256(SK, challenge) and send close
+    let response = crypto::hmac_sha256(&sk, &challenge);
     let close_req = types::Request::Close {
         session_id: session_id.clone(),
+        response_b64: crypto::b64_encode(&response),
     }
     .to_json()?;
     let close_resp = post_json(&server, &close_req)?;

client/src/types.rs

@@ -24,8 +24,12 @@ pub enum Request {
         x: EncryptedBlob,
         y: EncryptedBlob,
     },
+    CloseChallenge {
+        session_id: String,
+    },
     Close {
         session_id: String,
+        response_b64: String,
     },
     Attest {
         session_id: Option<String>,
@@ -47,6 +51,9 @@ pub enum Response {
     Add {
         sum: EncryptedBlob,
     },
+    CloseChallenge {
+        challenge_b64: String,
+    },
     CloseOk {},
     Attest {
         attestation_document_b64: String,

docs/architecture.md

@@ -224,7 +224,8 @@ representation (`#[serde(tag = "type")]`).
 | `init` | *(none)* | Start a new session |
 | `key-exchange` | `session_id`, `client_pubkey_b64` | Complete ECDH and request attestation |
 | `add` | `session_id`, `x: EncryptedBlob`, `y: EncryptedBlob` | Addition |
-| `close` | `session_id` | Close session and wipe keys |
+| `close-challenge` | `session_id` | Request a close challenge (step 1) |
+| `close` | `session_id`, `response_b64` | Close session with challenge response (step 2) |
 | `attest` | `session_id?`, `user_data_b64?`, `nonce_b64?` | Request an attestation document on demand |
 
 #### Responses (Enclave → Client)
@@ -234,6 +235,7 @@ representation (`#[serde(tag = "type")]`).
 | `init` | `session_id`, `enclave_pubkey_b64` | Session created |
 | `key-exchange` | `attestation_document_b64` | Attestation document (COSE_Sign1, base64) |
 | `add` | `sum: EncryptedBlob` | Encrypted result |
+| `close-challenge` | `challenge_b64` | 32-byte random challenge (base64) |
 | `close-ok` | *(none)* | Session closed |
 | `attest` | `attestation_document_b64` | Attestation document |
 | `error` | `error` | Error message string |
@@ -310,10 +312,20 @@ sequenceDiagram
     P-->>C: HTTP 200 (same JSON)
     Note over C: Decrypt sum with MK
 
-    Note over C,E: Phase 5 — Session Close
+    Note over C,E: Phase 5 — Session Close (Challenge-Response)
 
-    C->>P: POST / {"type":"close", "session_id":"..."}
+    C->>P: POST / {"type":"close-challenge", "session_id":"..."}
     P->>E: vsock: (same JSON)
+    E->>N: GetRandom (32 bytes for challenge)
+    N-->>E: random bytes
+    Note over E: Store challenge in session
+    E-->>P: {"type":"close-challenge", "challenge_b64":"..."}
+    P-->>C: HTTP 200 (same JSON)
+
+    Note over C: Compute response =<br/>HMAC-SHA256(SK, challenge)
+    C->>P: POST / {"type":"close", "session_id":"...", "response_b64":"..."}
+    P->>E: vsock: (same JSON)
+    Note over E: Verify HMAC-SHA256(SK, challenge) == response
     Note over E: Delete session (wipe keys)
     E-->>P: {"type":"close-ok"}
     P-->>C: HTTP 200 (same JSON)
@@ -504,9 +516,15 @@ Behavior depends on which optional fields are provided:
 4. **Output confidentiality & authenticity** — Enclave outputs are returned as
    AES-128-GCM (MK). A network attacker (or the untrusted proxy) cannot modify
    ciphertexts without detection by the client.
-5. **Forward secrecy** — Ephemeral ECDH keypairs are generated per session;
+5. **Session close authentication** — Session close uses a challenge-response
+   protocol: the enclave issues a 32-byte random challenge, and the client must
+   respond with `HMAC-SHA256(SK, challenge)`. Since SK is derived from the ECDH
+   shared secret, only the legitimate client (who participated in the key
+   exchange) can compute the correct response. This prevents an attacker who
+   has observed the session ID from forcibly closing the session (DoS).
+6. **Forward secrecy** — Ephemeral ECDH keypairs are generated per session;
    compromising one session does not affect others.
-6. **Hardware-rooted trust** — The attestation document is signed by the Nitro
+7. **Hardware-rooted trust** — The attestation document is signed by the Nitro
    Hypervisor via the Nitro Secure Module, whose certificate chain roots to the
    AWS Nitro Enclaves root CA.
    The Nitro hypervisor is secure-booted with the aid of the hardware

enclave/src/crypto.rs

@@ -16,6 +16,12 @@ pub(crate) fn derive_key(shared_secret: &[u8], label: &[u8]) -> Vec<u8> {
     hmac::sign(&key, label).as_ref().to_vec()
 }
 
+/// Compute HMAC-SHA256(key, data) and return the 32-byte tag.
+pub(crate) fn hmac_sha256(key_bytes: &[u8], data: &[u8]) -> Vec<u8> {
+    let key = hmac::Key::new(hmac::HMAC_SHA256, key_bytes);
+    hmac::sign(&key, data).as_ref().to_vec()
+}
+
 /// AES-128-GCM encrypt.
 pub(crate) fn aes128gcm_encrypt(
     key_16: &[u8],

enclave/src/main.rs

@@ -213,7 +213,42 @@ pub(crate) fn handle_vsock_stream(stream: &mut VsockStream, max_request_size: us
                 let response_str = response.to_json()?;
                 write_response(stream, response_str.as_bytes())?;
             }
-            types::Request::Close { session_id } => {
+            types::Request::CloseChallenge { session_id } => {
+                let sess = session::get_session(&session_id)?;
+                // Require key exchange to have completed (SK must exist)
+                if sess.sk.is_none() {
+                    return Err(anyhow::anyhow!("key exchange not completed"));
+                }
+                let challenge = session::set_close_challenge(&session_id)?;
+                let response = types::Response::CloseChallenge {
+                    challenge_b64: crypto::b64_encode(&challenge),
+                };
+                let response_str = response.to_json()?;
+                write_response(stream, response_str.as_bytes())?;
+            }
+            types::Request::Close {
+                session_id,
+                response_b64,
+            } => {
+                let sess = session::get_session(&session_id)?;
+                let sk = sess
+                    .sk
+                    .as_ref()
+                    .ok_or_else(|| anyhow::anyhow!("SK not set"))?;
+                let challenge = sess
+                    .close_challenge
+                    .as_ref()
+                    .ok_or_else(|| anyhow::anyhow!("no pending close challenge"))?;
+
+                // Verify: response == HMAC-SHA256(SK, challenge)
+                let expected = crypto::hmac_sha256(sk, challenge);
+                let actual = crypto::b64_decode(&response_b64)?;
+                if expected != actual {
+                    return Err(anyhow::anyhow!(
+                        "close challenge-response verification failed"
+                    ));
+                }
+
                 session::delete_session(&session_id)?;
                 let response = types::Response::CloseOk {};
                 let response_str = response.to_json()?;

enclave/src/session.rs

@@ -23,6 +23,9 @@ pub(crate) struct Session {
     pub sk: Option<Vec<u8>>,
     pub mk: Option<Vec<u8>>,
     pub vk: Option<Vec<u8>>,
+
+    /// 32-byte close challenge (set by CloseChallenge, consumed by Close).
+    pub close_challenge: Option<Vec<u8>>,
 }
 
 static SESSIONS: OnceLock<Mutex<HashMap<String, Session>>> = OnceLock::new();
@@ -74,6 +77,7 @@ pub(crate) fn new_session() -> Result<Session> {
         sk: None,
         mk: None,
         vk: None,
+        close_challenge: None,
     })
 }
 
@@ -107,6 +111,7 @@ pub(crate) fn get_session(session_id: &str) -> Result<Session> {
             sk: s.sk.clone(),
             mk: s.mk.clone(),
             vk: s.vk.clone(),
+            close_challenge: s.close_challenge.clone(),
         })
         .ok_or_else(|| anyhow!("unknown session_id"))
 }
@@ -119,6 +124,19 @@ pub(crate) fn put_session(sess: Session) -> Result<()> {
     Ok(())
 }
 
+/// Generate and store a 32-byte close challenge for the session.
+pub(crate) fn set_close_challenge(session_id: &str) -> Result<Vec<u8>> {
+    let challenge = nsm_random_bytes(32)?;
+    let mut map = sessions()
+        .lock()
+        .map_err(|_| anyhow!("session mutex poisoned"))?;
+    let sess = map
+        .get_mut(session_id)
+        .ok_or_else(|| anyhow!("unknown session_id"))?;
+    sess.close_challenge = Some(challenge.clone());
+    Ok(challenge)
+}
+
 pub(crate) fn delete_session(session_id: &str) -> Result<()> {
     let mut map = sessions()
         .lock()

enclave/src/types.rs

@@ -26,8 +26,14 @@ pub enum Request {
         y: EncryptedBlob,
     },
 
-    /// Close the session and wipe keys.
-    Close { session_id: String },
+    /// Request a close challenge (step 1 of challenge-response close).
+    CloseChallenge { session_id: String },
+
+    /// Close the session with a challenge response (step 2).
+    Close {
+        session_id: String,
+        response_b64: String,
+    },
 
     /// Request an attestation document anytime.
     ///
@@ -60,6 +66,9 @@ pub enum Response {
     Add {
         sum: EncryptedBlob,
     },
+    CloseChallenge {
+        challenge_b64: String,
+    },
     CloseOk {},
     Attest {
         attestation_document_b64: String,