From a30bef11dab76433e129047d0481de8689284b62 Mon Sep 17 00:00:00 2001
From: CleverWild <cleverwilddev@gmail.com>
Date: Tue, 30 Jun 2026 20:07:01 +0200
Subject: [PATCH] fix: terminate VaultGate connection after vault lockout

---
 .../src/peers/operator/vault_gate/mod.rs      | 16 +++++---
 .../arbiter-server/tests/operator/unseal.rs   | 41 +++++++++++++++++++
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/server/crates/arbiter-server/src/peers/operator/vault_gate/mod.rs b/server/crates/arbiter-server/src/peers/operator/vault_gate/mod.rs
index d0e38e3..1320972 100644
--- a/server/crates/arbiter-server/src/peers/operator/vault_gate/mod.rs
+++ b/server/crates/arbiter-server/src/peers/operator/vault_gate/mod.rs
@@ -11,7 +11,7 @@ use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
 use state::State;
 
 use chacha20poly1305::{AeadInPlace, KeyInit as _, XChaCha20Poly1305, XNonce};
-use kameo::{Actor, error::SendError, messages, prelude::Message};
+use kameo::{Actor, error::SendError, messages, prelude::{Context, Message}};
 use kameo_actors::message_bus::Register;
 use tokio::sync::oneshot;
 use tracing::{error, info};
@@ -143,12 +143,13 @@ impl VaultGate {
         })
     }
 
-    #[message]
+    #[message(ctx)]
     pub async fn handle_unseal_encrypted_key(
         &mut self,
         nonce: Vec<u8>,
         ciphertext: Vec<u8>,
         associated_data: Vec<u8>,
+        ctx: &mut Context<Self, Result<(), Error>>,
     ) -> Result<(), Error> {
         let State::ReadyForExchange { secret, .. } = &self.state else {
             return Err(Error::State);
@@ -172,7 +173,12 @@ impl VaultGate {
                 Ok(())
             }
             Err(SendError::HandlerError(vault::Error::InvalidKey)) => Err(Error::InvalidKey),
-            Err(SendError::HandlerError(vault::Error::LockedOut)) => Err(Error::LockedOut),
+            Err(SendError::HandlerError(vault::Error::LockedOut)) => {
+                // Vault is permanently locked — terminate this gate so the
+                // run_vault_gate loop breaks and the connection is closed.
+                ctx.stop();
+                Err(Error::LockedOut)
+            }
             Err(SendError::HandlerError(err)) => {
                 error!(?err, "Vault failed to unseal key");
                 Err(Error::InvalidKey)
@@ -245,7 +251,7 @@ impl Message<events::Bootstrapped> for VaultGate {
     async fn handle(
         &mut self,
         _: events::Bootstrapped,
-        ctx: &mut kameo::prelude::Context<Self, Self::Reply>,
+        ctx: &mut Context<Self, Self::Reply>,
     ) -> Self::Reply {
         let result = async {
             let mut conn = self
@@ -281,7 +287,7 @@ impl Message<events::Unsealed> for VaultGate {
     async fn handle(
         &mut self,
         _: events::Unsealed,
-        ctx: &mut kameo::prelude::Context<Self, Self::Reply>,
+        ctx: &mut Context<Self, Self::Reply>,
     ) -> Self::Reply {
         if let Some(tx) = self.promotion_tx.take() {
             let _ = tx.send(Ok(()));
diff --git a/server/crates/arbiter-server/tests/operator/unseal.rs b/server/crates/arbiter-server/tests/operator/unseal.rs
index 365e6ec..37710fb 100644
--- a/server/crates/arbiter-server/tests/operator/unseal.rs
+++ b/server/crates/arbiter-server/tests/operator/unseal.rs
@@ -140,6 +140,47 @@ pub async fn unseal_corrupted_ciphertext() {
     ));
 }
 
+/// After MAX_UNSEAL_ATTEMPTS wrong keys, the vault locks and the VaultGate
+/// must stop itself so the connection is dropped.
+#[tokio::test]
+#[test_log::test]
+pub async fn lockout_stops_vault_gate() {
+    use kameo::error::SendError;
+
+    let seal_key = b"real-seal-key";
+    let (_db, gate, _promotion_rx) = setup_sealed_gate(seal_key).await;
+
+    // Exhaust all MAX_UNSEAL_ATTEMPTS (5) with wrong keys; each returns InvalidKey.
+    for _ in 0..5 {
+        let encrypted_key = client_dh_encrypt(&gate, b"wrong-key").await;
+        assert!(matches!(
+            gate.ask(encrypted_key).await,
+            Err(SendError::HandlerError(VaultGateError::InvalidKey))
+        ));
+    }
+
+    // Sixth attempt: vault is now locked, returns LockedOut, and the gate stops itself.
+    let encrypted_key = client_dh_encrypt(&gate, b"wrong-key").await;
+    assert!(matches!(
+        gate.ask(encrypted_key).await,
+        Err(SendError::HandlerError(VaultGateError::LockedOut))
+    ));
+
+    // Give the actor scheduler time to process the stop signal.
+    tokio::time::sleep(std::time::Duration::from_millis(50)).await;
+
+    // Any subsequent message must be rejected because the gate is stopped.
+    let client_secret = EphemeralSecret::random();
+    let client_public = PublicKey::from(&client_secret);
+    assert!(
+        matches!(
+            gate.ask(HandleHandshake { client_pubkey: client_public }).await,
+            Err(SendError::ActorNotRunning(_))
+        ),
+        "VaultGate must be stopped after LockedOut"
+    );
+}
+
 #[tokio::test]
 #[test_log::test]
 pub async fn unseal_retry_after_invalid_key() {