Files
arbiter/server/crates/arbiter-server/src/context.rs

404 lines
13 KiB
Rust

use std::collections::HashSet;
use std::sync::Arc;
use std::time::Duration;
use diesel::OptionalExtension as _;
use diesel_async::RunQueryDsl as _;
use ed25519_dalek::VerifyingKey;
use kameo::actor::{ActorRef, Spawn};
use miette::Diagnostic;
use rand::rngs::StdRng;
use secrecy::{ExposeSecret, SecretBox};
use smlang::statemachine;
use thiserror::Error;
use tokio::sync::{watch, RwLock};
use zeroize::Zeroizing;
use crate::{
context::{
bootstrap::{BootstrapActor, generate_token},
lease::LeaseHandler,
tls::{RotationState, RotationTask, TlsDataRaw, TlsManager},
},
db::{
self,
models::ArbiterSetting,
schema::{self, arbiter_settings},
},
};
pub(crate) mod bootstrap;
pub(crate) mod lease;
pub(crate) mod tls;
pub(crate) mod unseal;
#[derive(Error, Debug, Diagnostic)]
pub enum InitError {
#[error("Database setup failed: {0}")]
#[diagnostic(code(arbiter_server::init::database_setup))]
DatabaseSetup(#[from] db::DatabaseSetupError),
#[error("Connection acquire failed: {0}")]
#[diagnostic(code(arbiter_server::init::database_pool))]
DatabasePool(#[from] db::PoolError),
#[error("Database query error: {0}")]
#[diagnostic(code(arbiter_server::init::database_query))]
DatabaseQuery(#[from] diesel::result::Error),
#[error("TLS initialization failed: {0}")]
#[diagnostic(code(arbiter_server::init::tls_init))]
Tls(#[from] tls::TlsInitError),
#[error("Bootstrap token generation failed: {0}")]
#[diagnostic(code(arbiter_server::init::bootstrap_token))]
BootstrapToken(#[from] bootstrap::BootstrapError),
#[error("I/O Error: {0}")]
#[diagnostic(code(arbiter_server::init::io))]
Io(#[from] std::io::Error),
}
#[derive(Error, Debug, Diagnostic)]
pub enum UnsealError {
#[error("Database error: {0}")]
#[diagnostic(code(arbiter_server::unseal::database_pool))]
Database(#[from] db::PoolError),
#[error("Query error: {0}")]
#[diagnostic(code(arbiter_server::unseal::database_query))]
Query(#[from] diesel::result::Error),
#[error("Decryption failed: {0}")]
#[diagnostic(code(arbiter_server::unseal::decryption))]
DecryptionFailed(#[from] crate::crypto::CryptoError),
#[error("Invalid state for unseal")]
#[diagnostic(code(arbiter_server::unseal::invalid_state))]
InvalidState,
#[error("Missing salt in database")]
#[diagnostic(code(arbiter_server::unseal::missing_salt))]
MissingSalt,
#[error("No root key configured in database")]
#[diagnostic(code(arbiter_server::unseal::no_root_key))]
NoRootKey,
}
#[derive(Error, Debug, Diagnostic)]
pub enum SealError {
#[error("Invalid state for seal")]
#[diagnostic(code(arbiter_server::seal::invalid_state))]
InvalidState,
}
/// Secure in-memory storage for root encryption key
///
/// Uses `secrecy` crate for automatic zeroization on drop to prevent key material
/// from remaining in memory after use. SecretBox provides heap-allocated secret
/// storage that implements Send + Sync for safe use in async contexts.
pub struct KeyStorage {
/// 32-byte root key protected by SecretBox
key: SecretBox<[u8; 32]>,
}
impl KeyStorage {
/// Create new KeyStorage from a 32-byte root key
pub fn new(key: [u8; 32]) -> Self {
Self {
key: SecretBox::new(Box::new(key)),
}
}
/// Access the key for cryptographic operations
pub fn key(&self) -> &[u8; 32] {
self.key.expose_secret()
}
}
// Drop автоматически реализован через secrecy::Zeroize
// который зануляет память при освобождении
statemachine! {
name: Server,
transitions: {
*NotBootstrapped + Bootstrapped = Sealed,
Sealed + Unsealed(KeyStorage) / move_key = Ready(KeyStorage),
Ready(KeyStorage) + Sealed / dispose_key = Sealed,
}
}
pub struct _Context;
impl ServerStateMachineContext for _Context {
/// Move key from unseal event into Ready state
fn move_key(&mut self, event_data: KeyStorage) -> Result<KeyStorage, ()> {
// Просто перемещаем KeyStorage из event в state
// Без клонирования - event data consumed
Ok(event_data)
}
/// Securely dispose of key when sealing
#[allow(missing_docs)]
#[allow(clippy::unused_unit)]
fn dispose_key(&mut self, _state_data: &KeyStorage) -> Result<(), ()> {
// KeyStorage будет dropped после state transition
// secrecy::Zeroize зануляет память автоматически
Ok(())
}
}
pub(crate) struct _ServerContextInner {
pub db: db::DatabasePool,
pub state: RwLock<ServerStateMachine<_Context>>,
pub rng: StdRng,
pub tls: Arc<TlsManager>,
pub bootstrapper: ActorRef<BootstrapActor>,
pub rotation_state: RwLock<RotationState>,
pub rotation_acks: Arc<RwLock<HashSet<VerifyingKey>>>,
pub user_agent_leases: LeaseHandler<VerifyingKey>,
pub client_leases: LeaseHandler<VerifyingKey>,
}
#[derive(Clone)]
pub(crate) struct ServerContext(Arc<_ServerContextInner>);
impl std::ops::Deref for ServerContext {
type Target = _ServerContextInner;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl ServerContext {
/// Check if all active clients have acknowledged the rotation
pub async fn check_rotation_ready(&self) -> bool {
// TODO: Implement proper rotation readiness check
// For now, return false as placeholder
false
}
async fn load_tls(
db: &db::DatabasePool,
settings: Option<&ArbiterSetting>,
) -> Result<TlsManager, InitError> {
match settings {
Some(s) if s.current_cert_id.is_some() => {
// Load active certificate from tls_certificates table
Ok(TlsManager::load_from_db(
db.clone(),
s.current_cert_id.unwrap(),
)
.await?)
}
Some(s) => {
// Legacy migration: extract validity and save to new table
let tls_data_raw = TlsDataRaw {
cert: s.cert.clone(),
key: s.cert_key.clone(),
};
// For legacy certificates, use current time as not_before
// and current time + 90 days as not_after
let not_before = chrono::Utc::now().timestamp();
let not_after = not_before + (90 * 24 * 60 * 60); // 90 days
Ok(TlsManager::new_from_legacy(
db.clone(),
tls_data_raw,
not_before,
not_after,
)
.await?)
}
None => {
// First startup - generate new certificate
Ok(TlsManager::new(db.clone()).await?)
}
}
}
pub async fn new(db: db::DatabasePool) -> Result<Self, InitError> {
let mut conn = db.get().await?;
let rng = rand::make_rng();
let settings = arbiter_settings::table
.first::<ArbiterSetting>(&mut conn)
.await
.optional()?;
drop(conn);
// Load TLS manager
let tls = Self::load_tls(&db, settings.as_ref()).await?;
// Load rotation state from database
let rotation_state = RotationState::load_from_db(&db)
.await
.unwrap_or(RotationState::Normal);
let bootstrap_token = generate_token().await?;
let mut state = ServerStateMachine::new(_Context);
if let Some(settings) = &settings
&& settings.root_key_id.is_some()
{
// TODO: pass the encrypted root key to the state machine and let it handle decryption and transition to Sealed
let _ = state.process_event(ServerEvents::Bootstrapped);
}
// Create shutdown channel for rotation task
let (rotation_shutdown_tx, rotation_shutdown_rx) = watch::channel(false);
// Initialize bootstrap actor
let bootstrapper = BootstrapActor::spawn(BootstrapActor::new(&db).await?);
let context = Arc::new(_ServerContextInner {
db: db.clone(),
rng,
tls: Arc::new(tls),
state: RwLock::new(state),
bootstrapper,
rotation_state: RwLock::new(rotation_state),
rotation_acks: Arc::new(RwLock::new(HashSet::new())),
user_agent_leases: Default::default(),
client_leases: Default::default(),
});
Ok(Self(context))
}
/// Unseal vault with password
pub async fn unseal(&self, password: &str) -> Result<(), UnsealError> {
use crate::crypto::root_key;
use diesel::QueryDsl as _;
// 1. Get root_key_id from settings
let mut conn = self.db.get().await?;
let settings: db::models::ArbiterSetting = schema::arbiter_settings::table
.first(&mut conn)
.await?;
let root_key_id = settings.root_key_id.ok_or(UnsealError::NoRootKey)?;
// 2. Load encrypted root key
let encrypted: db::models::AeadEncrypted = schema::aead_encrypted::table
.find(root_key_id)
.first(&mut conn)
.await?;
let salt = encrypted
.argon2_salt
.as_ref()
.ok_or(UnsealError::MissingSalt)?;
drop(conn);
// 3. Decrypt root key using password
let root_key = root_key::decrypt_root_key(&encrypted, password, salt)
.map_err(UnsealError::DecryptionFailed)?;
// 4. Create secure storage
let key_storage = KeyStorage::new(root_key);
// 5. Transition state machine
let mut state = self.state.write().await;
state
.process_event(ServerEvents::Unsealed(key_storage))
.map_err(|_| UnsealError::InvalidState)?;
Ok(())
}
/// Seal the server (lock the key)
pub async fn seal(&self) -> Result<(), SealError> {
let mut state = self.state.write().await;
state
.process_event(ServerEvents::Sealed)
.map_err(|_| SealError::InvalidState)?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_keystorage_creation() {
let key = [42u8; 32];
let storage = KeyStorage::new(key);
assert_eq!(storage.key()[0], 42);
assert_eq!(storage.key().len(), 32);
}
#[test]
fn test_keystorage_zeroization() {
let key = [99u8; 32];
{
let _storage = KeyStorage::new(key);
// storage будет dropped здесь
}
// После drop SecretBox должен зануляеть память
// Это проверяется автоматически через secrecy::Zeroize
}
#[test]
fn test_state_machine_transitions() {
let mut state = ServerStateMachine::new(_Context);
// Начальное состояние
assert!(matches!(state.state(), &ServerStates::NotBootstrapped));
// Bootstrapped transition
state.process_event(ServerEvents::Bootstrapped).unwrap();
assert!(matches!(state.state(), &ServerStates::Sealed));
// Unsealed transition
let key_storage = KeyStorage::new([1u8; 32]);
state
.process_event(ServerEvents::Unsealed(key_storage))
.unwrap();
assert!(matches!(state.state(), &ServerStates::Ready(_)));
// Sealed transition
state.process_event(ServerEvents::Sealed).unwrap();
assert!(matches!(state.state(), &ServerStates::Sealed));
}
#[test]
fn test_move_key_callback() {
let mut ctx = _Context;
let key_storage = KeyStorage::new([7u8; 32]);
let result = ctx.move_key(key_storage);
assert!(result.is_ok());
assert_eq!(result.unwrap().key()[0], 7);
}
#[test]
fn test_dispose_key_callback() {
let mut ctx = _Context;
let key_storage = KeyStorage::new([13u8; 32]);
let result = ctx.dispose_key(&key_storage);
assert!(result.is_ok());
}
#[test]
fn test_invalid_state_transitions() {
let mut state = ServerStateMachine::new(_Context);
// Попытка unseal без bootstrap
let key_storage = KeyStorage::new([1u8; 32]);
let result = state.process_event(ServerEvents::Unsealed(key_storage));
assert!(result.is_err());
// Правильный путь
state.process_event(ServerEvents::Bootstrapped).unwrap();
// Попытка повторного bootstrap
let result = state.process_event(ServerEvents::Bootstrapped);
assert!(result.is_err());
}
}