MarketTakers/arbiter
5
0
Fork 0
Code Issues 27 Pull Requests 5 Packages Projects Releases Wiki Activity

feat(proto): add URL parsing and TLS certificate management

Browse Source
This commit is contained in:
hdbg
2026-02-17 14:00:39 +01:00
parent 7438d62695
commit cbbe1f8881
21 changed files with 814 additions and 204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
server/crates/arbiter-server/src/actors/bootstrap.rs
View File

@@ -1,28 +1,31 @@
use arbiter_proto::{BOOTSTRAP_TOKEN_PATH, home_path};
use arbiter_proto::{BOOTSTRAP_PATH, home_path};
use diesel::QueryDsl;
use diesel_async::RunQueryDsl;
use kameo::{Actor, messages};
use miette::Diagnostic;
use rand::{RngExt, distr::StandardUniform, make_rng, rngs::StdRng};
use rand::{
RngExt,
distr::{Alphanumeric},
make_rng,
rngs::StdRng,
};
use thiserror::Error;
use tracing::info;
use crate::db::{self, DatabasePool, 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| {
let token: String = rng.sample_iter(Alphanumeric).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?;
tokio::fs::write(home_path()?.join(BOOTSTRAP_PATH), token.as_str()).await?;
Ok(token)
}
@@ -58,10 +61,9 @@ impl Bootstrapper {
drop(conn);
let token = if row_count == 0 {
let token = generate_token().await?;
info!(%token, "Generated bootstrap token");
tokio::fs::write(home_path()?.join(BOOTSTRAP_TOKEN_PATH), token.as_str()).await?;
Some(token)
} else {
None

19
server/crates/arbiter-server/src/actors/keyholder/mod.rs
View File

@@ -345,30 +345,15 @@ impl KeyHolder {
#[cfg(test)]
mod tests {
use diesel::SelectableHelper;
use diesel::dsl::insert_into;
use diesel_async::RunQueryDsl;
use memsafe::MemSafe;
use crate::db::{self, models::ArbiterSetting};
use crate::db::{self};
use super::*;
async fn seed_settings(pool: &db::DatabasePool) {
let mut conn = pool.get().await.unwrap();
insert_into(schema::arbiter_settings::table)
.values(&ArbiterSetting {
id: 1,
root_key_id: None,
cert_key: vec![],
cert: vec![],
})
.execute(&mut conn)
.await
.unwrap();
}
async fn bootstrapped_actor(db: &db::DatabasePool) -> KeyHolder {
seed_settings(db).await;
let mut actor = KeyHolder::new(db.clone()).await.unwrap();
let seal_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
actor.bootstrap(seal_key).await.unwrap();

53
server/crates/arbiter-server/src/context/mod.rs
View File

@@ -1,14 +1,12 @@
use std::sync::Arc;
use diesel::OptionalExtension as _;
use diesel_async::RunQueryDsl as _;
use miette::Diagnostic;
use thiserror::Error;
use crate::{
actors::GlobalActors,
context::tls::{TlsDataRaw, TlsManager},
db::{self, models::ArbiterSetting, schema::arbiter_settings},
context::tls::TlsManager,
db::{self},
};
pub mod tls;
@@ -29,7 +27,7 @@ pub enum InitError {
#[error("TLS initialization failed: {0}")]
#[diagnostic(code(arbiter_server::init::tls_init))]
Tls(#[from] tls::TlsInitError),
Tls(#[from] tls::InitError),
#[error("Actor spawn failed: {0}")]
#[diagnostic(code(arbiter_server::init::actor_spawn))]
@@ -57,54 +55,11 @@ impl std::ops::Deref for ServerContext {
}
impl ServerContext {
async fn load_tls(
db: &mut db::DatabaseConnection,
settings: Option<&ArbiterSetting>,
) -> Result<TlsManager, InitError> {
match &settings {
Some(settings) => {
let tls_data_raw = TlsDataRaw {
cert: settings.cert.clone(),
key: settings.cert_key.clone(),
};
Ok(TlsManager::new(Some(tls_data_raw)).await?)
}
None => {
let tls = TlsManager::new(None).await?;
let tls_data_raw = tls.bytes();
diesel::insert_into(arbiter_settings::table)
.values(&ArbiterSetting {
id: 1,
root_key_id: None,
cert_key: tls_data_raw.key,
cert: tls_data_raw.cert,
})
.execute(db)
.await?;
Ok(tls)
}
}
}
pub async fn new(db: db::DatabasePool) -> Result<Self, InitError> {
let mut conn = db.get().await?;
let settings = arbiter_settings::table
.first::<ArbiterSetting>(&mut conn)
.await
.optional()?;
let tls = Self::load_tls(&mut conn, settings.as_ref()).await?;
drop(conn);
Ok(Self(Arc::new(_ServerContextInner {
actors: GlobalActors::spawn(db.clone()).await?,
tls: TlsManager::new(db.clone()).await?,
db,
tls,
})))
}
}

253
server/crates/arbiter-server/src/context/tls.rs
View File

@@ -1,12 +1,36 @@
use std::string::FromUtf8Error;
use diesel::{ExpressionMethods as _, QueryDsl, SelectableHelper as _};
use diesel_async::{AsyncConnection, RunQueryDsl};
use miette::Diagnostic;
use rcgen::{Certificate, KeyPair};
use rustls::pki_types::CertificateDer;
use pem::Pem;
use rcgen::{
BasicConstraints, Certificate, CertificateParams, CertifiedIssuer, DistinguishedName, DnType,
IsCa, Issuer, KeyPair, KeyUsagePurpose,
};
use rustls::pki_types::{pem::PemObject};
use thiserror::Error;
use tonic::transport::CertificateDer;
use crate::db::{
self,
models::{NewTlsHistory, TlsHistory},
schema::{
arbiter_settings,
tls_history::{self},
},
};
const ENCODE_CONFIG: pem::EncodeConfig = {
let line_ending = match cfg!(target_family = "windows") {
true => pem::LineEnding::CRLF,
false => pem::LineEnding::LF,
};
pem::EncodeConfig::new().set_line_ending(line_ending)
};
#[derive(Error, Debug, Diagnostic)]
pub enum TlsInitError {
pub enum InitError {
#[error("Key generation error during TLS initialization: {0}")]
#[diagnostic(code(arbiter_server::tls_init::key_generation))]
KeyGeneration(#[from] rcgen::Error),
@@ -18,68 +42,211 @@ pub enum TlsInitError {
#[error("Key deserialization error: {0}")]
#[diagnostic(code(arbiter_server::tls_init::key_deserialization))]
KeyDeserializationError(rcgen::Error),
#[error("Database error during TLS initialization: {0}")]
#[diagnostic(code(arbiter_server::tls_init::database_error))]
DatabaseError(#[from] diesel::result::Error),
#[error("Pem deserialization error during TLS initialization: {0}")]
#[diagnostic(code(arbiter_server::tls_init::pem_deserialization))]
PemDeserializationError(#[from] rustls::pki_types::pem::Error),
#[error("Database pool acquire error during TLS initialization: {0}")]
#[diagnostic(code(arbiter_server::tls_init::database_pool_acquire))]
DatabasePoolAcquire(#[from] db::PoolError),
}
pub struct TlsData {
pub cert: CertificateDer<'static>,
pub keypair: KeyPair,
pub type PemCert = String;
pub fn encode_cert_to_pem(cert: &CertificateDer) -> PemCert {
pem::encode_config(
&Pem::new("CERTIFICATE", cert.to_vec()),
ENCODE_CONFIG,
)
}
pub struct TlsDataRaw {
pub cert: Vec<u8>,
pub key: Vec<u8>,
#[allow(unused)]
struct SerializedTls {
cert_pem: PemCert,
cert_key_pem: String,
}
impl TlsDataRaw {
pub fn serialize(cert: &TlsData) -> Self {
Self {
cert: cert.cert.as_ref().to_vec(),
key: cert.keypair.serialize_pem().as_bytes().to_vec(),
}
struct TlsCa {
issuer: Issuer<'static, KeyPair>,
cert: CertificateDer<'static>,
}
impl TlsCa {
fn generate() -> Result<Self, InitError> {
let keypair = KeyPair::generate()?;
let mut params = CertificateParams::new(["Arbiter Instance CA".into()])?;
params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained);
params.key_usages = vec![
KeyUsagePurpose::KeyCertSign,
KeyUsagePurpose::CrlSign,
KeyUsagePurpose::DigitalSignature,
];
let mut dn = DistinguishedName::new();
dn.push(DnType::CommonName, "Arbiter Instance CA");
params.distinguished_name = dn;
let certified_issuer = CertifiedIssuer::self_signed(params, keypair)?;
let cert_key_pem = certified_issuer.key().serialize_pem();
let issuer = Issuer::from_ca_cert_pem(
&certified_issuer.pem(),
KeyPair::from_pem(cert_key_pem.as_ref()).unwrap(),
)
.unwrap();
Ok(Self {
issuer,
cert: certified_issuer.der().clone(),
})
}
fn generate_leaf(&self) -> Result<TlsCert, InitError> {
let cert_key = KeyPair::generate()?;
let mut params = CertificateParams::new(["Arbiter Instance Leaf".into()])?;
params.is_ca = IsCa::NoCa;
params.key_usages = vec![
KeyUsagePurpose::DigitalSignature,
KeyUsagePurpose::KeyEncipherment,
];
let mut dn = DistinguishedName::new();
dn.push(DnType::CommonName, "Arbiter Instance Leaf");
params.distinguished_name = dn;
let new_cert = params.signed_by(&cert_key, &self.issuer)?;
Ok(TlsCert {
cert: new_cert,
cert_key,
})
}
pub fn deserialize(&self) -> Result<TlsData, TlsInitError> {
let cert = CertificateDer::from_slice(&self.cert).into_owned();
#[allow(unused)]
fn serialize(&self) -> Result<SerializedTls, InitError> {
let cert_key_pem = self.issuer.key().serialize_pem();
Ok(SerializedTls {
cert_pem: encode_cert_to_pem(&self.cert),
cert_key_pem,
})
}
let key = String::from_utf8(self.key.clone()).map_err(TlsInitError::KeyInvalidFormat)?;
let keypair = KeyPair::from_pem(&key).map_err(TlsInitError::KeyDeserializationError)?;
Ok(TlsData { cert, keypair })
#[allow(unused)]
fn try_deserialize(cert_pem: &str, cert_key_pem: &str) -> Result<Self, InitError> {
let keypair =
KeyPair::from_pem(cert_key_pem).map_err(InitError::KeyDeserializationError)?;
let issuer = Issuer::from_ca_cert_pem(cert_pem, keypair)?;
Ok(Self {
issuer,
cert: CertificateDer::from_pem_slice(cert_pem.as_bytes())?,
})
}
}
fn generate_cert(key: &KeyPair) -> Result<Certificate, rcgen::Error> {
let params =
rcgen::CertificateParams::new(vec!["arbiter.local".to_string(), "localhost".to_string()])?;
params.self_signed(key)
struct TlsCert {
cert: Certificate,
cert_key: KeyPair,
}
// TODO: Implement cert rotation
pub struct TlsManager {
data: TlsData,
cert: CertificateDer<'static>,
keypair: KeyPair,
ca_cert: CertificateDer<'static>,
_db: db::DatabasePool,
}
impl TlsManager {
pub async fn new(data: Option<TlsDataRaw>) -> Result<Self, TlsInitError> {
match data {
Some(raw) => {
let tls_data = raw.deserialize()?;
Ok(Self { data: tls_data })
}
None => {
let keypair = KeyPair::generate()?;
let cert = generate_cert(&keypair)?;
let tls_data = TlsData {
cert: cert.der().clone(),
keypair,
pub async fn generate_new(db: &db::DatabasePool) -> Result<Self, InitError> {
let ca = TlsCa::generate()?;
let new_cert = ca.generate_leaf()?;
{
let mut conn = db.get().await?;
conn.transaction(|conn| {
Box::pin(async {
let new_tls_history = NewTlsHistory {
cert: new_cert.cert.pem(),
cert_key: new_cert.cert_key.serialize_pem(),
ca_cert: encode_cert_to_pem(&ca.cert),
ca_key: ca.issuer.key().serialize_pem(),
};
let inserted_tls_history: i32 = diesel::insert_into(tls_history::table)
.values(&new_tls_history)
.returning(tls_history::id)
.get_result(conn)
.await?;
diesel::update(arbiter_settings::table)
.set(arbiter_settings::tls_id.eq(inserted_tls_history))
.execute(conn)
.await?;
Result::<_, diesel::result::Error>::Ok(())
})
})
.await?;
}
Ok(Self {
cert: new_cert.cert.der().clone(),
keypair: new_cert.cert_key,
ca_cert: ca.cert,
_db: db.clone(),
})
}
pub async fn new(db: db::DatabasePool) -> Result<Self, InitError> {
let cert_data: Option<TlsHistory> = {
let mut conn = db.get().await?;
arbiter_settings::table
.left_join(tls_history::table)
.select(Option::<TlsHistory>::as_select())
.first(&mut conn)
.await?
};
match cert_data {
Some(data) => {
let try_load = || -> Result<_, Box<dyn std::error::Error>> {
let keypair = KeyPair::from_pem(&data.cert_key)?;
let cert = CertificateDer::from_pem_slice(data.cert.as_bytes())?;
let ca_cert = CertificateDer::from_pem_slice(data.ca_cert.as_bytes())?;
Ok(Self {
cert,
keypair,
ca_cert,
_db: db.clone(),
})
};
Ok(Self { data: tls_data })
match try_load() {
Ok(manager) => Ok(manager),
Err(e) => {
eprintln!("Failed to load existing TLS certs: {e}. Generating new ones.");
Self::generate_new(&db).await
}
}
}
None => Self::generate_new(&db).await,
}
}
pub fn bytes(&self) -> TlsDataRaw {
TlsDataRaw::serialize(&self.data)
pub fn cert(&self) -> &CertificateDer<'static> {
&self.cert
}
pub fn ca_cert(&self) -> &CertificateDer<'static> {
&self.ca_cert
}
pub fn cert_pem(&self) -> PemCert {
encode_cert_to_pem(&self.cert)
}
pub fn key_pem(&self) -> String {
self.keypair.serialize_pem()
}
}

20
server/crates/arbiter-server/src/db/mod.rs
View File

@@ -24,23 +24,23 @@ const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");
#[derive(Error, Diagnostic, Debug)]
pub enum DatabaseSetupError {
#[error("Failed to determine home directory")]
#[diagnostic(code(arbiter::db::home_dir_error))]
#[diagnostic(code(arbiter::db::home_dir))]
HomeDir(std::io::Error),
#[error(transparent)]
#[diagnostic(code(arbiter::db::connection_error))]
#[diagnostic(code(arbiter::db::connection))]
Connection(diesel::ConnectionError),
#[error(transparent)]
#[diagnostic(code(arbiter::db::concurrency_error))]
#[diagnostic(code(arbiter::db::concurrency))]
ConcurrencySetup(diesel::result::Error),
#[error(transparent)]
#[diagnostic(code(arbiter::db::migration_error))]
#[diagnostic(code(arbiter::db::migration))]
Migration(Box<dyn std::error::Error + Send + Sync>),
#[error(transparent)]
#[diagnostic(code(arbiter::db::pool_error))]
#[diagnostic(code(arbiter::db::pool))]
Pool(#[from] PoolInitError),
}
@@ -91,12 +91,12 @@ fn initialize_database(url: &str) -> Result<(), DatabaseSetupError> {
#[tracing::instrument(level = "info")]
pub async fn create_pool(url: Option<&str>) -> Result<DatabasePool, DatabaseSetupError> {
let database_url = url.map(String::from).unwrap_or(format!(
"{}?mode=rwc",
(database_path()?
let database_url = url.map(String::from).unwrap_or(
database_path()?
.to_str()
.expect("database path is not valid UTF-8"))
));
.expect("database path is not valid UTF-8")
.to_string(),
);
initialize_database(&database_url)?;

26
server/crates/arbiter-server/src/db/models.rs
View File

@@ -1,7 +1,7 @@
#![allow(unused)]
#![allow(clippy::all)]
use crate::db::schema::{self, aead_encrypted, arbiter_settings, root_key_history};
use crate::db::schema::{self, aead_encrypted, arbiter_settings, root_key_history, tls_history};
use diesel::{prelude::*, sqlite::Sqlite};
use restructed::Models;
@@ -46,13 +46,29 @@ pub struct RootKeyHistory {
pub salt: Vec<u8>,
}
#[derive(Queryable, Debug, Insertable)]
#[derive(Models, Queryable, Debug, Insertable, Selectable)]
#[diesel(table_name = tls_history, check_for_backend(Sqlite))]
#[view(
NewTlsHistory,
derive(Insertable),
omit(id, created_at),
attributes_with = "deriveless"
)]
pub struct TlsHistory {
pub id: i32,
pub cert: String,
pub cert_key: String, // PEM Encoded private key
pub ca_cert: String, // PEM Encoded certificate for cert signing
pub ca_key: String, // PEM Encoded public key for cert signing
pub created_at: i32,
}
#[derive(Queryable, Debug, Insertable, Selectable)]
#[diesel(table_name = arbiter_settings, check_for_backend(Sqlite))]
pub struct ArbiterSetting {
pub struct ArbiterSettings {
pub id: i32,
pub root_key_id: Option<i32>, // references root_key_history.id
pub cert_key: Vec<u8>,
pub cert: Vec<u8>,
pub tls_id: Option<i32>, // references tls_history.id
}
#[derive(Queryable, Debug)]

16
server/crates/arbiter-server/src/db/schema.rs
View File

@@ -16,8 +16,7 @@ diesel::table! {
arbiter_settings (id) {
id -> Integer,
root_key_id -> Nullable<Integer>,
cert_key -> Binary,
cert -> Binary,
tls_id -> Nullable<Integer>,
}
}
@@ -43,6 +42,17 @@ diesel::table! {
}
}
diesel::table! {
tls_history (id) {
id -> Integer,
cert -> Text,
cert_key -> Text,
ca_cert -> Text,
ca_key -> Text,
created_at -> Integer,
}
}
diesel::table! {
useragent_client (id) {
id -> Integer,
@@ -55,11 +65,13 @@ diesel::table! {
diesel::joinable!(aead_encrypted -> root_key_history (associated_root_key_id));
diesel::joinable!(arbiter_settings -> root_key_history (root_key_id));
diesel::joinable!(arbiter_settings -> tls_history (tls_id));
diesel::allow_tables_to_appear_in_same_query!(
aead_encrypted,
arbiter_settings,
program_client,
root_key_history,
tls_history,
useragent_client,
);

31
server/crates/arbiter-server/src/main.rs
View File

@@ -1,7 +1,13 @@
use arbiter_proto::proto::arbiter_service_server::ArbiterServiceServer;
use arbiter_server::{Server, context::ServerContext, db};
use std::net::SocketAddr;
use arbiter_proto::{proto::arbiter_service_server::ArbiterServiceServer, url::ArbiterUrl};
use arbiter_server::{Server, actors::bootstrap::GetToken, context::ServerContext, db};
use miette::miette;
use tonic::transport::{Identity, ServerTlsConfig};
use tracing::info;
const PORT: u16 = 50051;
#[tokio::main]
async fn main() -> miette::Result<()> {
tracing_subscriber::fmt()
@@ -13,18 +19,31 @@ async fn main() -> miette::Result<()> {
info!("Starting arbiter server");
info!("Initializing database");
let db = db::create_pool(None).await?;
info!("Database ready");
info!("Initializing server context");
let context = ServerContext::new(db).await?;
info!("Server context ready");
let addr = "[::1]:50051".parse().expect("valid address");
let addr: SocketAddr = format!("127.0.0.1:{PORT}").parse().expect("valid address");
info!(%addr, "Starting gRPC server");
let url = ArbiterUrl {
host: addr.ip().to_string(),
port: addr.port(),
ca_cert: context.tls.ca_cert().clone().into_owned(),
bootstrap_token: context.actors.bootstrapper.ask(GetToken).await.unwrap(),
};
info!(%url, "Server URL");
let tls = ServerTlsConfig::new().identity(Identity::from_pem(
context.tls.cert_pem(),
context.tls.key_pem(),
));
tonic::transport::Server::builder()
.tls_config(tls)
.map_err(|err| miette!("Faild to setup TLS: {err}"))?
.add_service(ArbiterServiceServer::new(Server::new(context)))
.serve(addr)
.await