WIP: kameo::messages wiring for transport generalization

server/Cargo.toml

@@ -26,7 +26,6 @@ thiserror = "2.0.18"
 async-trait = "0.1.89"
 futures = "0.3.32"
 tokio-stream = { version = "0.1.18", features = ["full"] }
-kameo = "0.20"
 prost-types = { version = "0.14.3", features = ["chrono"] }
 x25519-dalek = { version = "2.0.1", features = ["getrandom"] }
 rstest = "0.26.1"
@@ -47,3 +46,6 @@ miette = { version = "7.6.0", features = ["fancy", "serde"] }
 mutants = "0.0.4"
 ml-dsa = { version = "0.1.0-rc.8", features = ["zeroize"] }
 base64 = "0.22.1"
+kameo = {git = "https://github.com/hdbg/kameo.git", rev = "805b417"}
+kameo_actors = {git = "https://github.com/hdbg/kameo.git", rev = "805b417"}
+

server/crates/arbiter-proto/src/transport.rs

@@ -106,6 +106,36 @@ pub trait Receiver<Inbound>: Send + Sync {
 /// any built-in correlation mechanism between inbound and outbound items.
 pub trait Bi<Inbound, Outbound>: Sender<Outbound> + Receiver<Inbound> + Send + Sync {}
 
+#[async_trait]
+impl<T, Outbound> Sender<Outbound> for &mut T
+where
+    T: Sender<Outbound> + ?Sized,
+    Outbound: Send + 'static,
+{
+    async fn send(&mut self, item: Outbound) -> Result<(), Error> {
+        (**self).send(item).await
+    }
+}
+
+#[async_trait]
+impl<T, Inbound> Receiver<Inbound> for &mut T
+where
+    T: Receiver<Inbound> + ?Sized,
+    Inbound: Send + 'static,
+{
+    async fn recv(&mut self) -> Option<Inbound> {
+        (**self).recv().await
+    }
+}
+
+impl<T, Inbound, Outbound> Bi<Inbound, Outbound> for &mut T
+where
+    T: Bi<Inbound, Outbound> + ?Sized,
+    Inbound: Send + 'static,
+    Outbound: Send + 'static,
+{
+}
+
 pub trait SplittableBi<Inbound, Outbound>: Bi<Inbound, Outbound> {
     type Sender: Sender<Outbound>;
     type Receiver: Receiver<Inbound>;

server/crates/arbiter-server/Cargo.toml

@@ -58,7 +58,7 @@ ml-dsa.workspace = true
 ed25519-dalek.workspace = true
 x25519-dalek.workspace = true
 k256.workspace = true
-kameo_actors = "0.5.0"
+kameo_actors.workspace = true
 
 [dev-dependencies]
 insta = "1.46.3"

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>,
+    pub(super) bi: &'a mut GrpcBi<UserAgentRequest, UserAgentResponse>,
-    request_tracker: &'a mut RequestTracker,
+    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);
+    let mut transport = AuthTransportAdapter::new(bi, request_tracker);
-    auth::authenticate(conn, transport).await
+    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::{
+use arbiter_proto::{
-    user_agent_response::Payload as UserAgentResponsePayload,
+    proto::user_agent::{
-    vault::{
+        user_agent_request::Payload as UserAgentRequestPayload,
-        self as proto_vault,
+        user_agent_response::Payload as UserAgentResponsePayload,
-        bootstrap::{
+        vault::{
-            self as proto_bootstrap, BootstrapResult as ProtoBootstrapResult,
+            self as proto_vault,
-        },
+            bootstrap::{self as proto_bootstrap, BootstrapResult as ProtoBootstrapResult},
-        request::Payload as VaultRequestPayload,
+            request::Payload as VaultRequestPayload,
-        response::Payload as VaultResponsePayload,
+            response::Payload as VaultResponsePayload,
-        unseal::{
+            unseal::{
-            self as proto_unseal, UnsealResult as ProtoUnsealResult, UnsealStart,
+                self as proto_unseal, UnsealResult as ProtoUnsealResult,
-            request::Payload as UnsealRequestPayload, response::Payload as UnsealResponsePayload,
+                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(
+impl AuthTransportAdapter<'_> {
-    gate: &ActorRef<VaultGate>,
+    async fn send_query_state(&mut self) -> Result<(), TransportError> {
-    req: proto_vault::Request,
+        use arbiter_proto::proto::shared::VaultState as ProtoVaultState;
-) -> Result<Option<UserAgentResponsePayload>, Status> {
+        self.send_response_payload(wrap_vault_response(VaultResponsePayload::State(
-    let Some(payload) = req.payload else {
+            ProtoVaultState::Sealed.into(),
-        return Err(Status::invalid_argument("Missing vault request payload"));
+        )))
-    };
+        .await
-
-    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(
+#[async_trait]
-    gate: &ActorRef<VaultGate>,
+impl Receiver<vault_gate::Inbound> for AuthTransportAdapter<'_> {
-    req: proto_unseal::Request,
+    async fn recv(&mut self) -> Option<vault_gate::Inbound> {
-) -> Result<Option<UserAgentResponsePayload>, Status> {
+        loop {
-    let Some(payload) = req.payload else {
+            let request = match self.bi_mut().recv().await? {
-        return Err(Status::invalid_argument("Missing unseal request payload"));
+                Ok(request) => request,
-    };
+                Err(error) => {
+                    warn!(?error, "Failed to receive user agent request during vault gate");
+                    return None;
+                }
+            };
 
-    match payload {
+            if let Err(err) = self.tracker_mut().request(request.id) {
-        UnsealRequestPayload::Start(req) => handle_unseal_start(gate, req).await,
+                let _ = self.bi_mut().send(Err(err)).await;
-        UnsealRequestPayload::EncryptedKey(req) => handle_unseal_encrypted_key(gate, req).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(
+#[async_trait]
-    gate: &ActorRef<VaultGate>,
+impl Sender<Result<vault_gate::Outbound, vault_gate::Error>> for AuthTransportAdapter<'_> {
-    req: UnsealStart,
+    async fn send(
-) -> Result<Option<UserAgentResponsePayload>, Status> {
+        &mut self,
-    let client_pubkey = <[u8; 32]>::try_from(req.client_pubkey)
+        item: Result<vault_gate::Outbound, vault_gate::Error>,
-        .map(x25519_dalek::PublicKey::from)
+    ) -> Result<(), TransportError> {
-        .map_err(|_| Status::invalid_argument("Invalid X25519 public key"))?;
+        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
+        let payload = match outbound {
-        .ask(HandleHandshake { client_pubkey })
+            vault_gate::Outbound::HandleHandshake(Ok(response)) => {
-        .await
+                wrap_unseal_response(UnsealResponsePayload::Start(
-        .map_err(|err| {
+                    proto_unseal::UnsealStartResponse {
-            warn!(error = ?err, "Failed to handle unseal start");
+                        server_pubkey: response.server_pubkey.as_bytes().to_vec(),
-            Status::internal("Failed to start unseal flow")
+                    },
-        })?;
+                ))
+            }
+            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(
+        self.send_response_payload(payload).await
-        proto_unseal::UnsealStartResponse {
+    }
-            server_pubkey: response.server_pubkey.as_bytes().to_vec(),
-        },
-    ))))
 }
 
-async fn handle_unseal_encrypted_key(
+impl Bi<vault_gate::Inbound, Result<vault_gate::Outbound, vault_gate::Error>>
-    gate: &ActorRef<VaultGate>,
+    for AuthTransportAdapter<'_>
-    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

@@ -68,10 +68,10 @@ fn parse_auth_event(payload: Inbound) -> AuthEvents {
 
 pub async fn authenticate<T>(
     props: &mut UserAgentConnection,
-    transport: T,
+    transport: &mut T,
 ) -> Result<AuthCredentials, Error>
 where
-    T: Bi<Inbound, Result<Outbound, Error>> + Send,
+    T: Bi<Inbound, Result<Outbound, Error>> + Send + ?Sized,
 {
     let mut state = AuthStateMachine::new(AuthContext::new(props, transport));
 

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

@@ -174,20 +174,20 @@ async fn register_key(db: &DatabasePool, pubkey: &authn::PublicKey) -> Result<i3
     Ok(id)
 }
 
-pub struct AuthContext<'a, T> {
+pub struct AuthContext<'a, T: ?Sized> {
     pub(super) conn: &'a mut UserAgentConnection,
-    pub(super) transport: T,
+    pub(super) transport: &'a mut T,
 }
 
-impl<'a, T> AuthContext<'a, T> {
+impl<'a, T: ?Sized> AuthContext<'a, T> {
-    pub fn new(conn: &'a mut UserAgentConnection, transport: T) -> Self {
+    pub fn new(conn: &'a mut UserAgentConnection, transport: &'a mut T) -> Self {
         Self { conn, transport }
     }
 }
 
 impl<T> AuthStateMachineContext for AuthContext<'_, T>
 where
-    T: Bi<super::Inbound, Result<super::Outbound, Error>> + Send,
+    T: Bi<super::Inbound, Result<super::Outbound, Error>> + Send + ?Sized,
 {
     type Error = Error;
 

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

@@ -1,13 +1,25 @@
 use crate::{
-    actors::GlobalActors, crypto::integrity::Integrable, db, peers::client::ClientProfile,
+    actors::GlobalActors,
+    crypto::integrity::{self, Integrable},
+    db,
+    peers::client::ClientProfile,
 };
 use arbiter_crypto::authn;
 
+use arbiter_proto::transport::{Bi, Sender};
+pub use auth::authenticate;
+use kameo::actor::{ActorRef, Spawn as _};
+pub use session::UserAgentSession;
+use tokio::sync::oneshot;
+use tracing::warn;
+use vault_gate::VaultGate;
+
+use crate::crypto::integrity::hashing::Hashable;
+
 pub mod auth;
 pub mod session;
 pub mod vault_gate;
 
-
 #[derive(Debug, Clone, Hash)]
 pub struct Credentials {
     pub id: i32,
@@ -40,7 +52,6 @@ impl Hashable for AuthCredentials {
     }
 }
 
-
 impl Integrable for AuthCredentials {
     const KIND: &'static str = "useragent_credentials";
 }
@@ -52,6 +63,7 @@ pub enum OutOfBand {
     ClientConnectionCancel { pubkey: authn::PublicKey },
 }
 
+#[derive(Clone)]
 pub struct UserAgentConnection {
     pub(crate) db: db::DatabasePool,
     pub(crate) actors: GlobalActors,
@@ -63,9 +75,101 @@ impl UserAgentConnection {
     }
 }
 
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("authentication failed: {0:?}")]
+    Auth(auth::Error),
+    #[error("vault gate failed: {0}")]
+    VaultGate(#[from] vault_gate::Error),
+    #[error("transport closed unexpectedly")]
+    Transport,
+    #[error("internal: {0}")]
+    Internal(String),
+}
 
+impl From<auth::Error> for Error {
+    fn from(err: auth::Error) -> Self {
+        Self::Auth(err)
+    }
+}
 
-pub use auth::authenticate;
+pub async fn start<T>(
-pub use session::UserAgentSession;
+    props: &mut UserAgentConnection,
+    mut transport: T,
+    oob_sender: Box<dyn Sender<OutOfBand>>,
+) -> Result<ActorRef<UserAgentSession>, Error>
+where
+    T: Bi<auth::Inbound, Result<auth::Outbound, auth::Error>> + Send,
+    T: Bi<vault_gate::Inbound, Result<vault_gate::Outbound, vault_gate::Error>> + Send,
+{
+    let auth_creds = authenticate(props, &mut transport).await?;
 
-use crate::crypto::integrity::hashing::Hashable;
+    let creds = if integrity::is_signing_available(&props.actors.vault)
+        .await
+        .unwrap_or(false)
+    {
+        auth_creds.creds
+    } else {
+        run_vault_gate(props, &mut transport, auth_creds).await?
+    };
+
+    Ok(UserAgentSession::spawn(UserAgentSession::new(
+        props.clone(),
+        creds,
+        oob_sender,
+    )))
+}
+
+async fn run_vault_gate<T>(
+    props: &UserAgentConnection,
+    transport: &mut T,
+    auth_creds: AuthCredentials,
+) -> Result<Credentials, Error>
+where
+    T: Bi<vault_gate::Inbound, Result<vault_gate::Outbound, vault_gate::Error>> + Send + ?Sized,
+{
+    let (promotion_tx, mut promotion_rx) = oneshot::channel();
+    let gate = VaultGate::spawn(VaultGate::new(
+        auth_creds,
+        props.actors.clone(),
+        props.db.clone(),
+        promotion_tx,
+    ));
+
+    let result = loop {
+        tokio::select! {
+            promotion = &mut promotion_rx => {
+                break match promotion {
+                    Ok(Ok(creds)) => Ok(creds),
+                    Ok(Err(err)) => Err(Error::VaultGate(err)),
+                    Err(_) => Err(Error::Internal(
+                        "vault gate promotion channel closed".into(),
+                    )),
+                };
+            }
+
+            inbound = transport.recv() => {
+                let Some(inbound) = inbound else {
+                    break Err(Error::Transport);
+                };
+
+                match gate.ask(inbound).await {
+                    Ok(outbound) => {
+                        if transport.send(Ok(outbound)).await.is_err() {
+                            break Err(Error::Transport);
+                        }
+                    }
+                    Err(err) => {
+                        warn!(?err, "VaultGate failed to handle message");
+                        break Err(Error::Internal(format!(
+                            "vault gate ask failed: {err:?}"
+                        )));
+                    }
+                }
+            }
+        }
+    };
+
+    gate.kill();
+    result
+}

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

@@ -119,7 +119,7 @@ impl VaultGate {
     }
 }
 
-#[messages]
+#[messages(messages = Inbound, replies = Outbound)]
 impl VaultGate {
     #[message]
     pub async fn handle_handshake(
@@ -185,7 +185,7 @@ impl VaultGate {
     }
 
     #[message]
-    pub(crate) async fn handle_bootstrap_encrypted_key(
+    pub async fn handle_bootstrap_encrypted_key(
         &mut self,
         nonce: Vec<u8>,
         ciphertext: Vec<u8>,

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

@@ -2,17 +2,16 @@ 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 },
+    ReadyForExchange {
+        server_key: PublicKey,
+        secret: SharedSecret,
+    },
 }

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

@@ -42,11 +42,11 @@ pub async fn test_bootstrap_token_auth() {
         .unwrap();
     let token = actors.bootstrapper.ask(GetToken).await.unwrap().unwrap();
 
-    let (server_transport, mut test_transport) = ChannelTransport::new();
+    let (mut server_transport, mut test_transport) = ChannelTransport::new();
     let db_for_task = db.clone();
     let task = tokio::spawn(async move {
         let mut props = UserAgentConnection::new(db_for_task, actors);
-        auth::authenticate(&mut props, server_transport).await
+        auth::authenticate(&mut props, &mut server_transport).await
     });
 
     let new_key = MlDsa87::key_gen(&mut rand::rng());
@@ -84,11 +84,11 @@ pub async fn test_bootstrap_invalid_token_auth() {
     let db = db::create_test_pool().await;
     let actors = GlobalActors::spawn(db.clone()).await.unwrap();
 
-    let (server_transport, mut test_transport) = ChannelTransport::new();
+    let (mut server_transport, mut test_transport) = ChannelTransport::new();
     let db_for_task = db.clone();
     let task = tokio::spawn(async move {
         let mut props = UserAgentConnection::new(db_for_task, actors);
-        auth::authenticate(&mut props, server_transport).await
+        auth::authenticate(&mut props, &mut server_transport).await
     });
 
     let new_key = MlDsa87::key_gen(&mut rand::rng());
@@ -157,11 +157,11 @@ pub async fn test_challenge_auth() {
         .unwrap();
     }
 
-    let (server_transport, mut test_transport) = ChannelTransport::new();
+    let (mut server_transport, mut test_transport) = ChannelTransport::new();
     let db_for_task = db.clone();
     let task = tokio::spawn(async move {
         let mut props = UserAgentConnection::new(db_for_task, actors);
-        auth::authenticate(&mut props, server_transport).await
+        auth::authenticate(&mut props, &mut server_transport).await
     });
 
     test_transport
@@ -234,11 +234,11 @@ pub async fn test_challenge_auth_rejects_integrity_tag_mismatch_when_unsealed()
             .unwrap();
     }
 
-    let (server_transport, mut test_transport) = ChannelTransport::new();
+    let (mut server_transport, mut test_transport) = ChannelTransport::new();
     let db_for_task = db.clone();
     let task = tokio::spawn(async move {
         let mut props = UserAgentConnection::new(db_for_task, actors);
-        auth::authenticate(&mut props, server_transport).await
+        auth::authenticate(&mut props, &mut server_transport).await
     });
 
     test_transport
@@ -298,11 +298,11 @@ pub async fn test_challenge_auth_rejects_invalid_signature() {
         .unwrap();
     }
 
-    let (server_transport, mut test_transport) = ChannelTransport::new();
+    let (mut server_transport, mut test_transport) = ChannelTransport::new();
     let db_for_task = db.clone();
     let task = tokio::spawn(async move {
         let mut props = UserAgentConnection::new(db_for_task, actors);
-        auth::authenticate(&mut props, server_transport).await
+        auth::authenticate(&mut props, &mut server_transport).await
     });
 
     test_transport