2026-04-18 11:54:31 +00:00
3 changed files with 190 additions and 47 deletions
--- a/server/crates/arbiter-server/src/grpc/user_agent/vault_gate.rs
+++ b/server/crates/arbiter-server/src/grpc/user_agent/vault_gate.rs
@@ -1,16 +1,19 @@
 use arbiter_proto::{
-    proto::user_agent::{
-        user_agent_request::Payload as UserAgentRequestPayload,
-        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,
-                request::Payload as UnsealRequestPayload,
-                response::Payload as UnsealResponsePayload,
+    proto::{
+        shared::VaultState as ProtoVaultState,
+        user_agent::{
+            user_agent_request::Payload as UserAgentRequestPayload,
+            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,
+                    request::Payload as UnsealRequestPayload,
+                    response::Payload as UnsealResponsePayload,
+                },
            },
        },
    },
@@ -21,8 +24,11 @@ use tonic::Status;
 use tracing::warn;

 use super::auth::AuthTransportAdapter;
-use crate::peers::user_agent::vault_gate::{
-    self as vault_gate, HandleBootstrapEncryptedKey, HandleHandshake, HandleUnsealEncryptedKey,
+use crate::{
+    actors::vault::VaultState,
+    peers::user_agent::vault_gate::{
+        self as vault_gate, HandleBootstrapEncryptedKey, HandleHandshake, HandleUnsealEncryptedKey,
+    },
 };

 fn wrap_vault_response(payload: VaultResponsePayload) -> UserAgentResponsePayload {
@@ -43,15 +49,7 @@ fn wrap_bootstrap_response(result: ProtoBootstrapResult) -> UserAgentResponsePay
    }))
 }

-impl AuthTransportAdapter<'_> {
-    async fn send_query_state(&mut self) -> Result<(), TransportError> {
-        use arbiter_proto::proto::shared::VaultState as ProtoVaultState;
-        self.send_response_payload(wrap_vault_response(VaultResponsePayload::State(
-            ProtoVaultState::Sealed.into(),
-        )))
-        .await
-    }
-}
+impl AuthTransportAdapter<'_> {}

 #[async_trait]
 impl Receiver<vault_gate::Inbound> for AuthTransportAdapter<'_> {
@@ -60,7 +58,10 @@ impl Receiver<vault_gate::Inbound> for AuthTransportAdapter<'_> {
            let request = match self.bi_mut().recv().await? {
                Ok(request) => request,
                Err(error) => {
-                    warn!(?error, "Failed to receive user agent request during vault gate");
+                    warn!(
+                        ?error,
+                        "Failed to receive user agent request during vault gate"
+                    );
                    return None;
                }
            };
@@ -94,23 +95,24 @@ impl Receiver<vault_gate::Inbound> for AuthTransportAdapter<'_> {
            let Some(vault_payload) = vault_req.payload else {
                let _ = self
                    .bi_mut()
-                    .send(Err(Status::invalid_argument("Missing vault request payload")))
+                    .send(Err(Status::invalid_argument(
+                        "Missing vault request payload",
+                    )))
                    .await;
                return None;
            };

            match vault_payload {
                VaultRequestPayload::QueryState(_) => {
-                    if self.send_query_state().await.is_err() {
-                        return None;
-                    }
-                    continue;
+                    return Some(vault_gate::Inbound::HandleVaultState);
                }
                VaultRequestPayload::Unseal(req) => {
                    let Some(unseal_payload) = req.payload else {
                        let _ = self
                            .bi_mut()
-                            .send(Err(Status::invalid_argument("Missing unseal request payload")))
+                            .send(Err(Status::invalid_argument(
+                                "Missing unseal request payload",
+                            )))
                            .await;
                        return None;
                    };
@@ -181,13 +183,29 @@ impl Sender<Result<vault_gate::Outbound, vault_gate::Error>> for AuthTransportAd
        };

        let payload = match outbound {
-            vault_gate::Outbound::HandleHandshake(Ok(response)) => {
-                wrap_unseal_response(UnsealResponsePayload::Start(
-                    proto_unseal::UnsealStartResponse {
-                        server_pubkey: response.server_pubkey.as_bytes().to_vec(),
-                    },
-                ))
-            }
+            vault_gate::Outbound::HandleVaultState(result) => match result {
+                Ok(state) => {
+                    let state = match state {
+                        VaultState::Unbootstrapped => ProtoVaultState::Unbootstrapped,
+                        VaultState::Sealed => ProtoVaultState::Sealed,
+                        VaultState::Unsealed => ProtoVaultState::Unsealed,
+                    };
+
+                    wrap_vault_response(VaultResponsePayload::State(state.into()))
+                }
+                Err(err) => {
+                    warn!(?err, "vault state query failed");
+                    return self
+                        .bi_mut()
+                        .send(Err(Status::internal("Failed to query vault state")))
+                        .await;
+                }
+            },
+            vault_gate::Outbound::HandleHandshake(Ok(response)) => wrap_unseal_response(
+                UnsealResponsePayload::Start(proto_unseal::UnsealStartResponse {
+                    server_pubkey: response.server_pubkey.as_bytes().to_vec(),
+                }),
+            ),
            vault_gate::Outbound::HandleHandshake(Err(err)) => {
                warn!(?err, "handshake failed");
                return self
--- a/server/crates/arbiter-server/src/peers/user_agent/vault_gate/mod.rs
+++ b/server/crates/arbiter-server/src/peers/user_agent/vault_gate/mod.rs
@@ -13,7 +13,7 @@ use super::{AuthCredentials, Credentials};
 use crate::{
    actors::{
        GlobalActors,
-        vault::{self, Bootstrap, TryUnseal, events},
+        vault::{self, Bootstrap, GetState, TryUnseal, VaultState, events},
    },
    crypto::integrity::{self, AttestationStatus},
    db::DatabasePool,
@@ -228,6 +228,18 @@ impl VaultGate {
            }
        }
    }
+
+    #[message]
+    pub async fn handle_vault_state(&mut self) -> Result<VaultState, Error> {
+        let answer = self
+            .actors
+            .vault
+            .ask(GetState {})
+            .await
+            .map_err(|_| Error::internal("failed to query vault"))?;
+
+        Ok(answer)
+    }
 }

 impl Message<events::Bootstrapped> for VaultGate {
@@ -239,13 +251,22 @@ impl Message<events::Bootstrapped> for VaultGate {
        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)
+            let mut conn = self
+                .db
+                .get()
                .await
-                .map_err(|e| {
-                    error!(?e, "Failed to sign integrity envelope on bootstrap");
-                    Error::internal("Integrity sign failed")
-                })?;
+                .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;
@@ -266,7 +287,11 @@ impl Message<events::Unsealed> for VaultGate {
        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"))?;
+            let mut conn = self
+                .db
+                .get()
+                .await
+                .map_err(|_| Error::internal("DB unavailable"))?;
            match integrity::verify_entity(
                &mut conn,
                &self.actors.vault,
--- a/useragent/rust/Cargo.lock
+++ b/useragent/rust/Cargo.lock
@@ -236,6 +236,7 @@ dependencies = [
 "memsafe",
 "ml-dsa",
 "rand",
+ "x-wing",
 ]

 [[package]]
@@ -713,7 +714,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601"
 dependencies = [
 "cfg-if",
- "cpufeatures",
+ "cpufeatures 0.3.0",
 "rand_core",
 ]

@@ -837,6 +838,15 @@ dependencies = [
 "libc",
 ]

+[[package]]
+name = "cpufeatures"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "cpufeatures"
 version = "0.3.0"
@@ -884,6 +894,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "77727bb15fa921304124b128af125e7e3b968275d1b108b379190264f4423710"
 dependencies = [
 "hybrid-array",
+ "rand_core",
 ]

 [[package]]
@@ -901,6 +912,32 @@ version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f27ae1dd37df86211c42e150270f82743308803d90a6f6e6651cd730d5e1732f"

+[[package]]
+name = "curve25519-dalek"
+version = "5.0.0-pre.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "335f1947f241137a14106b6f5acc5918a5ede29c9d71d3f2cb1678d5075d9fc3"
+dependencies = [
+ "cfg-if",
+ "cpufeatures 0.2.17",
+ "curve25519-dalek-derive",
+ "fiat-crypto",
+ "rustc_version",
+ "subtle",
+ "zeroize",
+]
+
+[[package]]
+name = "curve25519-dalek-derive"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "dart-sys"
 version = "4.1.5"
@@ -1314,6 +1351,12 @@ dependencies = [
 "bytemuck",
 ]

+[[package]]
+name = "fiat-crypto"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "64cd1e32ddd350061ae6edb1b082d7c54915b5c672c389143b9a63403a109f24"
+
 [[package]]
 name = "flate2"
 version = "1.1.9"
@@ -1773,6 +1816,7 @@ version = "0.4.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3944cf8cf766b40e2a1a333ee5e9b563f854d5fa49d6a8ca2764e97c6eddb214"
 dependencies = [
+ "ctutils",
 "typenum",
 "zeroize",
 ]
@@ -2005,7 +2049,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9e24a010dd405bd7ed803e5253182815b41bf2e6a80cc3bfc066658e03a198aa"
 dependencies = [
 "cfg-if",
- "cpufeatures",
+ "cpufeatures 0.3.0",
+]
+
+[[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",
 ]

 [[package]]
@@ -2207,12 +2261,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",
+ "sha3",
+ "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",
@@ -3488,6 +3557,12 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6637bab7722d379c8b41ba849228d680cc12d0a45ba1fa2b48f2a30577a06731"

+[[package]]
+name = "subtle"
+version = "2.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
+
 [[package]]
 name = "syn"
 version = "2.0.117"
@@ -4761,6 +4836,20 @@ version = "0.6.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4"

+[[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",
+ "sha3",
+ "x25519-dalek",
+ "zeroize",
+]
+
 [[package]]
 name = "x11-dl"
 version = "2.21.0"
@@ -4793,6 +4882,17 @@ version = "0.13.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ea6fc2961e4ef194dcbfe56bb845534d0dc8098940c7e5c012a258bfec6701bd"

+[[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",
+ "rand_core",
+ "zeroize",
+]
+
 [[package]]
 name = "xcursor"
 version = "0.3.10"