From 832d8844576de07be27236078447e0feabdc2d64 Mon Sep 17 00:00:00 2001 From: hdbg Date: Fri, 13 Feb 2026 13:41:01 +0100 Subject: [PATCH] feat(auth): implement bootstrap token auth handling --- protobufs/auth.proto | 6 +- server/crates/arbiter-server/Cargo.toml | 1 + .../arbiter-server/src/actors/user_agent.rs | 98 +++++++++---------- server/crates/arbiter-server/src/lib.rs | 63 ++++++++++-- 4 files changed, 102 insertions(+), 66 deletions(-) diff --git a/protobufs/auth.proto b/protobufs/auth.proto index c290a66..b0b3cd8 100644 --- a/protobufs/auth.proto +++ b/protobufs/auth.proto @@ -5,10 +5,8 @@ package arbiter.auth; import "google/protobuf/timestamp.proto"; message AuthChallengeRequest { - oneof payload { - bytes pubkey = 1; - string bootstrap_token = 2; - } + bytes pubkey = 1; + optional string bootstrap_token = 2; } message AuthChallenge { diff --git a/server/crates/arbiter-server/Cargo.toml b/server/crates/arbiter-server/Cargo.toml index 6317470..f0f3e94 100644 --- a/server/crates/arbiter-server/Cargo.toml +++ b/server/crates/arbiter-server/Cargo.toml @@ -34,3 +34,4 @@ memsafe = "0.4.0" chacha20poly1305 = { version = "0.10.1", features = ["std"] } zeroize = { version = "1.8.2", features = ["std", "simd"] } kameo.workspace = true + diff --git a/server/crates/arbiter-server/src/actors/user_agent.rs b/server/crates/arbiter-server/src/actors/user_agent.rs index 57b23b7..62fdd95 100644 --- a/server/crates/arbiter-server/src/actors/user_agent.rs +++ b/server/crates/arbiter-server/src/actors/user_agent.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use arbiter_proto::{ proto::{ UserAgentRequest, UserAgentResponse, @@ -8,10 +10,12 @@ use arbiter_proto::{ }, transport::Bi, }; +use ed25519_dalek::VerifyingKey; use futures::StreamExt; use kameo::{Actor, message::StreamMessage, messages, prelude::Context}; use secrecy::{ExposeSecret, SecretBox}; use tokio::sync::mpsc; +use tokio::sync::mpsc::Sender; use tonic::{Status, transport::Server}; use tracing::error; @@ -59,78 +63,64 @@ impl UserAgentStateMachineContext for ServerContext { pub struct UserAgentActor { context: ServerContext, state: UserAgentStateMachine, - tx: mpsc::Sender>, + rx: Sender>, } impl UserAgentActor { pub(crate) fn new( context: ServerContext, - tx: mpsc::Sender>, + rx: Sender>, ) -> Self { Self { context: context.clone(), state: UserAgentStateMachine::new(context), - tx, + rx, } } - async fn handle_grpc( + async fn auth_with_bootstrap_token( &mut self, - msg: UserAgentRequest, - ctx: &mut Context, - ) -> Result { - let Some(msg) = msg.payload else { - error!(actor = "useragent", "Received message with no payload"); - ctx.stop(); - return Err(tonic::Status::invalid_argument( - "Message payload is required", - )); - }; - - let UserAgentRequestPayload::AuthMessage(ClientMessage { - payload: Some(client_message), - }) = msg - else { - error!( - actor = "useragent", - "Received unexpected message type during authentication" - ); - ctx.stop(); - return Err(tonic::Status::invalid_argument( - "Unexpected message type during authentication", - )); - }; - - match client_message { - ClientAuthPayload::AuthChallengeRequest(AuthChallengeRequest { - payload: Some(payload), - }) => match payload { - auth::auth_challenge_request::Payload::Pubkey(items) => todo!(), - auth::auth_challenge_request::Payload::BootstrapToken(_) => todo!(), - }, - ClientAuthPayload::AuthChallengeSolution(_auth_challenge_solution) => todo!(), - _ => { - error!( - actor = "useragent", - "Received unexpected message type during authentication" - ); - ctx.stop(); - return Err(tonic::Status::invalid_argument( - "Unexpected message type during authentication", - )); - } - } + pubkey: ed25519_dalek::VerifyingKey, + token: String, + ) -> Result { + todo!() } } +type Output = Result; + #[messages] impl UserAgentActor { #[message(ctx)] - pub async fn grpc(&mut self, msg: UserAgentRequest, ctx: &mut Context) { - let result = self.handle_grpc(msg, ctx).await; - self.tx.send(result).await.unwrap_or_else(|e| { - error!(handler = "useragent", "Failed to send response: {}", e); - ctx.stop(); - }); + async fn handle_auth_challenge_request( + &mut self, + req: AuthChallengeRequest, + ctx: &mut Context, + ) -> Output { + let pubkey = req.pubkey.as_array().ok_or(Status::invalid_argument( + "Expected pubkey to have specific length", + ))?; + let pubkey = VerifyingKey::from_bytes(pubkey).map_err(|err| { + error!(?pubkey, "Failed to convert to VerifyingKey"); + Status::invalid_argument("Failed to convert pubkey to VerifyingKey") + })?; + + if let Some(token) = req.bootstrap_token { + return self + .auth_with_bootstrap_token(pubkey, token) + .await + .map_err(|_| Status::internal("Failed to authenticate with bootstrap token")); + } + + todo!() + } + + #[message(ctx)] + async fn handle_auth_challenge_solution( + &mut self, + _solution: auth::AuthChallengeSolution, + ctx: &mut Context, + ) -> Output { + todo!() } } diff --git a/server/crates/arbiter-server/src/lib.rs b/server/crates/arbiter-server/src/lib.rs index e9d58b6..afb7d56 100644 --- a/server/crates/arbiter-server/src/lib.rs +++ b/server/crates/arbiter-server/src/lib.rs @@ -1,9 +1,18 @@ #![allow(unused)] +use std::sync::Arc; + use tracing::error; use arbiter_proto::{ - proto::{ClientRequest, ClientResponse, UserAgentRequest, UserAgentResponse}, + proto::{ + ClientRequest, ClientResponse, UserAgentRequest, UserAgentResponse, + auth::{ + self, AuthChallengeRequest, ClientMessage, client_message::Payload as ClientAuthPayload, + }, + user_agent_request::Payload as UserAgentRequestPayload, + user_agent_request::*, + }, transport::BiStream, }; use async_trait::async_trait; @@ -61,15 +70,53 @@ impl arbiter_proto::proto::arbiter_service_server::ArbiterService for Server { let mut req_stream = request.into_inner(); let (tx, rx) = mpsc::channel(DEFAULT_CHANNEL_SIZE); - let actor = UserAgentActor::spawn(UserAgentActor::new(self.context.clone(), tx)); + let actor = UserAgentActor::spawn(UserAgentActor::new(self.context.clone(), tx.clone())); tokio::task::spawn(async move { - while let Some(Ok(req)) = req_stream.next().await && actor.is_alive() { - if actor.tell(user_agent::Grpc {msg: req}).await.is_err() { - error!("Failed to send message to UserAgentActor"); - break; - } - } + while let Some(Ok(req)) = req_stream.next().await + && actor.is_alive() + { + let Some(msg) = req.payload else { + error!(actor = "useragent", "Received message with no payload"); + actor.kill(); + tx.send(Err(Status::invalid_argument( + "Expected message with payload", + ))) + .await; + return; + }; + + let UserAgentRequestPayload::AuthMessage(ClientMessage { + payload: Some(client_message), + }) = msg + else { + error!( + actor = "useragent", + "Received unexpected message type during authentication" + ); + actor.kill(); + tx.send(Err(Status::invalid_argument( + "Expected AuthMessage with ClientMessage payload", + ))) + .await; + return; + }; + + match client_message { + ClientAuthPayload::AuthChallengeRequest(req) => {} + ClientAuthPayload::AuthChallengeSolution(_auth_challenge_solution) => todo!(), + _ => { + error!(actor = "useragent", "Received unexpected message type"); + actor.kill(); + tx.send(Err(Status::invalid_argument( + "Expected AuthMessage with ClientMessage payload", + ))) + .await; + return; + } + } + todo!() + } }); Ok(Response::new(ReceiverStream::new(rx)))