WIP: kameo::messages wiring for transport generalization

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

@@ -1,4 +1,4 @@
-use tokio::sync::{mpsc, oneshot};
+use tokio::sync::mpsc;
 
 use arbiter_proto::{
     proto::user_agent::{
@@ -9,17 +9,13 @@ use arbiter_proto::{
     transport::{Error as TransportError, Receiver, Sender, grpc::GrpcBi},
 };
 use async_trait::async_trait;
-use kameo::actor::{ActorRef, Spawn as _};
+use kameo::actor::ActorRef;
 use tonic::Status;
 use tracing::{error, info, warn};
 
 use crate::{
-    crypto::integrity,
     grpc::request_tracker::RequestTracker,
-    peers::user_agent::{
-        Credentials, OutOfBand, UserAgentConnection, UserAgentSession,
-        vault_gate::VaultGate,
-    },
+    peers::user_agent::{OutOfBand, UserAgentConnection, UserAgentSession},
 };
 
 mod auth;
@@ -129,115 +125,23 @@ pub async fn start(
 ) {
     let mut request_tracker = RequestTracker::default();
 
-    let auth_creds = match auth::start(&mut conn, &mut bi, &mut request_tracker).await {
-        Ok(creds) => creds,
-        Err(e) => {
-            warn!(error = ?e, "Authentication failed");
-            return;
-        }
-    };
-
-    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(conn, creds, 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;
-                    }
-                }
+    let actor = {
+        let transport = auth::AuthTransportAdapter::new(&mut bi, &mut request_tracker);
+        match crate::peers::user_agent::start(&mut conn, transport, Box::new(oob_adapter)).await {
+            Ok(actor) => actor,
+            Err(e) => {
+                warn!(error = ?e, "User agent connection failed");
+                return;
             }
         }
-    }
+    };
+
+    info!("User agent session established");
+
+    dispatch_loop(bi, actor.clone(), oob_receiver, request_tracker).await;
+    actor.kill();
 }
+

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

@@ -23,8 +23,8 @@ use crate::{
 };
 
 pub struct AuthTransportAdapter<'a> {
-    bi: &'a mut GrpcBi<UserAgentRequest, UserAgentResponse>,
-    request_tracker: &'a mut RequestTracker,
+    pub(super) bi: &'a mut GrpcBi<UserAgentRequest, UserAgentResponse>,
+    pub(super) request_tracker: &'a mut RequestTracker,
 }
 
 impl<'a> AuthTransportAdapter<'a> {
@@ -38,19 +38,35 @@ impl<'a> AuthTransportAdapter<'a> {
         }
     }
 
-    async fn send_user_agent_response(
+    pub(super) fn bi_mut(&mut self) -> &mut GrpcBi<UserAgentRequest, UserAgentResponse> {
+        self.bi
+    }
+
+    pub(super) fn tracker_mut(&mut self) -> &mut RequestTracker {
+        self.request_tracker
+    }
+
+    pub(super) async fn send_response_payload(
         &mut self,
-        payload: AuthResponsePayload,
+        payload: UserAgentResponsePayload,
     ) -> Result<(), TransportError> {
         self.bi
             .send(Ok(UserAgentResponse {
                 id: Some(self.request_tracker.current_request_id()),
-                payload: Some(UserAgentResponsePayload::Auth(proto_auth::Response {
-                    payload: Some(payload),
-                })),
+                payload: Some(payload),
             }))
             .await
     }
+
+    async fn send_user_agent_response(
+        &mut self,
+        payload: AuthResponsePayload,
+    ) -> Result<(), TransportError> {
+        self.send_response_payload(UserAgentResponsePayload::Auth(proto_auth::Response {
+            payload: Some(payload),
+        }))
+        .await
+    }
 }
 
 #[async_trait]
@@ -168,6 +184,6 @@ pub async fn start(
     bi: &mut GrpcBi<UserAgentRequest, UserAgentResponse>,
     request_tracker: &mut RequestTracker,
 ) -> Result<AuthCredentials, auth::Error> {
-    let transport = AuthTransportAdapter::new(bi, request_tracker);
-    auth::authenticate(conn, transport).await
+    let mut transport = AuthTransportAdapter::new(bi, request_tracker);
+    auth::authenticate(conn, &mut transport).await
 }

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

@@ -1,25 +1,28 @@
-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 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,
+            },
         },
     },
+    transport::{Bi, Error as TransportError, Receiver, Sender},
 };
-use kameo::{actor::ActorRef, error::SendError};
+use async_trait::async_trait;
 use tonic::Status;
 use tracing::warn;
 
+use super::auth::AuthTransportAdapter;
 use crate::peers::user_agent::vault_gate::{
     self as vault_gate, HandleBootstrapEncryptedKey, HandleHandshake, HandleUnsealEncryptedKey,
-    VaultGate,
 };
 
 fn wrap_vault_response(payload: VaultResponsePayload) -> UserAgentResponsePayload {
@@ -40,112 +43,196 @@ fn wrap_bootstrap_response(result: ProtoBootstrapResult) -> UserAgentResponsePay
     }))
 }
 
-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,
+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 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"));
-    };
+#[async_trait]
+impl Receiver<vault_gate::Inbound> for AuthTransportAdapter<'_> {
+    async fn recv(&mut self) -> Option<vault_gate::Inbound> {
+        loop {
+            let request = match self.bi_mut().recv().await? {
+                Ok(request) => request,
+                Err(error) => {
+                    warn!(?error, "Failed to receive user agent request during vault gate");
+                    return None;
+                }
+            };
 
-    match payload {
-        UnsealRequestPayload::Start(req) => handle_unseal_start(gate, req).await,
-        UnsealRequestPayload::EncryptedKey(req) => handle_unseal_encrypted_key(gate, req).await,
+            if let Err(err) = self.tracker_mut().request(request.id) {
+                let _ = self.bi_mut().send(Err(err)).await;
+                return None;
+            }
+
+            let Some(payload) = request.payload else {
+                let _ = self
+                    .bi_mut()
+                    .send(Err(Status::invalid_argument("Missing request payload")))
+                    .await;
+                return None;
+            };
+
+            let vault_req = match payload {
+                UserAgentRequestPayload::Vault(req) => req,
+                _ => {
+                    let _ = self
+                        .bi_mut()
+                        .send(Err(Status::permission_denied(
+                            "Only vault operations are permitted before unsealing",
+                        )))
+                        .await;
+                    return None;
+                }
+            };
+
+            let Some(vault_payload) = vault_req.payload else {
+                let _ = self
+                    .bi_mut()
+                    .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;
+                }
+                VaultRequestPayload::Unseal(req) => {
+                    let Some(unseal_payload) = req.payload else {
+                        let _ = self
+                            .bi_mut()
+                            .send(Err(Status::invalid_argument("Missing unseal request payload")))
+                            .await;
+                        return None;
+                    };
+                    match unseal_payload {
+                        UnsealRequestPayload::Start(start) => {
+                            let Ok(bytes) = <[u8; 32]>::try_from(start.client_pubkey) else {
+                                let _ = self
+                                    .bi_mut()
+                                    .send(Err(Status::invalid_argument(
+                                        "Invalid X25519 public key",
+                                    )))
+                                    .await;
+                                return None;
+                            };
+                            return Some(vault_gate::Inbound::HandleHandshake(HandleHandshake {
+                                client_pubkey: x25519_dalek::PublicKey::from(bytes),
+                            }));
+                        }
+                        UnsealRequestPayload::EncryptedKey(key) => {
+                            return Some(vault_gate::Inbound::HandleUnsealEncryptedKey(
+                                HandleUnsealEncryptedKey {
+                                    nonce: key.nonce,
+                                    ciphertext: key.ciphertext,
+                                    associated_data: key.associated_data,
+                                },
+                            ));
+                        }
+                    }
+                }
+                VaultRequestPayload::Bootstrap(req) => {
+                    let Some(encrypted_key) = req.encrypted_key else {
+                        let _ = self
+                            .bi_mut()
+                            .send(Err(Status::invalid_argument(
+                                "Missing bootstrap encrypted key",
+                            )))
+                            .await;
+                        return None;
+                    };
+                    return Some(vault_gate::Inbound::HandleBootstrapEncryptedKey(
+                        HandleBootstrapEncryptedKey {
+                            nonce: encrypted_key.nonce,
+                            ciphertext: encrypted_key.ciphertext,
+                            associated_data: encrypted_key.associated_data,
+                        },
+                    ));
+                }
+            }
+        }
     }
 }
 
-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"))?;
+#[async_trait]
+impl Sender<Result<vault_gate::Outbound, vault_gate::Error>> for AuthTransportAdapter<'_> {
+    async fn send(
+        &mut self,
+        item: Result<vault_gate::Outbound, vault_gate::Error>,
+    ) -> Result<(), TransportError> {
+        let outbound = match item {
+            Ok(outbound) => outbound,
+            Err(err) => {
+                warn!(?err, "vault gate produced transport-level error");
+                return self
+                    .bi_mut()
+                    .send(Err(Status::internal(err.to_string())))
+                    .await;
+            }
+        };
 
-    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")
-        })?;
+        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::HandleHandshake(Err(err)) => {
+                warn!(?err, "handshake failed");
+                return self
+                    .bi_mut()
+                    .send(Err(Status::internal("Failed to start unseal flow")))
+                    .await;
+            }
+            vault_gate::Outbound::HandleUnsealEncryptedKey(result) => {
+                let proto_result = match result {
+                    Ok(()) => ProtoUnsealResult::Success,
+                    Err(vault_gate::Error::InvalidKey) => ProtoUnsealResult::InvalidKey,
+                    Err(err) => {
+                        warn!(?err, "unseal failed");
+                        return self
+                            .bi_mut()
+                            .send(Err(Status::internal("Failed to unseal vault")))
+                            .await;
+                    }
+                };
+                wrap_unseal_response(UnsealResponsePayload::Result(proto_result.into()))
+            }
+            vault_gate::Outbound::HandleBootstrapEncryptedKey(result) => {
+                let proto_result = match result {
+                    Ok(()) => ProtoBootstrapResult::Success,
+                    Err(vault_gate::Error::InvalidKey) => ProtoBootstrapResult::InvalidKey,
+                    Err(vault_gate::Error::AlreadyBootstrapped) => {
+                        ProtoBootstrapResult::AlreadyBootstrapped
+                    }
+                    Err(err) => {
+                        warn!(?err, "bootstrap failed");
+                        return self
+                            .bi_mut()
+                            .send(Err(Status::internal("Failed to bootstrap vault")))
+                            .await;
+                    }
+                };
+                wrap_bootstrap_response(proto_result)
+            }
+        };
 
-    Ok(Some(wrap_unseal_response(UnsealResponsePayload::Start(
-        proto_unseal::UnsealStartResponse {
-            server_pubkey: response.server_pubkey.as_bytes().to_vec(),
-        },
-    ))))
+        self.send_response_payload(payload).await
+    }
 }
 
-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)))
+impl Bi<vault_gate::Inbound, Result<vault_gate::Outbound, vault_gate::Error>>
+    for AuthTransportAdapter<'_>
+{
 }