fix(server): sending fixed vault state when on stage

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

@@ -1,5 +1,7 @@
 use arbiter_proto::{
-    proto::user_agent::{
+    proto::{
+        shared::VaultState as ProtoVaultState,
+        user_agent::{
             user_agent_request::Payload as UserAgentRequestPayload,
             user_agent_response::Payload as UserAgentResponsePayload,
             vault::{
@@ -14,6 +16,7 @@ use arbiter_proto::{
                 },
             },
         },
+    },
     transport::{Bi, Error as TransportError, Receiver, Sender},
 };
 use async_trait::async_trait;
@@ -21,8 +24,11 @@ use tonic::Status;
 use tracing::warn;
 
 use super::auth::AuthTransportAdapter;
-use crate::peers::user_agent::vault_gate::{
+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<'_> {
+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
-    }
-}
 
 #[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 Some(vault_gate::Inbound::HandleVaultState);
-                        return None;
-                    }
-                    continue;
                 }
                 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)) => {
+            vault_gate::Outbound::HandleVaultState(result) => match result {
-                wrap_unseal_response(UnsealResponsePayload::Start(
+                Ok(state) => {
-                    proto_unseal::UnsealStartResponse {
+                    let state = match state {
-                        server_pubkey: response.server_pubkey.as_bytes().to_vec(),
+                        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

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,8 +251,17 @@ 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"))?;
+            let mut conn = self
-            integrity::sign_entity(&mut conn, &self.actors.vault, &self.auth_creds, self.auth_creds.creds.id)
+                .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");
@@ -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,

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"