feat: actors experiment

This commit is contained in:
hdbg
2026-02-13 12:11:56 +01:00
parent bbbb4feaa0
commit 208bbbd540
13 changed files with 245 additions and 87 deletions

View File

@@ -5,7 +5,10 @@ package arbiter.auth;
import "google/protobuf/timestamp.proto"; import "google/protobuf/timestamp.proto";
message AuthChallengeRequest { message AuthChallengeRequest {
bytes pubkey = 1; oneof payload {
bytes pubkey = 1;
string bootstrap_token = 2;
}
} }
message AuthChallenge { message AuthChallenge {

37
server/Cargo.lock generated
View File

@@ -61,6 +61,7 @@ version = "0.1.0"
dependencies = [ dependencies = [
"bytes", "bytes",
"futures", "futures",
"kameo",
"prost", "prost",
"prost-build", "prost-build",
"prost-derive", "prost-derive",
@@ -89,6 +90,7 @@ dependencies = [
"ed25519", "ed25519",
"ed25519-dalek", "ed25519-dalek",
"futures", "futures",
"kameo",
"memsafe", "memsafe",
"miette", "miette",
"rand", "rand",
@@ -104,6 +106,7 @@ dependencies = [
"tokio-stream", "tokio-stream",
"tonic", "tonic",
"tracing", "tracing",
"zeroize",
] ]
[[package]] [[package]]
@@ -714,6 +717,12 @@ version = "1.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813"
[[package]]
name = "dyn-clone"
version = "1.0.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555"
[[package]] [[package]]
name = "ed25519" name = "ed25519"
version = "3.0.0-rc.4" version = "3.0.0-rc.4"
@@ -1216,6 +1225,33 @@ dependencies = [
"wasm-bindgen", "wasm-bindgen",
] ]
[[package]]
name = "kameo"
version = "0.19.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9c4af7638c67029fd6821d02813c3913c803784648725d4df4082c9b91d7cbb1"
dependencies = [
"downcast-rs",
"dyn-clone",
"futures",
"kameo_macros",
"serde",
"tokio",
"tracing",
]
[[package]]
name = "kameo_macros"
version = "0.19.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a13c324e2d8c8e126e63e66087448b4267e263e6cb8770c56d10a9d0d279d9e2"
dependencies = [
"heck",
"proc-macro2",
"quote",
"syn 2.0.114",
]
[[package]] [[package]]
name = "lazy_static" name = "lazy_static"
version = "1.5.0" version = "1.5.0"
@@ -2419,6 +2455,7 @@ dependencies = [
"signal-hook-registry", "signal-hook-registry",
"socket2", "socket2",
"tokio-macros", "tokio-macros",
"tracing",
"windows-sys 0.61.2", "windows-sys 0.61.2",
] ]

View File

@@ -25,3 +25,4 @@ thiserror = "2.0.18"
async-trait = "0.1.89" async-trait = "0.1.89"
futures = "0.3.31" futures = "0.3.31"
tokio-stream = { version = "0.1.18", features = ["full"] } tokio-stream = { version = "0.1.18", features = ["full"] }
kameo = "0.19.2"

View File

@@ -14,6 +14,7 @@ tonic-prost = "0.14.3"
rkyv = "0.8.15" rkyv = "0.8.15"
tokio.workspace = true tokio.workspace = true
futures.workspace = true futures.workspace = true
kameo.workspace = true

View File

@@ -4,6 +4,7 @@ static PROTOBUF_DIR: &str = "../../../protobufs";
fn main() -> Result<(), Box<dyn std::error::Error>> { fn main() -> Result<(), Box<dyn std::error::Error>> {
configure() configure()
.message_attribute(".", "#[derive(::kameo::Reply)]")
.compile_protos( .compile_protos(
&[ &[
format!("{}/arbiter.proto", PROTOBUF_DIR), format!("{}/arbiter.proto", PROTOBUF_DIR),
@@ -11,6 +12,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
], ],
&[PROTOBUF_DIR.to_string()], &[PROTOBUF_DIR.to_string()],
) )
.unwrap(); .unwrap();
Ok(()) Ok(())
} }

View File

@@ -32,3 +32,5 @@ chrono.workspace = true
bytes = "1.11.1" bytes = "1.11.1"
memsafe = "0.4.0" memsafe = "0.4.0"
chacha20poly1305 = { version = "0.10.1", features = ["std"] } chacha20poly1305 = { version = "0.10.1", features = ["std"] }
zeroize = { version = "1.8.2", features = ["std", "simd"] }
kameo.workspace = true

View File

@@ -0,0 +1,136 @@
use arbiter_proto::{
proto::{
UserAgentRequest, UserAgentResponse,
auth::{
self, AuthChallengeRequest, ClientMessage, client_message::Payload as ClientAuthPayload,
},
user_agent_request::Payload as UserAgentRequestPayload,
},
transport::Bi,
};
use futures::StreamExt;
use kameo::{Actor, message::StreamMessage, messages, prelude::Context};
use secrecy::{ExposeSecret, SecretBox};
use tokio::sync::mpsc;
use tonic::{Status, transport::Server};
use tracing::error;
use crate::ServerContext;
#[derive(Debug)]
pub struct ChallengeContext {
challenge: auth::AuthChallenge,
key: ed25519_dalek::SigningKey,
}
smlang::statemachine!(
name: UserAgent,
derive_states: [Debug],
transitions: {
*Init + ReceivedRequest(ed25519_dalek::VerifyingKey) [async check_key_existence] / provide_challenge = WaitingForChallengeSolution(ChallengeContext),
Init + ReceivedBootstrapToken(String) = Authenticated,
WaitingForChallengeSolution(ChallengeContext) + ReceivedGoodSolution = Authenticated,
WaitingForChallengeSolution(ChallengeContext) + ReceivedBadSolution = Error,
}
);
impl UserAgentStateMachineContext for ServerContext {
#[allow(missing_docs)]
#[allow(clippy::unused_unit)]
fn provide_challenge(
&mut self,
event_data: ed25519_dalek::VerifyingKey,
) -> Result<ChallengeContext, ()> {
todo!()
}
#[allow(missing_docs)]
#[allow(clippy::result_unit_err)]
async fn check_key_existence(
&self,
event_data: &ed25519_dalek::VerifyingKey,
) -> Result<bool, ()> {
todo!()
}
}
#[derive(Actor)]
pub struct UserAgentActor {
context: ServerContext,
state: UserAgentStateMachine<ServerContext>,
tx: mpsc::Sender<Result<UserAgentResponse, tonic::Status>>,
}
impl UserAgentActor {
pub(crate) fn new(
context: ServerContext,
tx: mpsc::Sender<Result<UserAgentResponse, tonic::Status>>,
) -> Self {
Self {
context: context.clone(),
state: UserAgentStateMachine::new(context),
tx,
}
}
async fn handle_grpc(
&mut self,
msg: UserAgentRequest,
ctx: &mut Context<Self, ()>,
) -> Result<UserAgentResponse, tonic::Status> {
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",
));
}
}
}
}
#[messages]
impl UserAgentActor {
#[message(ctx)]
pub async fn grpc(&mut self, msg: UserAgentRequest, ctx: &mut Context<Self, ()>) {
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();
});
}
}

View File

@@ -11,6 +11,7 @@ use tokio::sync::RwLock;
use crate::{ use crate::{
context::{ context::{
bootstrap::generate_token,
lease::LeaseHandler, lease::LeaseHandler,
tls::{TlsDataRaw, TlsManager}, tls::{TlsDataRaw, TlsManager},
}, },
@@ -21,11 +22,9 @@ use crate::{
}, },
}; };
pub(crate) mod bootstrap;
pub(crate) mod lease; pub(crate) mod lease;
pub(crate) mod tls; pub(crate) mod tls;
pub(crate) mod bootstrap {
}
#[derive(Error, Debug, Diagnostic)] #[derive(Error, Debug, Diagnostic)]
pub enum InitError { pub enum InitError {
@@ -44,6 +43,10 @@ pub enum InitError {
#[error("TLS initialization failed: {0}")] #[error("TLS initialization failed: {0}")]
#[diagnostic(code(arbiter_server::init::tls_init))] #[diagnostic(code(arbiter_server::init::tls_init))]
Tls(#[from] tls::TlsInitError), Tls(#[from] tls::TlsInitError),
#[error("I/O Error: {0}")]
#[diagnostic(code(arbiter_server::init::io))]
Io(#[from] std::io::Error),
} }
// TODO: Placeholder for secure root key cell implementation // TODO: Placeholder for secure root key cell implementation
@@ -52,7 +55,7 @@ pub struct KeyStorage;
statemachine! { statemachine! {
name: Server, name: Server,
transitions: { transitions: {
*NotBootstrapped + Bootstrapped = Sealed, *NotBootstrapped(String) + Bootstrapped = Sealed,
Sealed + Unsealed(KeyStorage) / move_key = Ready(KeyStorage), Sealed + Unsealed(KeyStorage) / move_key = Ready(KeyStorage),
Ready(KeyStorage) + Sealed / dispose_key = Sealed, Ready(KeyStorage) + Sealed / dispose_key = Sealed,
} }
@@ -135,7 +138,9 @@ impl ServerContext {
drop(conn); drop(conn);
let mut state = ServerStateMachine::new(_Context); let bootstrap_token = generate_token().await?;
let mut state = ServerStateMachine::new(_Context, bootstrap_token);
if let Some(settings) = &settings if let Some(settings) = &settings
&& settings.root_key_id.is_some() && settings.root_key_id.is_some()
@@ -144,7 +149,6 @@ impl ServerContext {
let _ = state.process_event(ServerEvents::Bootstrapped); let _ = state.process_event(ServerEvents::Bootstrapped);
} }
Ok(Self(Arc::new(_ServerContextInner { Ok(Self(Arc::new(_ServerContextInner {
db, db,
rng, rng,

View File

@@ -0,0 +1,30 @@
use arbiter_proto::{BOOTSTRAP_TOKEN_PATH, home_path};
use diesel::{QueryDsl, dsl::exists, select};
use diesel_async::RunQueryDsl;
use memsafe::MemSafe;
use miette::Diagnostic;
use rand::{RngExt, distr::StandardUniform, make_rng, rngs::StdRng};
use secrecy::SecretString;
use thiserror::Error;
use tracing::info;
use zeroize::{Zeroize, Zeroizing};
use crate::db::{self, schema};
const TOKEN_LENGTH: usize = 64;
pub async fn generate_token() -> Result<String, std::io::Error> {
let rng: StdRng = make_rng();
let token: String = rng
.sample_iter::<char, _>(StandardUniform)
.take(TOKEN_LENGTH)
.fold(Default::default(), |mut accum, char| {
accum += char.to_string().as_str();
accum
});
tokio::fs::write(home_path()?.join(BOOTSTRAP_TOKEN_PATH), token.as_str()).await?;
Ok(token)
}

View File

@@ -1,69 +0,0 @@
use arbiter_proto::{
proto::{
UserAgentRequest, UserAgentResponse,
auth::{
self, AuthChallengeRequest, ClientMessage, client_message::Payload as ClientAuthPayload
},
user_agent_request::Payload as UserAgentRequestPayload,
},
transport::Bi,
};
use futures::StreamExt;
use tracing::error;
use crate::ServerContext;
smlang::statemachine!(
name: UserAgentAuth,
derive_states: [Debug],
derive_events: [Clone, Debug],
transitions: {
*Init + ReceivedRequest(ed25519_dalek::VerifyingKey) / provide_challenge = WaitingForChallengeSolution(auth::AuthChallenge),
WaitingForChallengeSolution(auth::AuthChallenge) + ReceivedGoodSolution = Authenticated,
WaitingForChallengeSolution(auth::AuthChallenge) + ReceivedBadSolution = Error,
}
);
impl UserAgentAuthStateMachineContext for ServerContext {
#[allow(missing_docs)]
#[allow(clippy::unused_unit)]
fn provide_challenge< >(&mut self,_event_data:ed25519_dalek::VerifyingKey) -> Result<auth::AuthChallenge,()> {
todo!()
}
}
pub(crate) async fn handle_user_agent(
context: ServerContext,
mut bistream: impl Bi<UserAgentRequest, UserAgentResponse> + Unpin,
) {
let auth_sm = UserAgentAuthStateMachine::new(context);
while let Some(Ok(msg)) = bistream.next().await
&& auth_sm.state() != &UserAgentAuthStates::Authenticated
{
let Some(msg) = msg.payload else {
error!(handler = "useragent", "Received message with no payload");
return;
};
let UserAgentRequestPayload::AuthMessage(ClientMessage {
payload: Some(client_message),
}) = msg
else {
error!(
handler = "useragent",
"Received unexpected message type during authentication"
);
return;
};
match client_message {
ClientAuthPayload::AuthChallengeRequest(auth_challenge_request) => {
let AuthChallengeRequest { pubkey } = auth_challenge_request;
},
ClientAuthPayload::AuthChallengeSolution(_auth_challenge_solution) => todo!(),
}
}
}

View File

@@ -1,23 +1,30 @@
#![allow(unused)] #![allow(unused)]
use tracing::error;
use arbiter_proto::{ use arbiter_proto::{
proto::{ClientRequest, ClientResponse, UserAgentRequest, UserAgentResponse}, proto::{ClientRequest, ClientResponse, UserAgentRequest, UserAgentResponse},
transport::BiStream, transport::BiStream,
}; };
use async_trait::async_trait; use async_trait::async_trait;
use futures::StreamExt;
use kameo::actor::Spawn;
use tokio_stream::wrappers::ReceiverStream; use tokio_stream::wrappers::ReceiverStream;
use tokio::sync::mpsc; use tokio::sync::mpsc;
use tonic::{Request, Response, Status}; use tonic::{Request, Response, Status};
use crate::{ use crate::{
handlers::{client::handle_client, user_agent::handle_user_agent}, actors::{
client::handle_client,
user_agent::{self, UserAgentActor},
},
context::ServerContext, context::ServerContext,
}; };
mod db; pub mod actors;
pub mod handlers;
mod context; mod context;
mod db;
const DEFAULT_CHANNEL_SIZE: usize = 1000; const DEFAULT_CHANNEL_SIZE: usize = 1000;
@@ -51,16 +58,20 @@ impl arbiter_proto::proto::arbiter_service_server::ArbiterService for Server {
&self, &self,
request: Request<tonic::Streaming<UserAgentRequest>>, request: Request<tonic::Streaming<UserAgentRequest>>,
) -> Result<Response<Self::UserAgentStream>, Status> { ) -> Result<Response<Self::UserAgentStream>, Status> {
let req_stream = request.into_inner(); let mut req_stream = request.into_inner();
let (tx, rx) = mpsc::channel(DEFAULT_CHANNEL_SIZE); let (tx, rx) = mpsc::channel(DEFAULT_CHANNEL_SIZE);
tokio::spawn(handle_user_agent( let actor = UserAgentActor::spawn(UserAgentActor::new(self.context.clone(), tx));
self.context.clone(),
BiStream { tokio::task::spawn(async move {
request_stream: req_stream, while let Some(Ok(req)) = req_stream.next().await && actor.is_alive() {
response_sender: tx, if actor.tell(user_agent::Grpc {msg: req}).await.is_err() {
}, error!("Failed to send message to UserAgentActor");
)); break;
}
}
});
Ok(Response::new(ReceiverStream::new(rx))) Ok(Response::new(ReceiverStream::new(rx)))
} }
} }