feat(user-agent): add VaultGate for sealed vault authentication

server/Cargo.lock

@@ -700,6 +700,7 @@ dependencies = [
  "memsafe",
  "ml-dsa",
  "rand 0.10.0",
+ "x-wing",
 ]
 
 [[package]]
@@ -779,7 +780,7 @@ dependencies = [
  "tonic",
  "tracing",
  "tracing-subscriber",
- "x25519-dalek",
+ "x25519-dalek 2.0.1",
  "zeroize",
 ]
 
@@ -1621,6 +1622,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "77727bb15fa921304124b128af125e7e3b968275d1b108b379190264f4423710"
 dependencies = [
  "hybrid-array",
+ "rand_core 0.10.0",
 ]
 
 [[package]]
@@ -2598,6 +2600,7 @@ version = "0.4.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8655f91cd07f2b9d0c24137bd650fe69617773435ee5ec83022377777ce65ef1"
 dependencies = [
+ "ctutils",
  "typenum",
  "zeroize",
 ]
@@ -3030,6 +3033,16 @@ dependencies = [
  "sha3-asm",
 ]
 
+[[package]]
+name = "kem"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "01737161ba802849cfd486b5bd209d38ba4943494c249a8126005170c7621edd"
+dependencies = [
+ "crypto-common 0.2.1",
+ "rand_core 0.10.0",
+]
+
 [[package]]
 name = "lazy_static"
 version = "1.5.0"
@@ -3250,12 +3263,27 @@ dependencies = [
  "zeroize",
 ]
 
+[[package]]
+name = "ml-kem"
+version = "0.3.0-rc.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "04437cb1a66c0b78740927b76cc61f218344b9f6ef3dd430e283274a718ef0e9"
+dependencies = [
+ "hybrid-array",
+ "kem",
+ "module-lattice",
+ "rand_core 0.10.0",
+ "sha3 0.11.0",
+ "zeroize",
+]
+
 [[package]]
 name = "module-lattice"
 version = "0.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "164eb3faeaecbd14b0b2a917c1b4d0c035097a9c559b0bed85c2cdd032bc8faa"
 dependencies = [
+ "ctutils",
  "hybrid-array",
  "num-traits",
  "zeroize",
@@ -6105,6 +6133,20 @@ dependencies = [
  "tap",
 ]
 
+[[package]]
+name = "x-wing"
+version = "0.1.0-rc.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e17d0d5f4d1f26b9b9e7477af1d3bef960e1d1fb64edab7912fde472a8a8432e"
+dependencies = [
+ "kem",
+ "ml-kem",
+ "rand_core 0.10.0",
+ "sha3 0.11.0",
+ "x25519-dalek 3.0.0-pre.6",
+ "zeroize",
+]
+
 [[package]]
 name = "x25519-dalek"
 version = "2.0.1"
@@ -6117,6 +6159,17 @@ dependencies = [
  "zeroize",
 ]
 
+[[package]]
+name = "x25519-dalek"
+version = "3.0.0-pre.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3d5d6ff67acd3945b933e592bfa7143db4fcbb2f871754b6b9fbd7847fc5aea"
+dependencies = [
+ "curve25519-dalek 5.0.0-pre.6",
+ "rand_core 0.10.0",
+ "zeroize",
+]
+
 [[package]]
 name = "x509-parser"
 version = "0.18.1"

server/crates/arbiter-crypto/Cargo.toml

@@ -8,6 +8,7 @@ ml-dsa = {workspace = true, optional = true }
 rand = {workspace = true, optional = true}
 base64 = {workspace = true, optional = true }
 memsafe = {version = "0.4.0", optional = true}
+x-wing = { version = "0.1.0-rc.0", features = ["zeroize"] }
 
 [lints]
 workspace = true
@@ -15,4 +16,4 @@ workspace = true
 [features]
 default = ["authn", "safecell"]
 authn = ["dep:ml-dsa", "dep:rand", "dep:base64"]
-safecell = ["dep:memsafe"]
+safecell = ["dep:memsafe"]

server/crates/arbiter-crypto/src/authn/v1.rs

@@ -1,3 +1,5 @@
+use std::hash::Hash;
+
 use base64::{Engine as _, prelude::BASE64_STANDARD};
 use ml_dsa::{
     EncodedVerifyingKey, Error, KeyGen, MlDsa87, Seed, Signature as MlDsaSignature,
@@ -17,6 +19,12 @@ pub type KeyParams = MlDsa87;
 #[derive(Clone, Debug, PartialEq)]
 pub struct PublicKey(Box<MlDsaVerifyingKey<KeyParams>>);
 
+impl Hash for PublicKey {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.to_bytes().hash(state);
+    }
+}
+
 #[derive(Clone, Debug, PartialEq)]
 pub struct Signature(Box<MlDsaSignature<KeyParams>>);
 

server/crates/arbiter-crypto/src/lib.rs

@@ -3,3 +3,5 @@ pub mod authn;
 
 #[cfg(feature = "safecell")]
 pub mod safecell;
+
+pub use x_wing;

server/crates/arbiter-server/src/actors/mod.rs

@@ -13,6 +13,7 @@ pub mod bootstrap;
 pub mod evm;
 pub mod flow_coordinator;
 pub mod vault;
+pub mod useragent_registry;
 
 #[derive(Error, Debug)]
 pub enum SpawnError {

server/crates/arbiter-server/src/actors/useragent_registry.rs

@@ -0,0 +1,57 @@
+use alloy::primitives::map::HashMap;
+use arbiter_crypto::authn;
+use kameo::{error::Infallible, prelude::*};
+
+use crate::{db::DatabasePool, peers::user_agent::{Credentials, UserAgentSession}};
+
+use super::vault::{Vault, events as vault_events};
+
+pub struct Args {
+    pub vault: ActorRef<Vault>,
+    pub pool: DatabasePool,
+}
+
+pub struct UserAgentRegistry {
+    vault: ActorRef<Vault>,
+    pool: DatabasePool,
+    connected: HashMap<Credentials, ActorRef<UserAgentSession>>,
+}
+
+impl Message<vault_events::Bootstrapped> for UserAgentRegistry {
+    type Reply = ();
+
+    async fn handle(
+        &mut self,
+        msg: vault_events::Bootstrapped,
+        ctx: &mut Context<Self, Self::Reply>,
+    ) -> Self::Reply {
+        todo!()
+    }
+}
+
+impl Message<vault_events::Unsealed> for UserAgentRegistry {
+    type Reply = ();
+
+    async fn handle(
+        &mut self,
+        msg: vault_events::Unsealed,
+        ctx: &mut Context<Self, Self::Reply>,
+    ) -> Self::Reply {
+        todo!()
+    }
+}
+impl Actor for UserAgentRegistry {
+    type Args = Args;
+
+    type Error = Infallible;
+
+    async fn on_start(args: Self::Args, actor_ref: ActorRef<Self>) -> Result<Self, Self::Error> {
+        Ok(Self {
+            vault: args.vault,
+            pool: args.pool,
+            connected: HashMap::default(),
+        })
+    }
+
+    
+}

server/crates/arbiter-server/src/actors/vault/mod.rs

@@ -25,10 +25,10 @@ use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
 pub mod events {
 
     #[derive(Clone, Copy)]
-    pub struct VaultBootstrapped;
+    pub struct Bootstrapped;
 
     #[derive(Clone, Copy)]
-    pub struct VaultUnsealed;
+    pub struct Unsealed;
 
     #[derive(Clone, Copy)]
     pub struct VaultResealed;
@@ -213,7 +213,7 @@ impl Vault {
         });
 
         info!("Vault bootstrapped successfully");
-        self.events.tell(Publish(events::VaultBootstrapped)).await;
+        self.events.tell(Publish(events::Bootstrapped)).await;
 
         Ok(())
     }
@@ -269,7 +269,7 @@ impl Vault {
         });
 
         info!("Vault unsealed successfully");
-        self.events.tell(Publish(events::VaultUnsealed)).await;
+        self.events.tell(Publish(events::Unsealed)).await;
 
         Ok(())
     }

server/crates/arbiter-server/src/grpc/mod.rs

@@ -10,7 +10,6 @@ use tonic::{Request, Response, Status, async_trait};
 use tracing::info;
 
 use crate::{
-    grpc::user_agent::start,
     peers::{client::ClientConnection, user_agent::UserAgentConnection},
 };
 
@@ -63,7 +62,7 @@ impl arbiter_proto::proto::arbiter_service_server::ArbiterService for super::Ser
 
         let (bi, rx) = GrpcBi::from_bi_stream(req_stream);
 
-        tokio::spawn(start(
+        tokio::spawn(user_agent::start(
             UserAgentConnection {
                 db: self.context.db.clone(),
                 actors: self.context.actors.clone(),

server/crates/arbiter-server/src/grpc/user_agent.rs

@@ -1,4 +1,4 @@
-use tokio::sync::mpsc;
+use tokio::sync::{mpsc, oneshot};
 
 use arbiter_proto::{
     proto::user_agent::{
@@ -14,8 +14,12 @@ use tonic::Status;
 use tracing::{error, info, warn};
 
 use crate::{
+    crypto::integrity,
     grpc::request_tracker::RequestTracker,
-    peers::user_agent::{OutOfBand, UserAgentConnection, UserAgentSession},
+    peers::user_agent::{
+        Credentials, OutOfBand, UserAgentConnection, UserAgentSession,
+        vault_gate::VaultGate,
+    },
 };
 
 mod auth;
@@ -24,6 +28,7 @@ mod inbound;
 mod outbound;
 mod sdk_client;
 mod vault;
+mod vault_gate;
 
 pub struct OutOfBandAdapter(mpsc::Sender<OutOfBand>);
 
@@ -124,27 +129,115 @@ pub async fn start(
 ) {
     let mut request_tracker = RequestTracker::default();
 
-    let (id, pubkey) = match auth::start(&mut conn, &mut bi, &mut request_tracker).await {
+    let auth_creds = match auth::start(&mut conn, &mut bi, &mut request_tracker).await {
-        Ok(pubkey) => pubkey,
+        Ok(creds) => creds,
         Err(e) => {
             warn!(error = ?e, "Authentication failed");
             return;
         }
     };
 
-    info!(?pubkey, "User authenticated successfully");
+    info!(pubkey = ?auth_creds.creds.pubkey, "User authenticated successfully");
+
+    let creds = if integrity::is_signing_available(&conn.actors.vault)
+        .await
+        .unwrap_or(false)
+    {
+        // Vault is unsealed; integrity was verified during auth — promote directly.
+        auth_creds.creds
+    } else {
+        // Vault is sealed/unbootstrapped; run the VaultGate phase.
+        let (promotion_tx, promotion_rx) = oneshot::channel();
+        let gate = VaultGate::spawn(VaultGate::new(
+            auth_creds,
+            conn.actors.clone(),
+            conn.db.clone(),
+            promotion_tx,
+        ));
+
+        let result = vault_gate_loop(&mut bi, &gate, &mut request_tracker, promotion_rx).await;
+        gate.kill();
+
+        match result {
+            Some(creds) => creds,
+            None => return,
+        }
+    };
 
     let (oob_sender, oob_receiver) = mpsc::channel(16);
     let oob_adapter = OutOfBandAdapter(oob_sender);
 
-    let actor = UserAgentSession::spawn(UserAgentSession::new(
+    let actor = UserAgentSession::spawn(UserAgentSession::new(conn, creds, Box::new(oob_adapter)));
-        conn,
-        id,
-        pubkey,
-        Box::new(oob_adapter),
-    ));
     let actor_for_cleanup = actor.clone();
 
     dispatch_loop(bi, actor, oob_receiver, request_tracker).await;
     actor_for_cleanup.kill();
 }
+
+async fn vault_gate_loop(
+    bi: &mut GrpcBi<UserAgentRequest, UserAgentResponse>,
+    gate: &ActorRef<VaultGate>,
+    request_tracker: &mut RequestTracker,
+    mut promotion_rx: oneshot::Receiver<Result<Credentials, crate::peers::user_agent::vault_gate::Error>>,
+) -> Option<Credentials> {
+    loop {
+        tokio::select! {
+            result = &mut promotion_rx => {
+                return match result {
+                    Ok(Ok(creds)) => Some(creds),
+                    Ok(Err(e)) => {
+                        warn!(error = ?e, "VaultGate promotion failed");
+                        None
+                    }
+                    Err(_) => {
+                        warn!("VaultGate promotion channel closed unexpectedly");
+                        None
+                    }
+                };
+            }
+
+            message = bi.recv() => {
+                let Some(message) = message else { return None; };
+
+                let conn = match message {
+                    Ok(conn) => conn,
+                    Err(err) => {
+                        warn!(error = ?err, "Failed to receive request during vault gate phase");
+                        return None;
+                    }
+                };
+
+                let request_id = match request_tracker.request(conn.id) {
+                    Ok(id) => id,
+                    Err(err) => {
+                        let _ = bi.send(Err(err)).await;
+                        return None;
+                    }
+                };
+
+                let Some(payload) = conn.payload else {
+                    let _ = bi.send(Err(Status::invalid_argument("Missing request payload"))).await;
+                    return None;
+                };
+
+                let response = match payload {
+                    UserAgentRequestPayload::Vault(req) => vault_gate::dispatch(gate, req).await,
+                    _ => Err(Status::permission_denied("Only vault operations are permitted before unsealing")),
+                };
+
+                match response {
+                    Ok(Some(payload)) => {
+                        if bi.send(Ok(UserAgentResponse { id: Some(request_id), payload: Some(payload) })).await.is_err() {
+                            return None;
+                        }
+                    }
+                    Ok(None) => {}
+                    Err(status) => {
+                        let _ = bi.send(Err(status)).await;
+                        return None;
+                    }
+                }
+            }
+        }
+    }
+}

server/crates/arbiter-server/src/grpc/user_agent/auth.rs

@@ -19,7 +19,7 @@ use tracing::warn;
 
 use crate::{
     grpc::request_tracker::RequestTracker,
-    peers::user_agent::{UserAgentConnection, auth},
+    peers::user_agent::{AuthCredentials, UserAgentConnection, auth},
 };
 
 pub struct AuthTransportAdapter<'a> {
@@ -167,7 +167,7 @@ pub async fn start(
     conn: &mut UserAgentConnection,
     bi: &mut GrpcBi<UserAgentRequest, UserAgentResponse>,
     request_tracker: &mut RequestTracker,
-) -> Result<(i32, authn::PublicKey), auth::Error> {
+) -> Result<AuthCredentials, auth::Error> {
     let transport = AuthTransportAdapter::new(bi, request_tracker);
     auth::authenticate(conn, transport).await
 }

server/crates/arbiter-server/src/grpc/user_agent/evm.rs

@@ -29,7 +29,7 @@ use crate::{
     },
     peers::user_agent::{
         UserAgentSession,
-        session::connection::{
+        session::handlers::{
             GrantMutationError, HandleEvmWalletCreate, HandleEvmWalletList, HandleGrantCreate,
             HandleGrantDelete, HandleGrantList, HandleSignTransaction,
             SignTransactionError as SessionSignTransactionError,

server/crates/arbiter-server/src/grpc/user_agent/sdk_client.rs

@@ -25,7 +25,7 @@ use crate::{
     grpc::Convert,
     peers::user_agent::{
         OutOfBand, UserAgentSession,
-        session::connection::{
+        session::handlers::{
             HandleGrantEvmWalletAccess, HandleListWalletAccess, HandleNewClientApprove,
             HandleRevokeEvmWalletAccess, HandleSdkClientList,
         },

server/crates/arbiter-server/src/grpc/user_agent/vault.rs

@@ -1,34 +1,15 @@
 use arbiter_proto::proto::shared::VaultState as ProtoVaultState;
 use arbiter_proto::proto::user_agent::{
     user_agent_response::Payload as UserAgentResponsePayload,
-    vault::{
+    vault::{self as proto_vault, request::Payload as VaultRequestPayload, response::Payload as VaultResponsePayload},
-        self as proto_vault,
-        bootstrap::{
-            self as proto_bootstrap, BootstrapEncryptedKey as ProtoBootstrapEncryptedKey,
-            BootstrapResult as ProtoBootstrapResult,
-        },
-        request::Payload as VaultRequestPayload,
-        response::Payload as VaultResponsePayload,
-        unseal::{
-            self as proto_unseal, UnsealEncryptedKey as ProtoUnsealEncryptedKey,
-            UnsealResult as ProtoUnsealResult, UnsealStart,
-            request::Payload as UnsealRequestPayload, response::Payload as UnsealResponsePayload,
-        },
-    },
 };
-use kameo::{actor::ActorRef, error::SendError};
+use kameo::actor::ActorRef;
 use tonic::Status;
 use tracing::warn;
 
 use crate::{
     actors::vault::VaultState,
-    peers::user_agent::{
+    peers::user_agent::{UserAgentSession, session::handlers::HandleQueryVaultState},
-        UserAgentSession,
-        session::connection::{
-            BootstrapError, HandleBootstrapEncryptedKey, HandleQueryVaultState,
-            HandleUnsealEncryptedKey, HandleUnsealRequest, UnsealError,
-        },
-    },
 };
 
 fn wrap_vault_response(payload: VaultResponsePayload) -> UserAgentResponsePayload {
@@ -37,18 +18,6 @@ fn wrap_vault_response(payload: VaultResponsePayload) -> UserAgentResponsePayloa
     })
 }
 
-fn wrap_unseal_response(payload: UnsealResponsePayload) -> UserAgentResponsePayload {
-    wrap_vault_response(VaultResponsePayload::Unseal(proto_unseal::Response {
-        payload: Some(payload),
-    }))
-}
-
-fn wrap_bootstrap_response(result: ProtoBootstrapResult) -> UserAgentResponsePayload {
-    wrap_vault_response(VaultResponsePayload::Bootstrap(proto_bootstrap::Response {
-        result: result.into(),
-    }))
-}
-
 pub(super) async fn dispatch(
     actor: &ActorRef<UserAgentSession>,
     req: proto_vault::Request,
@@ -59,109 +28,14 @@ pub(super) async fn dispatch(
 
     match payload {
         VaultRequestPayload::QueryState(_) => handle_query_vault_state(actor).await,
-        VaultRequestPayload::Unseal(req) => dispatch_unseal_request(actor, req).await,
+        VaultRequestPayload::Unseal(_) | VaultRequestPayload::Bootstrap(_) => {
-        VaultRequestPayload::Bootstrap(req) => handle_bootstrap_request(actor, req).await,
+            Err(Status::permission_denied(
+                "Vault is already unsealed; unseal/bootstrap not permitted in session",
+            ))
+        }
     }
 }
 
-async fn dispatch_unseal_request(
-    actor: &ActorRef<UserAgentSession>,
-    req: proto_unseal::Request,
-) -> Result<Option<UserAgentResponsePayload>, Status> {
-    let Some(payload) = req.payload else {
-        return Err(Status::invalid_argument("Missing unseal request payload"));
-    };
-
-    match payload {
-        UnsealRequestPayload::Start(req) => handle_unseal_start(actor, req).await,
-        UnsealRequestPayload::EncryptedKey(req) => handle_unseal_encrypted_key(actor, req).await,
-    }
-}
-
-async fn handle_unseal_start(
-    actor: &ActorRef<UserAgentSession>,
-    req: UnsealStart,
-) -> Result<Option<UserAgentResponsePayload>, Status> {
-    let client_pubkey = <[u8; 32]>::try_from(req.client_pubkey)
-        .map(x25519_dalek::PublicKey::from)
-        .map_err(|_| Status::invalid_argument("Invalid X25519 public key"))?;
-
-    let response = actor
-        .ask(HandleUnsealRequest { client_pubkey })
-        .await
-        .map_err(|err| {
-            warn!(error = ?err, "Failed to handle unseal start request");
-            Status::internal("Failed to start unseal flow")
-        })?;
-
-    Ok(Some(wrap_unseal_response(UnsealResponsePayload::Start(
-        proto_unseal::UnsealStartResponse {
-            server_pubkey: response.server_pubkey.as_bytes().to_vec(),
-        },
-    ))))
-}
-
-async fn handle_unseal_encrypted_key(
-    actor: &ActorRef<UserAgentSession>,
-    req: ProtoUnsealEncryptedKey,
-) -> Result<Option<UserAgentResponsePayload>, Status> {
-    let result = match actor
-        .ask(HandleUnsealEncryptedKey {
-            nonce: req.nonce,
-            ciphertext: req.ciphertext,
-            associated_data: req.associated_data,
-        })
-        .await
-    {
-        Ok(()) => ProtoUnsealResult::Success,
-        Err(SendError::HandlerError(UnsealError::InvalidKey)) => ProtoUnsealResult::InvalidKey,
-        Err(err) => {
-            warn!(error = ?err, "Failed to handle unseal request");
-            return Err(Status::internal("Failed to unseal vault"));
-        }
-    };
-    Ok(Some(wrap_unseal_response(UnsealResponsePayload::Result(
-        result.into(),
-    ))))
-}
-
-async fn handle_bootstrap_request(
-    actor: &ActorRef<UserAgentSession>,
-    req: proto_bootstrap::Request,
-) -> Result<Option<UserAgentResponsePayload>, Status> {
-    let encrypted_key = req
-        .encrypted_key
-        .ok_or_else(|| Status::invalid_argument("Missing bootstrap encrypted key"))?;
-    handle_bootstrap_encrypted_key(actor, encrypted_key).await
-}
-
-async fn handle_bootstrap_encrypted_key(
-    actor: &ActorRef<UserAgentSession>,
-    req: ProtoBootstrapEncryptedKey,
-) -> Result<Option<UserAgentResponsePayload>, Status> {
-    let result = match actor
-        .ask(HandleBootstrapEncryptedKey {
-            nonce: req.nonce,
-            ciphertext: req.ciphertext,
-            associated_data: req.associated_data,
-        })
-        .await
-    {
-        Ok(()) => ProtoBootstrapResult::Success,
-        Err(SendError::HandlerError(BootstrapError::InvalidKey)) => {
-            ProtoBootstrapResult::InvalidKey
-        }
-        Err(SendError::HandlerError(BootstrapError::AlreadyBootstrapped)) => {
-            ProtoBootstrapResult::AlreadyBootstrapped
-        }
-        Err(err) => {
-            warn!(error = ?err, "Failed to handle bootstrap request");
-            return Err(Status::internal("Failed to bootstrap vault"));
-        }
-    };
-    Ok(Some(wrap_bootstrap_response(result)))
-}
-
 async fn handle_query_vault_state(
     actor: &ActorRef<UserAgentSession>,
 ) -> Result<Option<UserAgentResponsePayload>, Status> {

server/crates/arbiter-server/src/grpc/user_agent/vault_gate.rs

@@ -0,0 +1,151 @@
+use arbiter_proto::proto::user_agent::{
+    user_agent_response::Payload as UserAgentResponsePayload,
+    vault::{
+        self as proto_vault,
+        bootstrap::{
+            self as proto_bootstrap, BootstrapResult as ProtoBootstrapResult,
+        },
+        request::Payload as VaultRequestPayload,
+        response::Payload as VaultResponsePayload,
+        unseal::{
+            self as proto_unseal, UnsealResult as ProtoUnsealResult, UnsealStart,
+            request::Payload as UnsealRequestPayload, response::Payload as UnsealResponsePayload,
+        },
+    },
+};
+use kameo::{actor::ActorRef, error::SendError};
+use tonic::Status;
+use tracing::warn;
+
+use crate::peers::user_agent::vault_gate::{
+    self as vault_gate, HandleBootstrapEncryptedKey, HandleHandshake, HandleUnsealEncryptedKey,
+    VaultGate,
+};
+
+fn wrap_vault_response(payload: VaultResponsePayload) -> UserAgentResponsePayload {
+    UserAgentResponsePayload::Vault(proto_vault::Response {
+        payload: Some(payload),
+    })
+}
+
+fn wrap_unseal_response(payload: UnsealResponsePayload) -> UserAgentResponsePayload {
+    wrap_vault_response(VaultResponsePayload::Unseal(proto_unseal::Response {
+        payload: Some(payload),
+    }))
+}
+
+fn wrap_bootstrap_response(result: ProtoBootstrapResult) -> UserAgentResponsePayload {
+    wrap_vault_response(VaultResponsePayload::Bootstrap(proto_bootstrap::Response {
+        result: result.into(),
+    }))
+}
+
+pub(super) async fn dispatch(
+    gate: &ActorRef<VaultGate>,
+    req: proto_vault::Request,
+) -> Result<Option<UserAgentResponsePayload>, Status> {
+    let Some(payload) = req.payload else {
+        return Err(Status::invalid_argument("Missing vault request payload"));
+    };
+
+    match payload {
+        VaultRequestPayload::QueryState(_) => {
+            use arbiter_proto::proto::shared::VaultState as ProtoVaultState;
+            Ok(Some(wrap_vault_response(VaultResponsePayload::State(
+                ProtoVaultState::Sealed.into(),
+            ))))
+        }
+        VaultRequestPayload::Unseal(req) => dispatch_unseal(gate, req).await,
+        VaultRequestPayload::Bootstrap(req) => dispatch_bootstrap(gate, req).await,
+    }
+}
+
+async fn dispatch_unseal(
+    gate: &ActorRef<VaultGate>,
+    req: proto_unseal::Request,
+) -> Result<Option<UserAgentResponsePayload>, Status> {
+    let Some(payload) = req.payload else {
+        return Err(Status::invalid_argument("Missing unseal request payload"));
+    };
+
+    match payload {
+        UnsealRequestPayload::Start(req) => handle_unseal_start(gate, req).await,
+        UnsealRequestPayload::EncryptedKey(req) => handle_unseal_encrypted_key(gate, req).await,
+    }
+}
+
+async fn handle_unseal_start(
+    gate: &ActorRef<VaultGate>,
+    req: UnsealStart,
+) -> Result<Option<UserAgentResponsePayload>, Status> {
+    let client_pubkey = <[u8; 32]>::try_from(req.client_pubkey)
+        .map(x25519_dalek::PublicKey::from)
+        .map_err(|_| Status::invalid_argument("Invalid X25519 public key"))?;
+
+    let response = gate
+        .ask(HandleHandshake { client_pubkey })
+        .await
+        .map_err(|err| {
+            warn!(error = ?err, "Failed to handle unseal start");
+            Status::internal("Failed to start unseal flow")
+        })?;
+
+    Ok(Some(wrap_unseal_response(UnsealResponsePayload::Start(
+        proto_unseal::UnsealStartResponse {
+            server_pubkey: response.server_pubkey.as_bytes().to_vec(),
+        },
+    ))))
+}
+
+async fn handle_unseal_encrypted_key(
+    gate: &ActorRef<VaultGate>,
+    req: arbiter_proto::proto::user_agent::vault::unseal::UnsealEncryptedKey,
+) -> Result<Option<UserAgentResponsePayload>, Status> {
+    let result = match gate
+        .ask(HandleUnsealEncryptedKey {
+            nonce: req.nonce,
+            ciphertext: req.ciphertext,
+            associated_data: req.associated_data,
+        })
+        .await
+    {
+        Ok(()) => ProtoUnsealResult::Success,
+        Err(SendError::HandlerError(vault_gate::Error::InvalidKey)) => ProtoUnsealResult::InvalidKey,
+        Err(err) => {
+            warn!(error = ?err, "Failed to handle unseal request");
+            return Err(Status::internal("Failed to unseal vault"));
+        }
+    };
+    Ok(Some(wrap_unseal_response(UnsealResponsePayload::Result(
+        result.into(),
+    ))))
+}
+
+async fn dispatch_bootstrap(
+    gate: &ActorRef<VaultGate>,
+    req: proto_bootstrap::Request,
+) -> Result<Option<UserAgentResponsePayload>, Status> {
+    let encrypted_key = req
+        .encrypted_key
+        .ok_or_else(|| Status::invalid_argument("Missing bootstrap encrypted key"))?;
+
+    let result = match gate
+        .ask(HandleBootstrapEncryptedKey {
+            nonce: encrypted_key.nonce,
+            ciphertext: encrypted_key.ciphertext,
+            associated_data: encrypted_key.associated_data,
+        })
+        .await
+    {
+        Ok(()) => ProtoBootstrapResult::Success,
+        Err(SendError::HandlerError(vault_gate::Error::InvalidKey)) => ProtoBootstrapResult::InvalidKey,
+        Err(SendError::HandlerError(vault_gate::Error::AlreadyBootstrapped)) => {
+            ProtoBootstrapResult::AlreadyBootstrapped
+        }
+        Err(err) => {
+            warn!(error = ?err, "Failed to handle bootstrap request");
+            return Err(Status::internal("Failed to bootstrap vault"));
+        }
+    };
+    Ok(Some(wrap_bootstrap_response(result)))
+}

server/crates/arbiter-server/src/peers/user_agent/auth/mod.rs

@@ -5,7 +5,7 @@ use tracing::error;
 mod state;
 use state::*;
 
-use super::UserAgentConnection;
+use super::{AuthCredentials, UserAgentConnection};
 
 #[derive(Debug, Clone)]
 pub enum Inbound {
@@ -69,7 +69,7 @@ fn parse_auth_event(payload: Inbound) -> AuthEvents {
 pub async fn authenticate<T>(
     props: &mut UserAgentConnection,
     transport: T,
-) -> Result<(i32, authn::PublicKey), Error>
+) -> Result<AuthCredentials, Error>
 where
     T: Bi<Inbound, Result<Outbound, Error>> + Send,
 {
@@ -82,7 +82,7 @@ where
         };
 
         match state.process_event(parse_auth_event(payload)).await {
-            Ok(AuthStates::AuthOk(result)) => return Ok((result.id, result.pubkey.clone())),
+            Ok(AuthStates::AuthOk(result)) => return Ok(result.clone()),
             Err(AuthError::ActionFailed(err)) => {
                 error!(?err, "State machine action failed");
                 return Err(err);

server/crates/arbiter-server/src/peers/user_agent/auth/state.rs

@@ -1,7 +1,7 @@
-use super::super::{UserAgentConnection, UserAgentCredentials};
+use super::super::{AuthCredentials, Credentials, UserAgentConnection};
 use arbiter_crypto::authn::{self, USERAGENT_CONTEXT};
 use arbiter_proto::transport::Bi;
-use diesel::{ExpressionMethods as _, OptionalExtension as _, QueryDsl, update};
+use diesel::{ExpressionMethods as _, OptionalExtension as _, QueryDsl, sqlite::Sqlite, update};
 use diesel_async::{AsyncConnection, RunQueryDsl};
 use kameo::actor::ActorRef;
 use tracing::error;
@@ -33,21 +33,18 @@ pub struct ChallengeSolution {
     pub solution: Vec<u8>,
 }
 
-pub struct AuthOk {
-    pub id: i32,
-    pub pubkey: authn::PublicKey,
-}
-
 smlang::statemachine!(
     name: Auth,
     custom_error: true,
     transitions: {
         *Init + AuthRequest(ChallengeRequest) / async prepare_challenge = SentChallenge(ChallengeContext),
-        Init + BootstrapAuthRequest(BootstrapAuthRequest) / async verify_bootstrap_token = AuthOk(AuthOk),
+        Init + BootstrapAuthRequest(BootstrapAuthRequest) / async verify_bootstrap_token = AuthOk(AuthCredentials),
-        SentChallenge(ChallengeContext) + ReceivedSolution(ChallengeSolution) / async verify_solution = AuthOk(AuthOk),
+        SentChallenge(ChallengeContext) + ReceivedSolution(ChallengeSolution) / async verify_solution = AuthOk(AuthCredentials),
     }
 );
 
+const NONCE_START: i32 = 1;
+
 /// Returns the current nonce, ready to use for the challenge nonce.
 async fn get_current_nonce_and_id(
     db: &DatabasePool,
@@ -94,9 +91,12 @@ async fn verify_integrity(
     let _result = integrity::verify_entity(
         &mut db_conn,
         vault,
-        &UserAgentCredentials {
+        &AuthCredentials {
-            pubkey: pubkey.clone(),
+            creds: Credentials {
-            nonce,
+                id,
+                pubkey: pubkey.clone(),
+            },
+            new_nonce: nonce,
         },
         id,
     )
@@ -109,49 +109,46 @@ async fn verify_integrity(
     Ok(())
 }
 
-async fn create_nonce(
+async fn compute_current_nonce(
-    db: &DatabasePool,
+    conn: &mut impl AsyncConnection<Backend = Sqlite>,
-    vault: &ActorRef<Vault>,
     pubkey: &authn::PublicKey,
 ) -> Result<(i32, i32), Error> {
-    let mut db_conn = db.get().await.map_err(|e| {
+    update(useragent_client::table)
-        error!(error = ?e, "Database pool error");
+        .filter(useragent_client::public_key.eq(pubkey.to_bytes()))
-        Error::internal("Database unavailable")
+        .set(useragent_client::nonce.eq(useragent_client::nonce + 1))
-    })?;
+        .returning((useragent_client::id, useragent_client::nonce))
-    let (id, new_nonce) = db_conn
+        .get_result(conn)
-        .exclusive_transaction(|conn| {
+        .await
-            Box::pin(async move {
+        .map_err(|e| {
-                let (id, new_nonce): (i32, i32) = update(useragent_client::table)
+            error!(error = ?e, "Database error incrementing nonce");
-                    .filter(useragent_client::public_key.eq(pubkey.to_bytes()))
+            Error::internal("Database operation failed")
-                    .set(useragent_client::nonce.eq(useragent_client::nonce + 1))
-                    .returning((useragent_client::id, useragent_client::nonce))
-                    .get_result(conn)
-                    .await
-                    .map_err(|e| {
-                        error!(error = ?e, "Database error");
-                        Error::internal("Database operation failed")
-                    })?;
-
-                integrity::sign_entity(
-                    conn,
-                    vault,
-                    &UserAgentCredentials {
-                        pubkey: pubkey.clone(),
-                        nonce: new_nonce,
-                    },
-                    id,
-                )
-                .await
-                .map_err(|e| {
-                    error!(?e, "Integrity signature update failed");
-                    Error::internal("Database error")
-                })?;
-
-                Result::<_, Error>::Ok((id, new_nonce))
-            })
         })
-        .await?;
+}
-    Ok((id, new_nonce))
+
+async fn resign_credentials(
+    conn: &mut impl AsyncConnection<Backend = Sqlite>,
+    vault: &ActorRef<Vault>,
+    id: i32,
+    pubkey: &authn::PublicKey,
+    new_nonce: i32,
+) -> Result<(), Error> {
+    integrity::sign_entity(
+        conn,
+        vault,
+        &AuthCredentials {
+            creds: Credentials {
+                id,
+                pubkey: pubkey.clone(),
+            },
+            new_nonce,
+        },
+        id,
+    )
+    .await
+    .map_err(|e| {
+        error!(?e, "Integrity signature update failed");
+        Error::internal("Database error")
+    })
 }
 
 async fn register_key(db: &DatabasePool, pubkey: &authn::PublicKey) -> Result<i32, Error> {
@@ -161,8 +158,6 @@ async fn register_key(db: &DatabasePool, pubkey: &authn::PublicKey) -> Result<i3
         Error::internal("Database unavailable")
     })?;
 
-    const NONCE_START: i32 = 1;
-
     let id: i32 = diesel::insert_into(useragent_client::table)
         .values((
             useragent_client::public_key.eq(pubkey_bytes),
@@ -200,9 +195,33 @@ where
         &mut self,
         ChallengeRequest { pubkey }: ChallengeRequest,
     ) -> Result<ChallengeContext, Self::Error> {
-        verify_integrity(&self.conn.db, &self.conn.actors.vault, &pubkey).await?;
+        let is_signing = integrity::is_signing_available(&self.conn.actors.vault)
+            .await
+            .unwrap_or(false);
 
-        let (id, nonce) = create_nonce(&self.conn.db, &self.conn.actors.vault, &pubkey).await?;
+        if is_signing {
+            verify_integrity(&self.conn.db, &self.conn.actors.vault, &pubkey).await?;
+        }
+
+        let vault = self.conn.actors.vault.clone();
+        let mut conn = self.conn.db.get().await.map_err(|e| {
+            error!(error = ?e, "Database pool error");
+            Error::internal("Database unavailable")
+        })?;
+
+        let (id, nonce) = conn
+            .exclusive_transaction(|conn| {
+                let pubkey = pubkey.clone();
+                let vault = vault.clone();
+                Box::pin(async move {
+                    let (id, new_nonce) = compute_current_nonce(conn, &pubkey).await?;
+                    if is_signing {
+                        resign_credentials(conn, &vault, id, &pubkey, new_nonce).await?;
+                    }
+                    Result::<_, Error>::Ok((id, new_nonce))
+                })
+            })
+            .await?;
 
         self.transport
             .send(Ok(Outbound::AuthChallenge { nonce }))
@@ -224,7 +243,7 @@ where
     async fn verify_bootstrap_token(
         &mut self,
         BootstrapAuthRequest { pubkey, token }: BootstrapAuthRequest,
-    ) -> Result<AuthOk, Self::Error> {
+    ) -> Result<AuthCredentials, Self::Error> {
         let token_ok: bool = self
             .conn
             .actors
@@ -245,12 +264,15 @@ where
 
         match token_ok {
             true => {
-                let id = register_key(&self.conn.db,  &pubkey).await?;
+                let id = register_key(&self.conn.db, &pubkey).await?;
                 self.transport
                     .send(Ok(Outbound::AuthSuccess))
                     .await
                     .map_err(|_| Error::Transport)?;
-                Ok(AuthOk { id, pubkey })
+                Ok(AuthCredentials {
+                    creds: Credentials { id, pubkey },
+                    new_nonce: NONCE_START,
+                })
             }
             false => {
                 error!("Invalid bootstrap token provided");
@@ -273,7 +295,7 @@ where
             key,
         }: &ChallengeContext,
         ChallengeSolution { solution }: ChallengeSolution,
-    ) -> Result<AuthOk, Self::Error> {
+    ) -> Result<AuthCredentials, Self::Error> {
         let signature = authn::Signature::try_from(solution.as_slice()).map_err(|_| {
             error!("Failed to decode signature in challenge solution");
             Error::InvalidChallengeSolution
@@ -287,7 +309,13 @@ where
                     .send(Ok(Outbound::AuthSuccess))
                     .await
                     .map_err(|_| Error::Transport)?;
-                Ok(AuthOk { id: *id, pubkey: key.clone() })
+                Ok(AuthCredentials {
+                    creds: Credentials {
+                        id: *id,
+                        pubkey: key.clone(),
+                    },
+                    new_nonce: *challenge_nonce,
+                })
             }
             false => {
                 self.transport

server/crates/arbiter-server/src/peers/user_agent/mod.rs

@@ -3,13 +3,45 @@ use crate::{
 };
 use arbiter_crypto::authn;
 
-#[derive(Debug)]
+pub mod auth;
-pub struct UserAgentCredentials {
+pub mod session;
+pub mod vault_gate;
+
+
+#[derive(Debug, Clone, Hash)]
+pub struct Credentials {
+    pub id: i32,
     pub pubkey: authn::PublicKey,
-    pub nonce: i32,
+}
+impl Hashable for Credentials {
+    fn hash<H: sha2::Digest>(&self, hasher: &mut H) {
+        self.id.hash(hasher);
+        self.pubkey.hash(hasher);
+    }
 }
 
-impl Integrable for UserAgentCredentials {
+#[derive(Debug, Clone)]
+pub struct AuthCredentials {
+    pub creds: Credentials,
+    // denotes new nonce, not current
+    pub new_nonce: i32,
+}
+
+impl Hashable for authn::PublicKey {
+    fn hash<H: sha2::Digest>(&self, hasher: &mut H) {
+        hasher.update(self.to_bytes());
+    }
+}
+
+impl Hashable for AuthCredentials {
+    fn hash<H: sha2::Digest>(&self, hasher: &mut H) {
+        self.creds.hash(hasher);
+        self.new_nonce.hash(hasher);
+    }
+}
+
+
+impl Integrable for AuthCredentials {
     const KIND: &'static str = "useragent_credentials";
 }
 
@@ -31,23 +63,9 @@ impl UserAgentConnection {
     }
 }
 
-pub mod auth;
+
-pub mod session;
 
 pub use auth::authenticate;
 pub use session::UserAgentSession;
 
-use crate::crypto::integrity::hashing::Hashable;
+use crate::crypto::integrity::hashing::Hashable;
-
-impl Hashable for authn::PublicKey {
-    fn hash<H: sha2::Digest>(&self, hasher: &mut H) {
-        hasher.update(self.to_bytes());
-    }
-}
-
-impl Hashable for UserAgentCredentials {
-    fn hash<H: sha2::Digest>(&self, hasher: &mut H) {
-        self.pubkey.hash(hasher);
-        self.nonce.hash(hasher);
-    }
-}

server/crates/arbiter-server/src/peers/user_agent/session/handlers.rs

@@ -28,88 +28,10 @@ use crate::db::models::{
 use crate::evm::policies::{Grant, SpecificGrant};
 use crate::{
     actors::vault::VaultState,
-    peers::user_agent::session::state::{UnsealContext, UserAgentEvents},
 };
 
-use super::{Error, UserAgentSession, state};
+use super::{Error, UserAgentSession};
 
-impl UserAgentSession {
-    fn take_unseal_secret(&mut self) -> Result<(EphemeralSecret, PublicKey), Error> {
-        let state::UserAgentStates::WaitingForUnsealKey(unseal_context) = self.state.state() else {
-            error!("Received encrypted key in invalid state");
-            return Err(Error::internal("Invalid state for unseal encrypted key"));
-        };
-
-        let ephemeral_secret = {
-            #[allow(
-                clippy::unwrap_used,
-                reason = "Mutex poison is unrecoverable and should panic"
-            )]
-            let mut secret_lock = unseal_context.secret.lock().unwrap();
-            let secret = secret_lock.take();
-            match secret {
-                Some(secret) => secret,
-                None => {
-                    drop(secret_lock);
-                    error!("Ephemeral secret already taken");
-                    return Err(Error::internal("Ephemeral secret already taken"));
-                }
-            }
-        };
-
-        Ok((ephemeral_secret, unseal_context.client_public_key))
-    }
-
-    fn decrypt_client_key_material(
-        ephemeral_secret: EphemeralSecret,
-        client_public_key: PublicKey,
-        nonce: &[u8],
-        ciphertext: &[u8],
-        associated_data: &[u8],
-    ) -> Result<SafeCell<Vec<u8>>, ()> {
-        let nonce = XNonce::from_slice(nonce);
-
-        let shared_secret = ephemeral_secret.diffie_hellman(&client_public_key);
-        let cipher = XChaCha20Poly1305::new(shared_secret.as_bytes().into());
-
-        let mut key_buffer = SafeCell::new(ciphertext.to_vec());
-
-        let decryption_result = key_buffer.write_inline(|write_handle| {
-            cipher.decrypt_in_place(nonce, associated_data, write_handle)
-        });
-
-        match decryption_result {
-            Ok(_) => Ok(key_buffer),
-            Err(err) => {
-                error!(?err, "Failed to decrypt encrypted key material");
-                Err(())
-            }
-        }
-    }
-}
-
-pub struct UnsealStartResponse {
-    pub server_pubkey: PublicKey,
-}
-
-#[derive(Debug, Error)]
-pub enum UnsealError {
-    #[error("Invalid key provided for unsealing")]
-    InvalidKey,
-    #[error("Internal error during unsealing process")]
-    General(#[from] super::Error),
-}
-
-#[derive(Debug, Error)]
-pub enum BootstrapError {
-    #[error("Invalid key provided for bootstrapping")]
-    InvalidKey,
-    #[error("Vault is already bootstrapped")]
-    AlreadyBootstrapped,
-
-    #[error("Internal error during bootstrapping process")]
-    General(#[from] super::Error),
-}
 
 #[derive(Debug, Error)]
 pub enum SignTransactionError {
@@ -129,153 +51,6 @@ pub enum GrantMutationError {
     Internal,
 }
 
-#[messages]
-impl UserAgentSession {
-    #[message]
-    pub async fn handle_unseal_request(
-        &mut self,
-        client_pubkey: x25519_dalek::PublicKey,
-    ) -> Result<UnsealStartResponse, Error> {
-        let secret = EphemeralSecret::random();
-        let public_key = PublicKey::from(&secret);
-
-        self.transition(UserAgentEvents::UnsealRequest(UnsealContext {
-            secret: Mutex::new(Some(secret)),
-            client_public_key: client_pubkey,
-        }))?;
-
-        Ok(UnsealStartResponse {
-            server_pubkey: public_key,
-        })
-    }
-
-    #[message]
-    pub async fn handle_unseal_encrypted_key(
-        &mut self,
-        nonce: Vec<u8>,
-        ciphertext: Vec<u8>,
-        associated_data: Vec<u8>,
-    ) -> Result<(), UnsealError> {
-        let (ephemeral_secret, client_public_key) = match self.take_unseal_secret() {
-            Ok(values) => values,
-            Err(Error::State) => {
-                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                return Err(UnsealError::InvalidKey);
-            }
-            Err(_err) => {
-                return Err(Error::internal("Failed to take unseal secret").into());
-            }
-        };
-
-        let seal_key_buffer = match Self::decrypt_client_key_material(
-            ephemeral_secret,
-            client_public_key,
-            &nonce,
-            &ciphertext,
-            &associated_data,
-        ) {
-            Ok(buffer) => buffer,
-            Err(()) => {
-                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                return Err(UnsealError::InvalidKey);
-            }
-        };
-
-        match self
-            .props
-            .actors
-            .vault
-            .ask(TryUnseal {
-                seal_key_raw: seal_key_buffer,
-            })
-            .await
-        {
-            Ok(_) => {
-                info!("Successfully unsealed key with client-provided key");
-                self.transition(UserAgentEvents::ReceivedValidKey)?;
-                Ok(())
-            }
-            Err(SendError::HandlerError(vault::Error::InvalidKey)) => {
-                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                Err(UnsealError::InvalidKey)
-            }
-            Err(SendError::HandlerError(err)) => {
-                error!(?err, "Vault failed to unseal key");
-                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                Err(UnsealError::InvalidKey)
-            }
-            Err(err) => {
-                error!(?err, "Failed to send unseal request to vault");
-                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                Err(Error::internal("Vault actor error").into())
-            }
-        }
-    }
-
-    #[message]
-    pub(crate) async fn handle_bootstrap_encrypted_key(
-        &mut self,
-        nonce: Vec<u8>,
-        ciphertext: Vec<u8>,
-        associated_data: Vec<u8>,
-    ) -> Result<(), BootstrapError> {
-        let (ephemeral_secret, client_public_key) = match self.take_unseal_secret() {
-            Ok(values) => values,
-            Err(Error::State) => {
-                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                return Err(BootstrapError::InvalidKey);
-            }
-            Err(err) => return Err(err.into()),
-        };
-
-        let seal_key_buffer = match Self::decrypt_client_key_material(
-            ephemeral_secret,
-            client_public_key,
-            &nonce,
-            &ciphertext,
-            &associated_data,
-        ) {
-            Ok(buffer) => buffer,
-            Err(()) => {
-                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                return Err(BootstrapError::InvalidKey);
-            }
-        };
-
-        match self
-            .props
-            .actors
-            .vault
-            .ask(Bootstrap {
-                seal_key_raw: seal_key_buffer,
-            })
-            .await
-        {
-            Ok(_) => {
-                info!("Successfully bootstrapped vault with client-provided key");
-                self.transition(UserAgentEvents::ReceivedValidKey)?;
-                Ok(())
-            }
-            Err(SendError::HandlerError(vault::Error::AlreadyBootstrapped)) => {
-                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                Err(BootstrapError::AlreadyBootstrapped)
-            }
-            Err(SendError::HandlerError(err)) => {
-                error!(?err, "Vault failed to bootstrap vault");
-                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                Err(BootstrapError::InvalidKey)
-            }
-            Err(err) => {
-                error!(?err, "Failed to send bootstrap request to vault");
-                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                Err(BootstrapError::General(Error::internal(
-                    "Vault actor error",
-                )))
-            }
-        }
-    }
-}
-
 #[messages]
 impl UserAgentSession {
     #[message]

server/crates/arbiter-server/src/peers/user_agent/session/mod.rs

@@ -1,12 +1,11 @@
 use arbiter_crypto::authn;
 use diesel::{ExpressionMethods, QueryDsl};
-use diesel_async::{AsyncConnection, RunQueryDsl};
+use diesel_async::{RunQueryDsl};
 use kameo_actors::message_bus::Register;
 
 use std::{borrow::Cow, collections::HashMap};
 
 use arbiter_proto::transport::Sender;
-use async_trait::async_trait;
 use kameo::{Actor, actor::ActorRef, messages, prelude::Message};
 use thiserror::Error;
 use tracing::error;
@@ -15,10 +14,8 @@ use crate::{
     actors::{
         flow_coordinator::{RegisterUserAgent, client_connect_approval::ClientApprovalController},
         vault::events,
-    }, crypto::integrity, db::schema::useragent_client, peers::{client::ClientProfile, user_agent::UserAgentCredentials}
+    }, crypto::integrity, db::schema::useragent_client, peers::{client::ClientProfile, user_agent::{AuthCredentials, Credentials}}
 };
-mod state;
-use state::{DummyContext, UserAgentEvents, UserAgentStateMachine};
 
 use super::{OutOfBand, UserAgentConnection};
 
@@ -58,41 +55,28 @@ pub struct PendingClientApproval {
 }
 
 pub struct UserAgentSession {
-    id: i32,
+    creds: Credentials,
-    pubkey: authn::PublicKey,
     props: UserAgentConnection,
-    state: UserAgentStateMachine<DummyContext>,
     sender: Box<dyn Sender<OutOfBand>>,
 
     pending_client_approvals: HashMap<Vec<u8>, PendingClientApproval>,
 }
 
-pub mod connection;
+pub mod handlers;
 
 impl UserAgentSession {
     pub(crate) fn new(
         props: UserAgentConnection,
-        id: i32,
+        creds: Credentials,
-        pubkey: authn::PublicKey,
         sender: Box<dyn Sender<OutOfBand>>,
     ) -> Self {
         Self {
-            id,
+            creds,
             props,
-            pubkey,
-            state: UserAgentStateMachine::new(DummyContext),
             sender,
             pending_client_approvals: Default::default(),
         }
     }
-
-    fn transition(&mut self, event: UserAgentEvents) -> Result<(), Error> {
-        self.state.process_event(event).map_err(|e| {
-            error!(?e, "State transition failed");
-            Error::State
-        })?;
-        Ok(())
-    }
 }
 
 #[messages]
@@ -128,61 +112,6 @@ impl UserAgentSession {
     }
 }
 
-impl Message<events::VaultBootstrapped> for UserAgentSession {
-    type Reply = Result<(), Error>;
-
-    async fn handle(
-        &mut self,
-        _: events::VaultBootstrapped,
-        ctx: &mut kameo::prelude::Context<Self, Self::Reply>,
-    ) -> Self::Reply {
-        let Ok(mut conn) = self.props.db.get().await else {
-            error!("Failed to get database connection for vault bootstrapped event");
-            ctx.stop();
-            return Err(Error::internal("Failed to get database connection"));
-        };
-
-
-        let result = conn.exclusive_transaction(|conn| {
-            Box::pin(async {
-                let nonce: i32 = useragent_client::table
-                    .filter(useragent_client::id.eq(self.id))
-                    .select(useragent_client::nonce)
-                    .first::<i32>(conn)
-                    .await
-                    .map_err(|e| {
-                        error!(?e, "Failed to get nonce for useragent bootstrapping");
-                        Error::internal("Failed to sign user agent credentials")
-                    })?;
-
-                let entity = UserAgentCredentials {
-                    pubkey: self.pubkey.clone(),
-                    nonce,
-                };
-
-                integrity::sign_entity(conn, &self.props.actors.vault, &entity, self.id)
-                    .await
-                    .map_err(|e| {
-                        error!(?e, "Failed to sign user agent credentials during vault bootstrapping");
-                        Error::internal("Failed to sign user agent credentials")
-                    })?;
-
-                Result::<_, Error>::Ok(())
-            })
-        }).await;
-
-        match result {
-            Ok(_) => Ok(()),
-            Err(err) => {
-                error!(?err, "Error during vault bootstrapping");
-                ctx.stop();
-                Err(err)
-            },
-        }
-
-    }
-}
-
 impl Actor for UserAgentSession {
     type Args = Self;
 
@@ -192,21 +121,6 @@ impl Actor for UserAgentSession {
         args: Self::Args,
         this: kameo::prelude::ActorRef<Self>,
     ) -> Result<Self, Self::Error> {
-        args.props
-            .actors
-            .events
-            .tell(Register(
-                this.clone().recipient::<events::VaultBootstrapped>(),
-            ))
-            .await
-            .map_err(|err| {
-                error!(
-                    ?err,
-                    "Failed to register user agent connection with event bus"
-                );
-                Error::internal("Failed to register user agent connection with event bus")
-            })?;
-
         args.props
             .actors
             .flow_coordinator

server/crates/arbiter-server/src/peers/user_agent/session/state.rs

@@ -1,27 +0,0 @@
-use std::sync::Mutex;
-
-use x25519_dalek::{EphemeralSecret, PublicKey};
-
-pub struct UnsealContext {
-    pub client_public_key: PublicKey,
-    pub secret: Mutex<Option<EphemeralSecret>>,
-}
-
-smlang::statemachine!(
-    name: UserAgent,
-    custom_error: false,
-    transitions: {
-        *Idle + UnsealRequest(UnsealContext) / generate_temp_keypair = WaitingForUnsealKey(UnsealContext),
-        WaitingForUnsealKey(UnsealContext) + ReceivedValidKey = Unsealed,
-        WaitingForUnsealKey(UnsealContext) + ReceivedInvalidKey = Idle,
-    }
-);
-
-pub struct DummyContext;
-impl UserAgentStateMachineContext for DummyContext {
-    #[allow(missing_docs)]
-    #[allow(clippy::unused_unit)]
-    fn generate_temp_keypair(&mut self, event_data: UnsealContext) -> Result<UnsealContext, ()> {
-        Ok(event_data)
-    }
-}

server/crates/arbiter-server/src/peers/user_agent/vault_gate/mod.rs

@@ -0,0 +1,295 @@
+use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
+use chacha20poly1305::{AeadInPlace, KeyInit as _, XChaCha20Poly1305, XNonce};
+use kameo::{Actor, error::SendError, messages, prelude::Message};
+use kameo_actors::message_bus::Register;
+use tokio::sync::oneshot;
+use tracing::{error, info};
+use x25519_dalek::{EphemeralSecret, PublicKey, SharedSecret};
+
+pub mod state;
+use state::*;
+
+use super::{AuthCredentials, Credentials};
+use crate::{
+    actors::{
+        GlobalActors,
+        vault::{self, Bootstrap, TryUnseal, events},
+    },
+    crypto::integrity::{self, AttestationStatus},
+    db::DatabasePool,
+};
+
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("Vault is already bootstrapped")]
+    AlreadyBootstrapped,
+    #[error("Invalid key provided")]
+    InvalidKey,
+
+    #[error("State transition failed")]
+    State,
+
+    #[error("Internal error: {0}")]
+    Internal(String),
+}
+impl Error {
+    fn internal(message: impl Into<String>) -> Self {
+        Self::Internal(message.into())
+    }
+}
+
+pub struct HandshakeResponse {
+    pub server_pubkey: PublicKey,
+}
+
+pub struct VaultGate {
+    pub auth_creds: AuthCredentials,
+    pub promotion_tx: Option<oneshot::Sender<Result<Credentials, Error>>>,
+    pub state: State,
+    pub actors: GlobalActors,
+    pub db: DatabasePool,
+}
+
+impl VaultGate {
+    pub fn new(
+        auth_creds: AuthCredentials,
+        actors: GlobalActors,
+        db: DatabasePool,
+        promotion_tx: oneshot::Sender<Result<Credentials, Error>>,
+    ) -> Self {
+        Self {
+            auth_creds,
+            state: State::default(),
+            actors,
+            db,
+            promotion_tx: Some(promotion_tx),
+        }
+    }
+}
+
+impl Actor for VaultGate {
+    type Args = Self;
+
+    type Error = ();
+
+    async fn on_start(
+        args: Self::Args,
+        actor_ref: kameo::prelude::ActorRef<Self>,
+    ) -> Result<Self, Self::Error> {
+        let _ = args
+            .actors
+            .events
+            .tell(Register(
+                actor_ref.clone().recipient::<events::Bootstrapped>(),
+            ))
+            .await;
+        let _ = args
+            .actors
+            .events
+            .tell(Register(actor_ref.recipient::<events::Unsealed>()))
+            .await;
+        Ok(args)
+    }
+}
+
+impl VaultGate {
+    fn decrypt_key(
+        secret: &SharedSecret,
+        nonce: &[u8],
+        ciphertext: &[u8],
+        associated_data: &[u8],
+    ) -> Result<SafeCell<Vec<u8>>, ()> {
+        let nonce = XNonce::from_slice(nonce);
+
+        let cipher = XChaCha20Poly1305::new(secret.as_bytes().into());
+
+        let mut key_buffer = SafeCell::new(ciphertext.to_vec());
+
+        let decryption_result = key_buffer.write_inline(|write_handle| {
+            cipher.decrypt_in_place(nonce, associated_data, write_handle)
+        });
+
+        match decryption_result {
+            Ok(_) => Ok(key_buffer),
+            Err(err) => {
+                error!(?err, "Failed to decrypt encrypted key material");
+                Err(())
+            }
+        }
+    }
+}
+
+#[messages]
+impl VaultGate {
+    #[message]
+    pub async fn handle_handshake(
+        &mut self,
+        client_pubkey: x25519_dalek::PublicKey,
+    ) -> Result<HandshakeResponse, Error> {
+        let ephemeral_secret = EphemeralSecret::random();
+        let public_key = PublicKey::from(&ephemeral_secret);
+
+        let secret = ephemeral_secret.diffie_hellman(&client_pubkey);
+
+        self.state = State::ReadyForExchange {
+            server_key: public_key.clone(),
+            secret,
+        };
+
+        Ok(HandshakeResponse {
+            server_pubkey: public_key,
+        })
+    }
+
+    #[message]
+    pub async fn handle_unseal_encrypted_key(
+        &mut self,
+        nonce: Vec<u8>,
+        ciphertext: Vec<u8>,
+        associated_data: Vec<u8>,
+    ) -> Result<(), Error> {
+        let State::ReadyForExchange { secret, .. } = &self.state else {
+            return Err(Error::State);
+        };
+
+        let seal_key_buffer = match Self::decrypt_key(secret, &nonce, &ciphertext, &associated_data)
+        {
+            Ok(buffer) => buffer,
+            Err(()) => {
+                return Err(Error::InvalidKey);
+            }
+        };
+
+        match self
+            .actors
+            .vault
+            .ask(TryUnseal {
+                seal_key_raw: seal_key_buffer,
+            })
+            .await
+        {
+            Ok(_) => {
+                info!("Successfully unsealed key with client-provided key");
+                Ok(())
+            }
+            Err(SendError::HandlerError(vault::Error::InvalidKey)) => Err(Error::InvalidKey),
+            Err(SendError::HandlerError(err)) => {
+                error!(?err, "Vault failed to unseal key");
+                Err(Error::InvalidKey)
+            }
+            Err(err) => {
+                error!(?err, "Failed to send unseal request to vault");
+                Err(Error::internal("Vault actor error").into())
+            }
+        }
+    }
+
+    #[message]
+    pub(crate) async fn handle_bootstrap_encrypted_key(
+        &mut self,
+        nonce: Vec<u8>,
+        ciphertext: Vec<u8>,
+        associated_data: Vec<u8>,
+    ) -> Result<(), Error> {
+        let State::ReadyForExchange { secret, .. } = &self.state else {
+            return Err(Error::State);
+        };
+
+        let seal_key_buffer = match Self::decrypt_key(secret, &nonce, &ciphertext, &associated_data)
+        {
+            Ok(buffer) => buffer,
+            Err(()) => {
+                return Err(Error::InvalidKey);
+            }
+        };
+
+        match self
+            .actors
+            .vault
+            .ask(Bootstrap {
+                seal_key_raw: seal_key_buffer,
+            })
+            .await
+        {
+            Ok(_) => {
+                info!("Successfully bootstrapped vault with client-provided key");
+                Ok(())
+            }
+            Err(SendError::HandlerError(vault::Error::AlreadyBootstrapped)) => {
+                Err(Error::AlreadyBootstrapped)
+            }
+            Err(SendError::HandlerError(err)) => {
+                error!(?err, "Vault failed to bootstrap vault");
+                Err(Error::InvalidKey)
+            }
+            Err(err) => {
+                error!(?err, "Failed to send bootstrap request to vault");
+                Err(Error::internal("Vault error"))
+            }
+        }
+    }
+}
+
+impl Message<events::Bootstrapped> for VaultGate {
+    type Reply = ();
+
+    async fn handle(
+        &mut self,
+        _: events::Bootstrapped,
+        ctx: &mut kameo::prelude::Context<Self, Self::Reply>,
+    ) -> Self::Reply {
+        let result = async {
+            let mut conn = self.db.get().await.map_err(|_| Error::internal("DB unavailable"))?;
+            integrity::sign_entity(&mut conn, &self.actors.vault, &self.auth_creds, self.auth_creds.creds.id)
+                .await
+                .map_err(|e| {
+                    error!(?e, "Failed to sign integrity envelope on bootstrap");
+                    Error::internal("Integrity sign failed")
+                })?;
+            Ok(self.auth_creds.creds.clone())
+        }
+        .await;
+
+        if let Some(tx) = self.promotion_tx.take() {
+            let _ = tx.send(result);
+        }
+        ctx.stop();
+    }
+}
+
+impl Message<events::Unsealed> for VaultGate {
+    type Reply = ();
+
+    async fn handle(
+        &mut self,
+        _: events::Unsealed,
+        ctx: &mut kameo::prelude::Context<Self, Self::Reply>,
+    ) -> Self::Reply {
+        let result = async {
+            let mut conn = self.db.get().await.map_err(|_| Error::internal("DB unavailable"))?;
+            match integrity::verify_entity(
+                &mut conn,
+                &self.actors.vault,
+                &self.auth_creds,
+                self.auth_creds.creds.id,
+            )
+            .await
+            {
+                Ok(AttestationStatus::Attested) => Ok(self.auth_creds.creds.clone()),
+                Ok(AttestationStatus::Unavailable) => {
+                    Err(Error::internal("Vault sealed during promotion"))
+                }
+                Err(e) => {
+                    error!(?e, "Integrity verification failed during unseal promotion");
+                    Err(Error::InvalidKey)
+                }
+            }
+        }
+        .await;
+
+        if let Some(tx) = self.promotion_tx.take() {
+            let _ = tx.send(result);
+        }
+        ctx.stop();
+    }
+}

server/crates/arbiter-server/src/peers/user_agent/vault_gate/state.rs

@@ -0,0 +1,18 @@
+use std::sync::Mutex;
+
+use x25519_dalek::{EphemeralSecret, PublicKey, SharedSecret};
+
+
+
+pub struct Handshake {
+    client_pubkey: PublicKey,
+}
+
+
+
+#[derive(Default)]
+pub enum State {
+    #[default]
+    Idle, 
+    ReadyForExchange { server_key: PublicKey, secret: SharedSecret },
+}

server/crates/arbiter-server/tests/user_agent/auth.rs

@@ -8,7 +8,7 @@ use arbiter_server::{
     actors::{GlobalActors, bootstrap::GetToken, vault::Bootstrap},
     crypto::integrity,
     db::{self, schema},
-    peers::user_agent::{UserAgentConnection, UserAgentCredentials, auth},
+    peers::user_agent::{AuthCredentials, Credentials, UserAgentConnection, auth},
 };
 use diesel::{ExpressionMethods as _, QueryDsl, insert_into};
 use diesel_async::RunQueryDsl;
@@ -144,9 +144,12 @@ pub async fn test_challenge_auth() {
         integrity::sign_entity(
             &mut conn,
             &actors.vault,
-            &UserAgentCredentials {
+            &AuthCredentials {
-                pubkey: new_key.verifying_key().into(),
+                creds: Credentials {
-                nonce: 1,
+                    id,
+                    pubkey: new_key.verifying_key().into(),
+                },
+                new_nonce: 1,
             },
             id,
         )
@@ -282,9 +285,12 @@ pub async fn test_challenge_auth_rejects_invalid_signature() {
         integrity::sign_entity(
             &mut conn,
             &actors.vault,
-            &UserAgentCredentials {
+            &AuthCredentials {
-                pubkey: new_key.verifying_key().into(),
+                creds: Credentials {
-                nonce: 1,
+                    id,
+                    pubkey: new_key.verifying_key().into(),
+                },
+                new_nonce: 1,
             },
             id,
         )

server/crates/arbiter-server/tests/user_agent/unseal.rs

@@ -1,4 +1,7 @@
-use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
+use arbiter_crypto::{
+    authn,
+    safecell::{SafeCell, SafeCellHandle as _},
+};
 use arbiter_server::{
     actors::{
         GlobalActors,
@@ -6,18 +9,19 @@ use arbiter_server::{
     },
     db,
     peers::user_agent::{
-        UserAgentSession,
+        AuthCredentials, Credentials,
-        session::connection::{HandleUnsealEncryptedKey, HandleUnsealRequest, UnsealError},
+        vault_gate::{Error as VaultGateError, HandleHandshake, HandleUnsealEncryptedKey, VaultGate},
     },
 };
 
 use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit};
 use kameo::actor::Spawn as _;
+use tokio::sync::oneshot;
 use x25519_dalek::{EphemeralSecret, PublicKey};
 
-async fn setup_sealed_user_agent(
+async fn setup_sealed_gate(
     seal_key: &[u8],
-) -> (db::DatabasePool, kameo::actor::ActorRef<UserAgentSession>) {
+) -> (db::DatabasePool, kameo::actor::ActorRef<VaultGate>, oneshot::Receiver<Result<Credentials, VaultGateError>>) {
     let db = db::create_test_pool().await;
     let actors = GlobalActors::spawn(db.clone()).await.unwrap();
 
@@ -30,20 +34,26 @@ async fn setup_sealed_user_agent(
         .unwrap();
     actors.vault.ask(Seal).await.unwrap();
 
-    let session = UserAgentSession::spawn(UserAgentSession::new_test(db.clone(), actors));
+    let (promotion_tx, promotion_rx) = oneshot::channel();
+    let pubkey = authn::SigningKey::generate().public_key();
+    let auth_creds = AuthCredentials {
+        creds: Credentials { id: 1, pubkey },
+        new_nonce: 1,
+    };
+    let gate = VaultGate::spawn(VaultGate::new(auth_creds, actors, db.clone(), promotion_tx));
 
-    (db, session)
+    (db, gate, promotion_rx)
 }
 
 async fn client_dh_encrypt(
-    user_agent: &kameo::actor::ActorRef<UserAgentSession>,
+    gate: &kameo::actor::ActorRef<VaultGate>,
     key_to_send: &[u8],
 ) -> HandleUnsealEncryptedKey {
     let client_secret = EphemeralSecret::random();
     let client_public = PublicKey::from(&client_secret);
 
-    let response = user_agent
+    let response = gate
-        .ask(HandleUnsealRequest {
+        .ask(HandleHandshake {
             client_pubkey: client_public,
         })
         .await
@@ -71,26 +81,26 @@ async fn client_dh_encrypt(
 #[test_log::test]
 pub async fn test_unseal_success() {
     let seal_key = b"test-seal-key";
-    let (_db, user_agent) = setup_sealed_user_agent(seal_key).await;
+    let (_db, gate, _promotion_rx) = setup_sealed_gate(seal_key).await;
 
-    let encrypted_key = client_dh_encrypt(&user_agent, seal_key).await;
+    let encrypted_key = client_dh_encrypt(&gate, seal_key).await;
 
-    let response = user_agent.ask(encrypted_key).await;
+    let response = gate.ask(encrypted_key).await;
     assert!(matches!(response, Ok(())));
 }
 
 #[tokio::test]
 #[test_log::test]
 pub async fn test_unseal_wrong_seal_key() {
-    let (_db, user_agent) = setup_sealed_user_agent(b"correct-key").await;
+    let (_db, gate, _promotion_rx) = setup_sealed_gate(b"correct-key").await;
 
-    let encrypted_key = client_dh_encrypt(&user_agent, b"wrong-key").await;
+    let encrypted_key = client_dh_encrypt(&gate, b"wrong-key").await;
 
-    let response = user_agent.ask(encrypted_key).await;
+    let response = gate.ask(encrypted_key).await;
     assert!(matches!(
         response,
         Err(kameo::error::SendError::HandlerError(
-            UnsealError::InvalidKey
+            VaultGateError::InvalidKey
         ))
     ));
 }
@@ -98,19 +108,18 @@ pub async fn test_unseal_wrong_seal_key() {
 #[tokio::test]
 #[test_log::test]
 pub async fn test_unseal_corrupted_ciphertext() {
-    let (_db, user_agent) = setup_sealed_user_agent(b"test-key").await;
+    let (_db, gate, _promotion_rx) = setup_sealed_gate(b"test-key").await;
 
     let client_secret = EphemeralSecret::random();
     let client_public = PublicKey::from(&client_secret);
 
-    user_agent
+    gate.ask(HandleHandshake {
-        .ask(HandleUnsealRequest {
+        client_pubkey: client_public,
-            client_pubkey: client_public,
+    })
-        })
+    .await
-        .await
+    .unwrap();
-        .unwrap();
 
-    let response = user_agent
+    let response = gate
         .ask(HandleUnsealEncryptedKey {
             nonce: vec![0u8; 24],
             ciphertext: vec![0u8; 32],
@@ -121,7 +130,7 @@ pub async fn test_unseal_corrupted_ciphertext() {
     assert!(matches!(
         response,
         Err(kameo::error::SendError::HandlerError(
-            UnsealError::InvalidKey
+            VaultGateError::InvalidKey
         ))
     ));
 }
@@ -130,24 +139,24 @@ pub async fn test_unseal_corrupted_ciphertext() {
 #[test_log::test]
 pub async fn test_unseal_retry_after_invalid_key() {
     let seal_key = b"real-seal-key";
-    let (_db, user_agent) = setup_sealed_user_agent(seal_key).await;
+    let (_db, gate, _promotion_rx) = setup_sealed_gate(seal_key).await;
 
     {
-        let encrypted_key = client_dh_encrypt(&user_agent, b"wrong-key").await;
+        let encrypted_key = client_dh_encrypt(&gate, b"wrong-key").await;
 
-        let response = user_agent.ask(encrypted_key).await;
+        let response = gate.ask(encrypted_key).await;
         assert!(matches!(
             response,
             Err(kameo::error::SendError::HandlerError(
-                UnsealError::InvalidKey
+                VaultGateError::InvalidKey
             ))
         ));
     }
 
     {
-        let encrypted_key = client_dh_encrypt(&user_agent, seal_key).await;
+        let encrypted_key = client_dh_encrypt(&gate, seal_key).await;
 
-        let response = user_agent.ask(encrypted_key).await;
+        let response = gate.ask(encrypted_key).await;
         assert!(matches!(response, Ok(())));
     }
 }