protobufs/unseal.proto

@@ -4,13 +4,17 @@ package arbiter.unseal;
 
 import "google/protobuf/empty.proto";
 
-message UnsealStart {}
+message UnsealStart {
+  bytes client_pubkey = 1;
+}
 
 message UnsealStartResponse {
-  bytes pubkey = 1;
+  bytes server_pubkey = 1;
 }
 message UnsealEncryptedKey {
-  bytes key = 1;
+  bytes nonce = 1;
+  bytes ciphertext = 2;
+  bytes associated_data = 3;
 }
 
 enum UnsealResult {

server/crates/arbiter-proto/build.rs

@@ -3,6 +3,9 @@ use tonic_prost_build::configure;
 static PROTOBUF_DIR: &str = "../../../protobufs";
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
+
+    println!("cargo::rerun-if-changed={PROTOBUF_DIR}");
+
     configure()
         .message_attribute(".", "#[derive(::kameo::Reply)]")
         .compile_protos(

server/crates/arbiter-server/Cargo.toml

@@ -5,13 +5,7 @@ edition = "2024"
 repository = "https://git.markettakers.org/MarketTakers/arbiter"
 
 [dependencies]
-diesel = { version = "2.3.6", features = [
-    "sqlite",
-    "uuid",
-    "time",
-    "chrono",
-    "serde_json",
-] }
+diesel = { version = "2.3.6", features = ["chrono", "returning_clauses_for_sqlite_3_35", "serde_json", "time", "uuid"] }
 diesel-async = { version = "0.7.4", features = [
     "bb8",
     "migrations",
@@ -45,6 +39,12 @@ chrono.workspace = true
 memsafe = "0.4.0"
 zeroize = { version = "1.8.2", features = ["std", "simd"] }
 kameo.workspace = true
+x25519-dalek = { version = "2.0.1", features = ["getrandom"] }
+chacha20poly1305 = { version = "0.10.1", features = ["std"] }
+statig = { version = "0.4.1", features = ["async"] }
+argon2 = { version = "0.5.3", features = ["zeroize"] }
+restructed = "0.2.2"
+strum = { version = "0.27.2", features = ["derive"] }
 
 [dev-dependencies]
 insta = "1.46.3"

server/crates/arbiter-server/migrations/2026-02-14-171124-0000_init/up.sql

@@ -1,22 +1,34 @@
 create table if not exists aead_encrypted (
     id INTEGER not null PRIMARY KEY,
-    current_nonce integer not null default(1), -- if re-encrypted, this should be incremented
+    current_nonce blob not null default(1), -- if re-encrypted, this should be incremented
     ciphertext blob not null,
     tag blob not null,
-    schema_version integer not null default(1) -- server would need to reencrypt, because this means that we have changed algorithm
+    schema_version integer not null default(1), -- server would need to reencrypt, because this means that we have changed algorithm
+    created_at integer not null default(unixepoch ('now'))
+) STRICT;
+
+create table if not exists root_key_history (
+    id INTEGER not null PRIMARY KEY,
+    -- root key stored as aead encrypted artifact, with only difference that it's decrypted by unseal key (derived from user password)
+    root_key_encryption_nonce blob not null default(1), -- if re-encrypted, this should be incremented. Used for encrypting root key
+    data_encryption_nonce blob not null default(1), -- nonce used for encrypting with key itself
+    ciphertext blob not null,
+    tag blob not null,
+    schema_version integer not null default(1), -- server would need to reencrypt, because this means that we have changed algorithm
+    salt blob not null -- for key deriviation 
 ) STRICT;
 
 -- This is a singleton
 create table if not exists arbiter_settings (
     id INTEGER not null PRIMARY KEY CHECK (id = 1), -- singleton row, id must be 1
-    root_key_id integer references aead_encrypted (id) on delete RESTRICT, -- if null, means wasn't bootstrapped yet
+    root_key_id integer references root_key_history (id) on delete RESTRICT, -- if null, means wasn't bootstrapped yet
     cert_key blob not null,
     cert blob not null
 ) STRICT;
 
 create table if not exists useragent_client (
     id integer not null primary key,
-    nonce integer not null default (1), -- used for auth challenge
+    nonce integer not null default(1), -- used for auth challenge
     public_key blob not null,
     created_at integer not null default(unixepoch ('now')),
     updated_at integer not null default(unixepoch ('now'))
@@ -24,7 +36,7 @@ create table if not exists useragent_client (
 
 create table if not exists program_client (
     id integer not null primary key,
-    nonce integer not null default (1), -- used for auth challenge
+    nonce integer not null default(1), -- used for auth challenge
     public_key blob not null,
     created_at integer not null default(unixepoch ('now')),
     updated_at integer not null default(unixepoch ('now'))

server/crates/arbiter-server/src/actors.rs

@@ -1,3 +1,4 @@
 pub mod user_agent;
 pub mod client;
-pub(crate) mod bootstrap;
+pub(crate) mod bootstrap;
+pub(crate) mod keyholder;

server/crates/arbiter-server/src/actors/bootstrap.rs

@@ -1,19 +1,13 @@
 use arbiter_proto::{BOOTSTRAP_TOKEN_PATH, home_path};
-use diesel::{ExpressionMethods, QueryDsl};
+use diesel::QueryDsl;
 use diesel_async::RunQueryDsl;
 use kameo::{Actor, messages};
-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::{
-    context::{self, ServerContext},
-    db::{self, DatabasePool, schema},
-};
+use crate::db::{self, DatabasePool, schema};
 
 const TOKEN_LENGTH: usize = 64;
 

server/crates/arbiter-server/src/actors/keyholder.rs

@@ -0,0 +1,583 @@
+use diesel::{
+    ExpressionMethods as _, OptionalExtension, QueryDsl, SelectableHelper,
+    dsl::{insert_into, update},
+};
+use diesel_async::{AsyncConnection, RunQueryDsl};
+use kameo::{Actor, messages};
+use memsafe::MemSafe;
+use tracing::{error, info};
+
+use crate::{
+    actors::keyholder::v1::{KeyCell, Nonce},
+    db::{
+        self,
+        models::{self, RootKeyHistory},
+        schema::{self},
+    },
+};
+
+pub mod v1;
+
+#[derive(Default)]
+enum State {
+    #[default]
+    Unbootstrapped,
+    Sealed {
+        encrypted_root_key: RootKeyHistory,
+        data_encryption_nonce: v1::Nonce,
+        root_key_encryption_nonce: v1::Nonce,
+    },
+    Unsealed {
+        root_key_history_id: i32,
+        root_key: KeyCell,
+        nonce: v1::Nonce,
+    },
+}
+
+#[derive(Debug, thiserror::Error, miette::Diagnostic)]
+pub enum Error {
+    #[error("Keyholder is already bootstrapped")]
+    #[diagnostic(code(arbiter::keyholder::already_bootstrapped))]
+    AlreadyBootstrapped,
+    #[error("Keyholder is not bootstrapped")]
+    #[diagnostic(code(arbiter::keyholder::not_bootstrapped))]
+    NotBootstrapped,
+    #[error("Invalid key provided")]
+    #[diagnostic(code(arbiter::keyholder::invalid_key))]
+    InvalidKey,
+
+    #[error("Requested aead entry not found")]
+    #[diagnostic(code(arbiter::keyholder::aead_not_found))]
+    NotFound,
+
+    #[error("Encryption error: {0}")]
+    #[diagnostic(code(arbiter::keyholder::encryption_error))]
+    Encryption(#[from] chacha20poly1305::aead::Error),
+
+    #[error("Database error: {0}")]
+    #[diagnostic(code(arbiter::keyholder::database_error))]
+    DatabaseConnection(#[from] db::PoolError),
+
+    #[error("Database transaction error: {0}")]
+    #[diagnostic(code(arbiter::keyholder::database_transaction_error))]
+    DatabaseTransaction(#[from] diesel::result::Error),
+
+    #[error("Broken database")]
+    #[diagnostic(code(arbiter::keyholder::broken_database))]
+    BrokenDatabase,
+}
+
+/// Manages vault root key and tracks current state of the vault (bootstrapped/unbootstrapped, sealed/unsealed).
+/// Provides API for encrypting and decrypting data using the vault root key.
+/// Abstraction over database to make sure nonces are never reused and encryption keys are never exposed in plaintext outside of this actor.
+#[derive(Actor)]
+pub struct KeyHolderActor {
+    db: db::DatabasePool,
+    state: State,
+}
+
+#[messages]
+impl KeyHolderActor {
+    pub async fn new(db: db::DatabasePool) -> Result<Self, Error> {
+        let state = {
+            let mut conn = db.get().await?;
+
+            let (root_key_history,) = schema::arbiter_settings::table
+                .left_join(schema::root_key_history::table)
+                .select((Option::<RootKeyHistory>::as_select(),))
+                .get_result::<(Option<RootKeyHistory>,)>(&mut conn)
+                .await?;
+
+            match root_key_history {
+                Some(root_key_history) => State::Sealed {
+                    data_encryption_nonce: Nonce::try_from(
+                        root_key_history.data_encryption_nonce.as_slice(),
+                    )
+                    .map_err(|_| {
+                        error!("Broken database: invalid data encryption nonce");
+                        Error::BrokenDatabase
+                    })?,
+                    root_key_encryption_nonce: Nonce::try_from(
+                        root_key_history.root_key_encryption_nonce.as_slice(),
+                    )
+                    .map_err(|_| {
+                        error!("Broken database: invalid root key encryption nonce");
+                        Error::BrokenDatabase
+                    })?,
+                    encrypted_root_key: root_key_history,
+                },
+                None => State::Unbootstrapped,
+            }
+        };
+
+        Ok(Self { db, state })
+    }
+
+    #[message]
+    pub async fn bootstrap(&mut self, seal_key_raw: MemSafe<Vec<u8>>) -> Result<(), Error> {
+        if !matches!(self.state, State::Unbootstrapped) {
+            return Err(Error::AlreadyBootstrapped);
+        }
+        let salt = v1::generate_salt();
+        let mut seal_key = v1::derive_seal_key(seal_key_raw, &salt);
+        let mut root_key = KeyCell::new_secure_random();
+
+        let root_key_nonce = v1::Nonce::default();
+        let data_encryption_nonce = v1::Nonce::default();
+
+        let root_key_ciphertext: Vec<u8> = {
+            let root_key_reader = root_key.0.read().unwrap();
+            let root_key_reader = root_key_reader.as_slice();
+            seal_key
+                .encrypt(&root_key_nonce, v1::ROOT_KEY_TAG, root_key_reader)
+                .map_err(|err| {
+                    error!(?err, "Fatal bootstrap error");
+                    Error::Encryption(err)
+                })?
+        };
+
+        let mut conn = self.db.get().await?;
+
+        let data_encryption_nonce_bytes = data_encryption_nonce.to_vec();
+        let root_key_history_id = conn
+            .transaction(|conn| {
+                Box::pin(async move {
+                    let root_key_history_id: i32 = insert_into(schema::root_key_history::table)
+                        .values(&models::NewRootKeyHistory {
+                            ciphertext: root_key_ciphertext,
+                            tag: v1::ROOT_KEY_TAG.to_vec(),
+                            root_key_encryption_nonce: root_key_nonce.to_vec(),
+                            data_encryption_nonce: data_encryption_nonce_bytes,
+                            schema_version: 1,
+                            salt: salt.to_vec(),
+                        })
+                        .returning(schema::root_key_history::id)
+                        .get_result(conn)
+                        .await?;
+
+                    update(schema::arbiter_settings::table)
+                        .set(schema::arbiter_settings::root_key_id.eq(root_key_history_id))
+                        .execute(conn)
+                        .await?;
+
+                    Result::<_, diesel::result::Error>::Ok(root_key_history_id)
+                })
+            })
+            .await?;
+
+        self.state = State::Unsealed {
+            root_key,
+            root_key_history_id,
+            nonce: data_encryption_nonce,
+        };
+
+        info!("Keyholder bootstrapped successfully");
+
+        Ok(())
+    }
+
+    #[message]
+    pub async fn try_unseal(&mut self, seal_key_raw: MemSafe<Vec<u8>>) -> Result<(), Error> {
+        let State::Sealed {
+            encrypted_root_key,
+            data_encryption_nonce,
+            root_key_encryption_nonce,
+        } = &mut self.state
+        else {
+            return Err(Error::NotBootstrapped);
+        };
+
+        let salt = &encrypted_root_key.salt;
+        let salt = v1::Salt::try_from(salt.as_slice()).map_err(|_| {
+            error!("Broken database: invalid salt for root key");
+            Error::BrokenDatabase
+        })?;
+        let mut seal_key = v1::derive_seal_key(seal_key_raw, &salt);
+
+        let mut root_key = MemSafe::new(encrypted_root_key.ciphertext.clone()).unwrap();
+        seal_key
+            .decrypt_in_place(root_key_encryption_nonce, v1::ROOT_KEY_TAG, &mut root_key)
+            .map_err(|err| {
+                error!(?err, "Failed to unseal root key: invalid seal key");
+                Error::InvalidKey
+            })?;
+
+        self.state = State::Unsealed {
+            root_key_history_id: encrypted_root_key.id,
+            root_key: v1::KeyCell::try_from(root_key).map_err(|err| {
+                error!(?err, "Broken database: invalid encryption key size");
+                Error::BrokenDatabase
+            })?,
+            nonce: std::mem::take(data_encryption_nonce), // we are replacing state, so it's safe to take the nonce out of it
+        };
+
+        info!("Keyholder unsealed successfully");
+
+        Ok(())
+    }
+
+    // Decrypts the `aead_encrypted` entry with the given ID and returns the plaintext
+    #[message]
+    pub async fn decrypt(&mut self, aead_id: i32) -> Result<MemSafe<Vec<u8>>, Error> {
+        let State::Unsealed { root_key, .. } = &mut self.state else {
+            return Err(Error::NotBootstrapped);
+        };
+        let mut conn = self.db.get().await?;
+        let row: models::AeadEncrypted = schema::aead_encrypted::table
+            .select(models::AeadEncrypted::as_select())
+            .filter(schema::aead_encrypted::id.eq(aead_id))
+            .first(&mut conn)
+            .await
+            .optional()?
+            .ok_or(Error::NotFound)?;
+
+        let nonce = v1::Nonce::try_from(row.current_nonce.as_slice()).map_err(|_| {
+            error!(
+                "Broken database: invalid nonce for aead_encrypted id={}",
+                aead_id
+            );
+            Error::BrokenDatabase
+        })?;
+        let mut output = MemSafe::new(row.ciphertext).unwrap();
+        root_key.decrypt_in_place(&nonce, v1::TAG, &mut output)?;
+        Ok(output)
+    }
+
+    // Creates new `aead_encrypted` entry in the database and returns it's ID
+    #[message]
+    pub async fn create_new(&mut self, mut plaintext: MemSafe<Vec<u8>>) -> Result<i32, Error> {
+        let State::Unsealed {
+            root_key,
+            root_key_history_id,
+            nonce,
+        } = &mut self.state
+        else {
+            return Err(Error::NotBootstrapped);
+        };
+
+        let mut conn = self.db.get().await?;
+        nonce.increment();
+
+        let mut ciphertext_buffer = plaintext.write().unwrap();
+        let ciphertext_buffer: &mut Vec<u8> = ciphertext_buffer.as_mut();
+        root_key.encrypt_in_place(&nonce, v1::TAG, &mut *ciphertext_buffer)?;
+
+        let ciphertext = std::mem::take(ciphertext_buffer);
+
+        let aead_id: i32 = conn
+            .transaction(|conn| {
+                Box::pin(async move {
+                    let aead_id: i32 = insert_into(schema::aead_encrypted::table)
+                        .values(&models::NewAeadEncrypted {
+                            ciphertext,
+                            tag: v1::TAG.to_vec(),
+                            current_nonce: nonce.to_vec(),
+                            schema_version: 1,
+                            created_at: chrono::Utc::now().timestamp() as i32,
+                        })
+                        .returning(schema::aead_encrypted::id)
+                        .get_result(conn)
+                        .await?;
+
+                    update(schema::root_key_history::table)
+                        .filter(schema::root_key_history::id.eq(*root_key_history_id))
+                        .set(schema::root_key_history::data_encryption_nonce.eq(nonce.to_vec()))
+                        .execute(conn)
+                        .await?;
+
+                    Result::<_, diesel::result::Error>::Ok(aead_id)
+                })
+            })
+            .await?;
+
+        Ok(aead_id)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashSet;
+
+    use diesel::dsl::insert_into;
+    use diesel_async::RunQueryDsl;
+    use memsafe::MemSafe;
+
+    use crate::db::{self, models::ArbiterSetting};
+
+    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) -> KeyHolderActor {
+        seed_settings(db).await;
+        let mut actor = KeyHolderActor::new(db.clone()).await.unwrap();
+        let seal_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
+        actor.bootstrap(seal_key).await.unwrap();
+        actor
+    }
+
+    #[tokio::test]
+    #[test_log::test]
+    async fn test_bootstrap() {
+        let db = db::create_test_pool().await;
+        seed_settings(&db).await;
+        let mut actor = KeyHolderActor::new(db.clone()).await.unwrap();
+
+        assert!(matches!(actor.state, State::Unbootstrapped));
+
+        let seal_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
+        actor.bootstrap(seal_key).await.unwrap();
+
+        assert!(matches!(actor.state, State::Unsealed { .. }));
+
+        let mut conn = db.get().await.unwrap();
+        let row: models::RootKeyHistory = schema::root_key_history::table
+            .select(models::RootKeyHistory::as_select())
+            .first(&mut conn)
+            .await
+            .unwrap();
+
+        assert_eq!(row.schema_version, 1);
+        assert_eq!(row.tag, v1::ROOT_KEY_TAG);
+        assert!(!row.ciphertext.is_empty());
+        assert!(!row.salt.is_empty());
+        assert_eq!(row.data_encryption_nonce, v1::Nonce::default().to_vec());
+    }
+
+    #[tokio::test]
+    #[test_log::test]
+    async fn test_bootstrap_rejects_double() {
+        let db = db::create_test_pool().await;
+        let mut actor = bootstrapped_actor(&db).await;
+
+        let seal_key2 = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
+        let err = actor.bootstrap(seal_key2).await.unwrap_err();
+        assert!(matches!(err, Error::AlreadyBootstrapped));
+    }
+
+    #[tokio::test]
+    #[test_log::test]
+    async fn test_create_decrypt_roundtrip() {
+        let db = db::create_test_pool().await;
+        let mut actor = bootstrapped_actor(&db).await;
+
+        let plaintext = b"hello arbiter";
+        let aead_id = actor
+            .create_new(MemSafe::new(plaintext.to_vec()).unwrap())
+            .await
+            .unwrap();
+
+        let mut decrypted = actor.decrypt(aead_id).await.unwrap();
+        let decrypted = decrypted.read().unwrap();
+        assert_eq!(*decrypted, plaintext);
+    }
+
+    #[tokio::test]
+    #[test_log::test]
+    async fn test_create_new_before_bootstrap_fails() {
+        let db = db::create_test_pool().await;
+        seed_settings(&db).await;
+        let mut actor = KeyHolderActor::new(db).await.unwrap();
+
+        let err = actor
+            .create_new(MemSafe::new(b"data".to_vec()).unwrap())
+            .await
+            .unwrap_err();
+        assert!(matches!(err, Error::NotBootstrapped));
+    }
+
+    #[tokio::test]
+    #[test_log::test]
+    async fn test_decrypt_before_bootstrap_fails() {
+        let db = db::create_test_pool().await;
+        seed_settings(&db).await;
+        let mut actor = KeyHolderActor::new(db).await.unwrap();
+
+        let err = actor.decrypt(1).await.unwrap_err();
+        assert!(matches!(err, Error::NotBootstrapped));
+    }
+
+    #[tokio::test]
+    #[test_log::test]
+    async fn test_decrypt_nonexistent_returns_not_found() {
+        let db = db::create_test_pool().await;
+        let mut actor = bootstrapped_actor(&db).await;
+
+        let err = actor.decrypt(9999).await.unwrap_err();
+        assert!(matches!(err, Error::NotFound));
+    }
+
+    #[tokio::test]
+    #[test_log::test]
+    async fn test_new_restores_sealed_state() {
+        let db = db::create_test_pool().await;
+        let actor = bootstrapped_actor(&db).await;
+        drop(actor);
+
+        let actor2 = KeyHolderActor::new(db).await.unwrap();
+        assert!(matches!(actor2.state, State::Sealed { .. }));
+    }
+
+    #[tokio::test]
+    #[test_log::test]
+    async fn test_nonce_never_reused() {
+        let db = db::create_test_pool().await;
+        let mut actor = bootstrapped_actor(&db).await;
+
+        let n = 5;
+        let mut ids = Vec::with_capacity(n);
+        for i in 0..n {
+            let id = actor
+                .create_new(MemSafe::new(format!("secret {i}").into_bytes()).unwrap())
+                .await
+                .unwrap();
+            ids.push(id);
+        }
+
+        // read all stored nonces from DB
+        let mut conn = db.get().await.unwrap();
+        let rows: Vec<models::AeadEncrypted> = schema::aead_encrypted::table
+            .select(models::AeadEncrypted::as_select())
+            .load(&mut conn)
+            .await
+            .unwrap();
+
+        assert_eq!(rows.len(), n);
+
+        let nonces: Vec<&Vec<u8>> = rows.iter().map(|r| &r.current_nonce).collect();
+        let unique: HashSet<&Vec<u8>> = nonces.iter().copied().collect();
+        assert_eq!(nonces.len(), unique.len(), "all nonces must be unique");
+
+        // verify nonces are sequential increments from 1
+        for (i, row) in rows.iter().enumerate() {
+            let mut expected = v1::Nonce::default();
+            for _ in 0..=i {
+                expected.increment();
+            }
+            assert_eq!(row.current_nonce, expected.to_vec(), "nonce {i} mismatch");
+        }
+
+        // verify data_encryption_nonce on root_key_history tracks the latest nonce
+        let root_row: models::RootKeyHistory = schema::root_key_history::table
+            .select(models::RootKeyHistory::as_select())
+            .first(&mut conn)
+            .await
+            .unwrap();
+        let last_nonce = &rows.last().unwrap().current_nonce;
+        assert_eq!(
+            &root_row.data_encryption_nonce, last_nonce,
+            "root_key_history must track the latest nonce"
+        );
+    }
+
+    #[tokio::test]
+    #[test_log::test]
+    async fn test_unseal_correct_password() {
+        let db = db::create_test_pool().await;
+        let mut actor = bootstrapped_actor(&db).await;
+
+        let plaintext = b"survive a restart";
+        let aead_id = actor
+            .create_new(MemSafe::new(plaintext.to_vec()).unwrap())
+            .await
+            .unwrap();
+        drop(actor);
+
+        let mut actor = KeyHolderActor::new(db.clone()).await.unwrap();
+        assert!(matches!(actor.state, State::Sealed { .. }));
+
+        let seal_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
+        actor.try_unseal(seal_key).await.unwrap();
+        assert!(matches!(actor.state, State::Unsealed { .. }));
+
+        // previously encrypted data is still decryptable
+        let mut decrypted = actor.decrypt(aead_id).await.unwrap();
+        assert_eq!(*decrypted.read().unwrap(), plaintext);
+    }
+
+    #[tokio::test]
+    #[test_log::test]
+    async fn test_unseal_wrong_then_correct_password() {
+        let db = db::create_test_pool().await;
+        let mut actor = bootstrapped_actor(&db).await;
+
+        let plaintext = b"important data";
+        let aead_id = actor
+            .create_new(MemSafe::new(plaintext.to_vec()).unwrap())
+            .await
+            .unwrap();
+        drop(actor);
+
+        let mut actor = KeyHolderActor::new(db.clone()).await.unwrap();
+        assert!(matches!(actor.state, State::Sealed { .. }));
+
+        // wrong password
+        let bad_key = MemSafe::new(b"wrong-password".to_vec()).unwrap();
+        let err = actor.try_unseal(bad_key).await.unwrap_err();
+        assert!(matches!(err, Error::InvalidKey));
+        assert!(
+            matches!(actor.state, State::Sealed { .. }),
+            "state must remain Sealed after failed attempt"
+        );
+
+        // correct password
+        let good_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
+        actor.try_unseal(good_key).await.unwrap();
+        assert!(matches!(actor.state, State::Unsealed { .. }));
+
+        let mut decrypted = actor.decrypt(aead_id).await.unwrap();
+        assert_eq!(*decrypted.read().unwrap(), plaintext);
+    }
+
+    #[tokio::test]
+    #[test_log::test]
+    async fn test_ciphertext_differs_across_entries() {
+        let db = db::create_test_pool().await;
+        let mut actor = bootstrapped_actor(&db).await;
+
+        let plaintext = b"same content";
+        let id1 = actor
+            .create_new(MemSafe::new(plaintext.to_vec()).unwrap())
+            .await
+            .unwrap();
+        let id2 = actor
+            .create_new(MemSafe::new(plaintext.to_vec()).unwrap())
+            .await
+            .unwrap();
+
+        // different nonces => different ciphertext, even for identical plaintext
+        let mut conn = db.get().await.unwrap();
+        let row1: models::AeadEncrypted = schema::aead_encrypted::table
+            .filter(schema::aead_encrypted::id.eq(id1))
+            .select(models::AeadEncrypted::as_select())
+            .first(&mut conn)
+            .await
+            .unwrap();
+        let row2: models::AeadEncrypted = schema::aead_encrypted::table
+            .filter(schema::aead_encrypted::id.eq(id2))
+            .select(models::AeadEncrypted::as_select())
+            .first(&mut conn)
+            .await
+            .unwrap();
+
+        assert_ne!(row1.ciphertext, row2.ciphertext);
+
+        // but both decrypt to the same plaintext
+        let mut d1 = actor.decrypt(id1).await.unwrap();
+        let mut d2 = actor.decrypt(id2).await.unwrap();
+        assert_eq!(*d1.read().unwrap(), plaintext);
+        assert_eq!(*d2.read().unwrap(), plaintext);
+    }
+}

server/crates/arbiter-server/src/actors/keyholder/v1.rs

@@ -0,0 +1,241 @@
+use std::ops::Deref as _;
+
+use argon2::{Algorithm, Argon2, password_hash::Salt as ArgonSalt};
+use chacha20poly1305::{
+    AeadInPlace, Key, KeyInit as _, XChaCha20Poly1305, XNonce,
+    aead::{AeadMut, Error, Payload},
+};
+use memsafe::MemSafe;
+use rand::{
+    Rng as _, SeedableRng,
+    rngs::{StdRng, SysRng},
+};
+
+pub const ROOT_KEY_TAG: &[u8] = "arbiter/seal/v1".as_bytes();
+pub const TAG: &[u8] = "arbiter/private-key/v1".as_bytes();
+pub const NONCE_LENGTH: usize = 24;
+
+#[derive(Default)]
+pub struct Nonce([u8; NONCE_LENGTH]);
+impl Nonce {
+    pub fn increment(&mut self) {
+        for i in (0..self.0.len()).rev() {
+            if self.0[i] == 0xFF {
+                self.0[i] = 0;
+            } else {
+                self.0[i] += 1;
+                break;
+            }
+        }
+    }
+
+    pub fn to_vec(&self) -> Vec<u8> {
+        self.0.to_vec()
+    }
+}
+impl<'a> TryFrom<&'a [u8]> for Nonce {
+    type Error = ();
+
+    fn try_from(value: &'a [u8]) -> Result<Self, Self::Error> {
+        if value.len() != NONCE_LENGTH {
+            return Err(());
+        }
+        let mut nonce = [0u8; NONCE_LENGTH];
+        nonce.copy_from_slice(&value);
+        Ok(Self(nonce))
+    }
+}
+
+pub struct KeyCell(pub(super) MemSafe<Key>);
+impl From<MemSafe<Key>> for KeyCell {
+    fn from(value: MemSafe<Key>) -> Self {
+        Self(value)
+    }
+}
+impl TryFrom<MemSafe<Vec<u8>>> for KeyCell {
+    type Error = ();
+
+    fn try_from(mut value: MemSafe<Vec<u8>>) -> Result<Self, Self::Error> {
+        let value = value.read().unwrap();
+        if value.len() != size_of::<Key>() {
+            return Err(());
+        }
+        let mut cell = MemSafe::new(Key::default()).unwrap();
+        {
+            let mut cell_write = cell.write().unwrap();
+            let cell_slice: &mut [u8] = cell_write.as_mut();
+            cell_slice.copy_from_slice(&value);
+        }
+        Ok(Self(cell))
+    }
+}
+
+impl KeyCell {
+    pub fn new_secure_random() -> Self {
+        let mut key = MemSafe::new(Key::default()).unwrap();
+        {
+            let mut key_buffer = key.write().unwrap();
+            let key_buffer: &mut [u8] = key_buffer.as_mut();
+
+            let mut rng = StdRng::try_from_rng(&mut SysRng).unwrap();
+            rng.fill_bytes(key_buffer);
+        }
+
+        key.into()
+    }
+
+    pub fn into_inner(self) -> MemSafe<Key> {
+        self.0
+    }
+
+    pub fn encrypt_in_place(
+        &mut self,
+        nonce: &Nonce,
+        associated_data: &[u8],
+        mut buffer: impl AsMut<Vec<u8>>,
+    ) -> Result<(), Error> {
+        let key_reader = self.0.read().unwrap();
+        let key_ref = key_reader.deref();
+        let cipher = XChaCha20Poly1305::new(key_ref);
+        let nonce = XNonce::from_slice(nonce.0.as_ref());
+        let buffer = buffer.as_mut();
+        cipher.encrypt_in_place(nonce, associated_data, buffer)
+    }
+    pub fn decrypt_in_place(
+        &mut self,
+        nonce: &Nonce,
+        associated_data: &[u8],
+        buffer: &mut MemSafe<Vec<u8>>,
+    ) -> Result<(), Error> {
+        let key_reader = self.0.read().unwrap();
+        let key_ref = key_reader.deref();
+        let cipher = XChaCha20Poly1305::new(key_ref);
+        let nonce = XNonce::from_slice(nonce.0.as_ref());
+        let mut buffer = buffer.write().unwrap();
+        let buffer: &mut Vec<u8> = buffer.as_mut();
+        cipher.decrypt_in_place(nonce, associated_data, buffer)
+    }
+
+    pub fn encrypt(
+        &mut self,
+        nonce: &Nonce,
+        associated_data: &[u8],
+        plaintext: impl AsRef<[u8]>,
+    ) -> Result<Vec<u8>, Error> {
+        let key_reader = self.0.read().unwrap();
+        let key_ref = key_reader.deref();
+        let mut cipher = XChaCha20Poly1305::new(key_ref);
+        let nonce = XNonce::from_slice(nonce.0.as_ref());
+
+
+        let ciphertext = cipher.encrypt(
+            &nonce,
+            Payload {
+                msg: plaintext.as_ref(),
+                aad: associated_data,
+            },
+        )?;
+        Ok(ciphertext)
+    }
+}
+
+pub type Salt = [u8; ArgonSalt::RECOMMENDED_LENGTH];
+
+pub(super) fn generate_salt() -> Salt {
+    let mut salt = Salt::default();
+    let mut rng = StdRng::try_from_rng(&mut SysRng).unwrap();
+    rng.fill_bytes(&mut salt);
+    salt
+}
+
+/// User password might be of different length, have not enough entropy, etc...
+/// Derive a fixed-length key from the password using Argon2id, which is designed for password hashing and key derivation.
+pub(super) fn derive_seal_key(mut password: MemSafe<Vec<u8>>, salt: &Salt) -> KeyCell {
+    let params = argon2::Params::new(262_144, 3, 4, None).unwrap();
+    let hasher = Argon2::new(Algorithm::Argon2id, argon2::Version::V0x13, params);
+    let mut key = MemSafe::new(Key::default()).unwrap();
+    {
+        let password_source = password.read().unwrap();
+        let mut key_buffer = key.write().unwrap();
+        let key_buffer: &mut [u8] = key_buffer.as_mut();
+
+        hasher
+            .hash_password_into(password_source.deref(), salt, key_buffer)
+            .unwrap();
+    }
+
+    key.into()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use memsafe::MemSafe;
+
+    #[test]
+    pub fn derive_seal_key_deterministic() {
+        static PASSWORD: &[u8] = b"password";
+        let password = MemSafe::new(PASSWORD.to_vec()).unwrap();
+        let password2 = MemSafe::new(PASSWORD.to_vec()).unwrap();
+        let salt = generate_salt();
+
+        let mut key1 = derive_seal_key(password, &salt);
+        let mut key2 = derive_seal_key(password2, &salt);
+
+        let key1_reader = key1.0.read().unwrap();
+        let key2_reader = key2.0.read().unwrap();
+
+        assert_eq!(key1_reader.deref(), key2_reader.deref());
+    }
+
+    #[test]
+    pub fn successful_derive() {
+        static PASSWORD: &[u8] = b"password";
+        let password = MemSafe::new(PASSWORD.to_vec()).unwrap();
+        let salt = generate_salt();
+
+        let mut key = derive_seal_key(password, &salt);
+        let key_reader = key.0.read().unwrap();
+        let key_ref = key_reader.deref();
+
+        assert_ne!(key_ref.as_slice(), &[0u8; 32][..]);
+    }
+
+    #[test]
+    pub fn encrypt_decrypt() {
+        static PASSWORD: &[u8] = b"password";
+        let password = MemSafe::new(PASSWORD.to_vec()).unwrap();
+        let salt = generate_salt();
+
+        let mut key = derive_seal_key(password, &salt);
+        let nonce = Nonce(*b"unique nonce 123 1231233"); // 24 bytes for XChaCha20Poly1305
+        let associated_data = b"associated data";
+        let mut buffer = b"secret data".to_vec();
+
+        key.encrypt_in_place(&nonce, associated_data, &mut buffer)
+            .unwrap();
+        assert_ne!(buffer, b"secret data");
+
+        let mut buffer = MemSafe::new(buffer).unwrap();
+
+        key.decrypt_in_place(&nonce, associated_data, &mut buffer)
+            .unwrap();
+
+        let buffer = buffer.read().unwrap();
+        assert_eq!(*buffer, b"secret data");
+    }
+
+    #[test]
+    // We should fuzz this
+    pub fn test_nonce_increment() {
+        let mut nonce = Nonce([0u8; NONCE_LENGTH]);
+        nonce.increment();
+
+        assert_eq!(
+            nonce.0,
+            [
+                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1
+            ]
+        );
+    }
+}

server/crates/arbiter-server/src/actors/user_agent.rs

@@ -1,104 +1,50 @@
+use std::{
+    ops::DerefMut,
+    sync::Mutex,
+};
+
 use arbiter_proto::proto::{
     UserAgentResponse,
     auth::{
-        self, AuthChallenge, AuthChallengeRequest, AuthOk, ClientMessage,
-        ServerMessage as AuthServerMessage, client_message::Payload as ClientAuthPayload,
+        self, AuthChallengeRequest, AuthOk, ServerMessage as AuthServerMessage,
         server_message::Payload as ServerAuthPayload,
     },
-    user_agent_request::Payload as UserAgentRequestPayload,
+    unseal::{UnsealEncryptedKey, UnsealResult, UnsealStart, UnsealStartResponse},
     user_agent_response::Payload as UserAgentResponsePayload,
 };
+use chacha20poly1305::{
+    AeadInPlace, XChaCha20Poly1305, XNonce,
+    aead::KeyInit,
+};
 use diesel::{ExpressionMethods as _, OptionalExtension as _, QueryDsl, dsl::update};
 use diesel_async::{AsyncConnection, RunQueryDsl};
 use ed25519_dalek::VerifyingKey;
-use kameo::{
-    Actor,
-    actor::{ActorRef, Spawn},
-    messages,
-    prelude::Context,
-};
+use kameo::{Actor, actor::ActorRef, messages};
+use memsafe::MemSafe;
 use tokio::sync::mpsc::Sender;
 use tonic::Status;
 use tracing::{error, info};
+use x25519_dalek::{EphemeralSecret, PublicKey};
 
 use crate::{
     ServerContext,
-    actors::bootstrap::{BootstrapActor, ConsumeToken},
+    actors::{
+        bootstrap::{BootstrapActor, ConsumeToken},
+        user_agent::state::{
+            AuthRequestContext, ChallengeContext, DummyContext, UnsealContext, UserAgentEvents,
+            UserAgentStateMachine, UserAgentStates,
+        },
+    },
     db::{self, schema},
     errors::GrpcStatusExt,
 };
 
-/// Context for state machine with validated key and sent challenge
-/// Challenge is then transformed to bytes using shared function and verified
-#[derive(Clone, Debug)]
-pub struct ChallengeContext {
-    challenge: AuthChallenge,
-    key: VerifyingKey,
-}
+mod state;
+#[cfg(test)]
+mod tests;
 
-// Request context with deserialized public key for state machine.
-// This intermediate struct is needed because the state machine branches depending on presence of bootstrap token,
-// but we want to have the deserialized key in both branches.
-#[derive(Clone, Debug)]
-pub struct AuthRequestContext {
-    pubkey: VerifyingKey,
-    bootstrap_token: Option<String>,
-}
-
-smlang::statemachine!(
-    name: UserAgent,
-    derive_states: [Debug],
-    custom_error: false,
-    transitions: {
-        *Init + AuthRequest(AuthRequestContext) / auth_request_context =  ReceivedAuthRequest(AuthRequestContext),
-        ReceivedAuthRequest(AuthRequestContext) + ReceivedBootstrapToken = Idle,
-
-        ReceivedAuthRequest(AuthRequestContext) + SentChallenge(ChallengeContext) / move_challenge = WaitingForChallengeSolution(ChallengeContext),
-
-        WaitingForChallengeSolution(ChallengeContext) + ReceivedGoodSolution = Idle,
-        WaitingForChallengeSolution(ChallengeContext) + ReceivedBadSolution = AuthError, // block further transitions, but connection should close anyway
-
-        Idle + UnsealRequest / generate_temp_keypair = UnsealStarted(ed25519_dalek::SigningKey),
-        UnsealStarted(ed25519_dalek::SigningKey) + SentTempKeypair / move_keypair = WaitingForUnsealKey(ed25519_dalek::SigningKey),
-    }
-);
-
-pub struct DummyContext;
-impl UserAgentStateMachineContext for DummyContext {
-    #[allow(missing_docs)]
-    #[allow(clippy::unused_unit)]
-    fn move_challenge(
-        &mut self,
-        _state_data: &AuthRequestContext,
-        event_data: ChallengeContext,
-    ) -> Result<ChallengeContext, ()> {
-        Ok(event_data)
-    }
-
-    #[allow(missing_docs)]
-    #[allow(clippy::unused_unit)]
-    fn auth_request_context(
-        &mut self,
-        event_data: AuthRequestContext,
-    ) -> Result<AuthRequestContext, ()> {
-        Ok(event_data)
-    }
-
-    #[allow(missing_docs)]
-    #[allow(clippy::unused_unit)]
-    fn move_keypair(
-        &mut self,
-        state_data: &ed25519_dalek::SigningKey,
-    ) -> Result<ed25519_dalek::SigningKey, ()> {
-        Ok(state_data.clone())
-    }
-
-    #[allow(missing_docs)]
-    #[allow(clippy::unused_unit)]
-    fn generate_temp_keypair(&mut self) -> Result<ed25519_dalek::SigningKey, ()> {
-        Ok(ed25519_dalek::SigningKey::generate(&mut rand::rng()))
-    }
-}
+mod transport;
+pub(crate) use transport::handle_user_agent;
 
 #[derive(Actor)]
 pub struct UserAgentActor {
@@ -272,18 +218,93 @@ fn auth_response(payload: ServerAuthPayload) -> UserAgentResponse {
     }
 }
 
+fn unseal_response(payload: UserAgentResponsePayload) -> UserAgentResponse {
+    UserAgentResponse {
+        payload: Some(payload),
+    }
+}
+
 #[messages]
 impl UserAgentActor {
-    #[message(ctx)]
-    pub async fn handle_auth_challenge_request(
-        &mut self,
-        req: AuthChallengeRequest,
-        ctx: &mut Context<Self, Output>,
-    ) -> Output {
+    #[message]
+    pub async fn handle_unseal_request(&mut self, req: UnsealStart) -> Output {
+        let secret = EphemeralSecret::random();
+        let public_key = PublicKey::from(&secret);
+
+        let client_pubkey_bytes: [u8; 32] = req
+            .client_pubkey
+            .try_into()
+            .map_err(|_| Status::invalid_argument("client_pubkey must be 32 bytes"))?;
+
+        let client_public_key = PublicKey::from(client_pubkey_bytes);
+
+        self.transition(UserAgentEvents::UnsealRequest(UnsealContext {
+            server_public_key: public_key,
+            secret: Mutex::new(Some(secret)),
+            client_public_key,
+        }))?;
+
+        Ok(unseal_response(
+            UserAgentResponsePayload::UnsealStartResponse(UnsealStartResponse {
+                server_pubkey: public_key.as_bytes().to_vec(),
+            }),
+        ))
+    }
+
+    #[message]
+    pub async fn handle_unseal_encrypted_key(&mut self, req: UnsealEncryptedKey) -> Output {
+        let UserAgentStates::WaitingForUnsealKey(unseal_context) = self.state.state() else {
+            error!("Received unseal encrypted key in invalid state");
+            return Err(Status::failed_precondition(
+                "Invalid state for unseal encrypted key",
+            ));
+        };
+        let ephemeral_secret = {
+            let mut secret_lock = unseal_context.secret.lock().unwrap();
+            let secret = secret_lock.take();
+            match secret {
+                Some(secret) => secret,
+                None => {
+                    drop(secret_lock);
+                    error!("Ephemeral secret already taken");
+                    self.transition(UserAgentEvents::ReceivedInvalidKey)?;
+                    return Ok(unseal_response(UserAgentResponsePayload::UnsealResult(
+                        UnsealResult::InvalidKey.into(),
+                    )));
+                }
+            }
+        };
+
+        let nonce = XNonce::from_slice(&req.nonce);
+
+        let shared_secret = ephemeral_secret.diffie_hellman(&unseal_context.client_public_key);
+        let cipher = XChaCha20Poly1305::new(shared_secret.as_bytes().into());
+
+        let mut root_key_buffer = MemSafe::new(req.ciphertext.clone()).unwrap();
+        let mut write_handle = root_key_buffer.write().unwrap();
+        let write_handle = write_handle.deref_mut();
+
+        let decryption_result = cipher
+            .decrypt_in_place(nonce, &req.associated_data, write_handle);
+
+        match decryption_result {
+            Ok(_) => todo!("Send key to the keyguarding"),
+            Err(err) => {
+                error!(?err, "Failed to decrypt unseal key");
+                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
+                return Ok(unseal_response(UserAgentResponsePayload::UnsealResult(
+                    UnsealResult::InvalidKey.into(),
+                )));
+            },
+        }
+    }
+
+    #[message]
+    pub async fn handle_auth_challenge_request(&mut self, req: AuthChallengeRequest) -> 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| {
+        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")
         })?;
@@ -299,11 +320,10 @@ impl UserAgentActor {
         }
     }
 
-    #[message(ctx)]
+    #[message]
     pub async fn handle_auth_challenge_solution(
         &mut self,
         solution: auth::AuthChallengeSolution,
-        ctx: &mut Context<Self, Output>,
     ) -> Output {
         let (valid, challenge_context) = self.verify_challenge_solution(&solution)?;
 
@@ -321,211 +341,3 @@ impl UserAgentActor {
         }
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use arbiter_proto::proto::{
-        UserAgentResponse,
-        auth::{self, AuthChallengeRequest, AuthOk},
-        user_agent_response::Payload as UserAgentResponsePayload,
-    };
-    use chrono::format;
-    use diesel::{ExpressionMethods as _, QueryDsl, insert_into};
-    use diesel_async::RunQueryDsl;
-    use ed25519_dalek::Signer as _;
-    use kameo::actor::Spawn;
-
-    use crate::{
-        actors::{
-            bootstrap::BootstrapActor,
-            user_agent::{HandleAuthChallengeRequest, HandleAuthChallengeSolution},
-        },
-        db::{self, schema},
-    };
-
-    use super::UserAgentActor;
-
-    #[tokio::test]
-    #[test_log::test]
-    pub async fn test_bootstrap_token_auth() {
-        let db = db::create_test_pool().await;
-        // explicitly not installing any user_agent pubkeys
-        let bootstrapper = BootstrapActor::new(&db).await.unwrap(); // this will create bootstrap token
-        let token = bootstrapper.get_token().unwrap();
-
-        let bootstrapper_ref = BootstrapActor::spawn(bootstrapper);
-        let user_agent = UserAgentActor::new_manual(
-            db.clone(),
-            bootstrapper_ref,
-            tokio::sync::mpsc::channel(1).0, // dummy channel, we won't actually send responses in this test
-        );
-        let user_agent_ref = UserAgentActor::spawn(user_agent);
-
-        // simulate client sending auth request with bootstrap token
-        let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
-        let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
-
-        let result = user_agent_ref
-            .ask(HandleAuthChallengeRequest {
-                req: AuthChallengeRequest {
-                    pubkey: pubkey_bytes,
-                    bootstrap_token: Some(token),
-                },
-            })
-            .await
-            .expect("Shouldn't fail to send message");
-
-        // auth succeeded
-        assert_eq!(
-            result,
-            UserAgentResponse {
-                payload: Some(UserAgentResponsePayload::AuthMessage(
-                    arbiter_proto::proto::auth::ServerMessage {
-                        payload: Some(arbiter_proto::proto::auth::server_message::Payload::AuthOk(
-                            AuthOk {},
-                        )),
-                    },
-                )),
-            }
-        );
-
-        // key is succesfully recorded in database
-        let mut conn = db.get().await.unwrap();
-        let stored_pubkey: Vec<u8> = schema::useragent_client::table
-            .select(schema::useragent_client::public_key)
-            .first::<Vec<u8>>(&mut conn)
-            .await
-            .unwrap();
-        assert_eq!(stored_pubkey, new_key.verifying_key().to_bytes().to_vec());
-    }
-
-    #[tokio::test]
-    #[test_log::test]
-    pub async fn test_bootstrap_invalid_token_auth() {
-        let db = db::create_test_pool().await;
-        // explicitly not installing any user_agent pubkeys
-        let bootstrapper = BootstrapActor::new(&db).await.unwrap(); // this will create bootstrap token
-
-        let bootstrapper_ref = BootstrapActor::spawn(bootstrapper);
-        let user_agent = UserAgentActor::new_manual(
-            db.clone(),
-            bootstrapper_ref,
-            tokio::sync::mpsc::channel(1).0, // dummy channel, we won't actually send responses in this test
-        );
-        let user_agent_ref = UserAgentActor::spawn(user_agent);
-
-        // simulate client sending auth request with bootstrap token
-        let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
-        let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
-
-        let result = user_agent_ref
-            .ask(HandleAuthChallengeRequest {
-                req: AuthChallengeRequest {
-                    pubkey: pubkey_bytes,
-                    bootstrap_token: Some("invalid_token".to_string()),
-                },
-            })
-            .await;
-
-        match result {
-            Err(kameo::error::SendError::HandlerError(status)) => {
-                assert_eq!(status.code(), tonic::Code::InvalidArgument);
-                insta::assert_debug_snapshot!(status, @r#"
-                Status {
-                    code: InvalidArgument,
-                    message: "Invalid bootstrap token",
-                    source: None,
-                }
-                "#);
-            }
-            Err(other) => {
-                panic!("Expected SendError::HandlerError, got {other:?}");
-            }
-            Ok(_) => {
-                panic!("Expected error due to invalid bootstrap token, but got success");
-            }
-        }
-    }
-
-    #[tokio::test]
-    #[test_log::test]
-    pub async fn test_challenge_auth() {
-        let db = db::create_test_pool().await;
-
-        let bootstrapper_ref = BootstrapActor::spawn(BootstrapActor::new(&db).await.unwrap());
-        let user_agent = UserAgentActor::new_manual(
-            db.clone(),
-            bootstrapper_ref,
-            tokio::sync::mpsc::channel(1).0, // dummy channel, we won't actually send responses in this test
-        );
-        let user_agent_ref = UserAgentActor::spawn(user_agent);
-
-        // simulate client sending auth request with bootstrap token
-        let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
-        let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
-
-        // insert pubkey into database to trigger challenge-response auth flow
-        {
-            let mut conn = db.get().await.unwrap();
-            insert_into(schema::useragent_client::table)
-                .values(schema::useragent_client::public_key.eq(pubkey_bytes.clone()))
-                .execute(&mut conn)
-                .await
-                .unwrap();
-        }
-
-        let result = user_agent_ref
-            .ask(HandleAuthChallengeRequest {
-                req: AuthChallengeRequest {
-                    pubkey: pubkey_bytes,
-                    bootstrap_token: None,
-                },
-            })
-            .await
-            .expect("Shouldn't fail to send message");
-
-        // auth challenge succeeded
-        let UserAgentResponse {
-            payload:
-                Some(UserAgentResponsePayload::AuthMessage(arbiter_proto::proto::auth::ServerMessage {
-                    payload:
-                        Some(arbiter_proto::proto::auth::server_message::Payload::AuthChallenge(
-                            challenge,
-                        )),
-                })),
-        } = result
-        else {
-            panic!("Expected auth challenge response, got {result:?}");
-        };
-
-        let formatted_challenge = arbiter_proto::format_challenge(&challenge);
-        let signature = new_key.sign(&formatted_challenge);
-        let serialized_signature = signature.to_bytes().to_vec();
-
-        let result = user_agent_ref
-            .ask(HandleAuthChallengeSolution {
-                solution: auth::AuthChallengeSolution {
-                    signature: serialized_signature,
-                },
-            })
-            .await
-            .expect("Shouldn't fail to send message");
-
-        // auth succeeded
-        assert_eq!(
-            result,
-            UserAgentResponse {
-                payload: Some(UserAgentResponsePayload::AuthMessage(
-                    arbiter_proto::proto::auth::ServerMessage {
-                        payload: Some(arbiter_proto::proto::auth::server_message::Payload::AuthOk(
-                            AuthOk {},
-                        )),
-                    },
-                )),
-            }
-        );
-    }
-}
-
-mod transport;
-pub(crate) use transport::handle_user_agent;

server/crates/arbiter-server/src/actors/user_agent/state.rs

@@ -0,0 +1,76 @@
+use std::sync::Mutex;
+
+use arbiter_proto::proto::auth::AuthChallenge;
+use ed25519_dalek::VerifyingKey;
+use x25519_dalek::{EphemeralSecret, PublicKey};
+
+/// Context for state machine with validated key and sent challenge
+/// Challenge is then transformed to bytes using shared function and verified
+#[derive(Clone, Debug)]
+pub struct ChallengeContext {
+    pub challenge: AuthChallenge,
+    pub key: VerifyingKey,
+}
+
+// Request context with deserialized public key for state machine.
+// This intermediate struct is needed because the state machine branches depending on presence of bootstrap token,
+// but we want to have the deserialized key in both branches.
+#[derive(Clone, Debug)]
+pub struct AuthRequestContext {
+    pub pubkey: VerifyingKey,
+    pub bootstrap_token: Option<String>,
+}
+
+pub struct UnsealContext {
+    pub server_public_key: PublicKey,
+    pub client_public_key: PublicKey,
+    pub secret: Mutex<Option<EphemeralSecret>>,
+}
+
+
+
+smlang::statemachine!(
+    name: UserAgent,
+    custom_error: false,
+    transitions: {
+        *Init + AuthRequest(AuthRequestContext) / auth_request_context =  ReceivedAuthRequest(AuthRequestContext),
+        ReceivedAuthRequest(AuthRequestContext) + ReceivedBootstrapToken = Idle,
+
+        ReceivedAuthRequest(AuthRequestContext) + SentChallenge(ChallengeContext) / move_challenge = WaitingForChallengeSolution(ChallengeContext),
+
+        WaitingForChallengeSolution(ChallengeContext) + ReceivedGoodSolution = Idle,
+        WaitingForChallengeSolution(ChallengeContext) + ReceivedBadSolution = AuthError, // block further transitions, but connection should close anyway
+
+        Idle + UnsealRequest(UnsealContext) / generate_temp_keypair = WaitingForUnsealKey(UnsealContext),
+        WaitingForUnsealKey(UnsealContext) + ReceivedValidKey = Unsealed,
+        WaitingForUnsealKey(UnsealContext) + ReceivedInvalidKey = Idle,
+    }
+);
+
+pub struct DummyContext;
+impl UserAgentStateMachineContext for DummyContext {
+    #[allow(missing_docs)]
+    #[allow(clippy::unused_unit)]
+    fn move_challenge(
+        &mut self,
+        _state_data: &AuthRequestContext,
+        event_data: ChallengeContext,
+    ) -> Result<ChallengeContext, ()> {
+        Ok(event_data)
+    }
+
+    #[allow(missing_docs)]
+    #[allow(clippy::unused_unit)]
+    fn auth_request_context(
+        &mut self,
+        event_data: AuthRequestContext,
+    ) -> Result<AuthRequestContext, ()> {
+        Ok(event_data)
+    }
+
+    #[allow(missing_docs)]
+    #[allow(clippy::unused_unit)]
+    fn generate_temp_keypair(&mut self, event_data: UnsealContext) -> Result<UnsealContext, ()> {
+        Ok(event_data)
+    }
+}

server/crates/arbiter-server/src/actors/user_agent/tests.rs

@@ -0,0 +1,199 @@
+use arbiter_proto::proto::{
+    UserAgentResponse,
+    auth::{self, AuthChallengeRequest, AuthOk},
+    user_agent_response::Payload as UserAgentResponsePayload,
+};
+use chrono::format;
+use diesel::{ExpressionMethods as _, QueryDsl, insert_into};
+use diesel_async::RunQueryDsl;
+use ed25519_dalek::Signer as _;
+use kameo::actor::Spawn;
+
+use crate::{
+    actors::{
+        bootstrap::BootstrapActor,
+        user_agent::{HandleAuthChallengeRequest, HandleAuthChallengeSolution},
+    },
+    db::{self, schema},
+};
+
+use super::UserAgentActor;
+
+#[tokio::test]
+#[test_log::test]
+pub async fn test_bootstrap_token_auth() {
+    let db = db::create_test_pool().await;
+    // explicitly not installing any user_agent pubkeys
+    let bootstrapper = BootstrapActor::new(&db).await.unwrap(); // this will create bootstrap token
+    let token = bootstrapper.get_token().unwrap();
+
+    let bootstrapper_ref = BootstrapActor::spawn(bootstrapper);
+    let user_agent = UserAgentActor::new_manual(
+        db.clone(),
+        bootstrapper_ref,
+        tokio::sync::mpsc::channel(1).0, // dummy channel, we won't actually send responses in this test
+    );
+    let user_agent_ref = UserAgentActor::spawn(user_agent);
+
+    // simulate client sending auth request with bootstrap token
+    let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
+    let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
+
+    let result = user_agent_ref
+        .ask(HandleAuthChallengeRequest {
+            req: AuthChallengeRequest {
+                pubkey: pubkey_bytes,
+                bootstrap_token: Some(token),
+            },
+        })
+        .await
+        .expect("Shouldn't fail to send message");
+
+    // auth succeeded
+    assert_eq!(
+        result,
+        UserAgentResponse {
+            payload: Some(UserAgentResponsePayload::AuthMessage(
+                arbiter_proto::proto::auth::ServerMessage {
+                    payload: Some(arbiter_proto::proto::auth::server_message::Payload::AuthOk(
+                        AuthOk {},
+                    )),
+                },
+            )),
+        }
+    );
+
+    // key is succesfully recorded in database
+    let mut conn = db.get().await.unwrap();
+    let stored_pubkey: Vec<u8> = schema::useragent_client::table
+        .select(schema::useragent_client::public_key)
+        .first::<Vec<u8>>(&mut conn)
+        .await
+        .unwrap();
+    assert_eq!(stored_pubkey, new_key.verifying_key().to_bytes().to_vec());
+}
+
+#[tokio::test]
+#[test_log::test]
+pub async fn test_bootstrap_invalid_token_auth() {
+    let db = db::create_test_pool().await;
+    // explicitly not installing any user_agent pubkeys
+    let bootstrapper = BootstrapActor::new(&db).await.unwrap(); // this will create bootstrap token
+
+    let bootstrapper_ref = BootstrapActor::spawn(bootstrapper);
+    let user_agent = UserAgentActor::new_manual(
+        db.clone(),
+        bootstrapper_ref,
+        tokio::sync::mpsc::channel(1).0, // dummy channel, we won't actually send responses in this test
+    );
+    let user_agent_ref = UserAgentActor::spawn(user_agent);
+
+    // simulate client sending auth request with bootstrap token
+    let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
+    let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
+
+    let result = user_agent_ref
+        .ask(HandleAuthChallengeRequest {
+            req: AuthChallengeRequest {
+                pubkey: pubkey_bytes,
+                bootstrap_token: Some("invalid_token".to_string()),
+            },
+        })
+        .await;
+
+    match result {
+        Err(kameo::error::SendError::HandlerError(status)) => {
+            assert_eq!(status.code(), tonic::Code::InvalidArgument);
+            insta::assert_debug_snapshot!(status, @r#"
+                Status {
+                    code: InvalidArgument,
+                    message: "Invalid bootstrap token",
+                    source: None,
+                }
+                "#);
+        }
+        Err(other) => {
+            panic!("Expected SendError::HandlerError, got {other:?}");
+        }
+        Ok(_) => {
+            panic!("Expected error due to invalid bootstrap token, but got success");
+        }
+    }
+}
+
+#[tokio::test]
+#[test_log::test]
+pub async fn test_challenge_auth() {
+    let db = db::create_test_pool().await;
+
+    let bootstrapper_ref = BootstrapActor::spawn(BootstrapActor::new(&db).await.unwrap());
+    let user_agent = UserAgentActor::new_manual(
+        db.clone(),
+        bootstrapper_ref,
+        tokio::sync::mpsc::channel(1).0, // dummy channel, we won't actually send responses in this test
+    );
+    let user_agent_ref = UserAgentActor::spawn(user_agent);
+
+    // simulate client sending auth request with bootstrap token
+    let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
+    let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
+
+    // insert pubkey into database to trigger challenge-response auth flow
+    {
+        let mut conn = db.get().await.unwrap();
+        insert_into(schema::useragent_client::table)
+            .values(schema::useragent_client::public_key.eq(pubkey_bytes.clone()))
+            .execute(&mut conn)
+            .await
+            .unwrap();
+    }
+
+    let result = user_agent_ref
+        .ask(HandleAuthChallengeRequest {
+            req: AuthChallengeRequest {
+                pubkey: pubkey_bytes,
+                bootstrap_token: None,
+            },
+        })
+        .await
+        .expect("Shouldn't fail to send message");
+
+    // auth challenge succeeded
+    let UserAgentResponse {
+        payload:
+            Some(UserAgentResponsePayload::AuthMessage(arbiter_proto::proto::auth::ServerMessage {
+                payload:
+                    Some(arbiter_proto::proto::auth::server_message::Payload::AuthChallenge(challenge)),
+            })),
+    } = result
+    else {
+        panic!("Expected auth challenge response, got {result:?}");
+    };
+
+    let formatted_challenge = arbiter_proto::format_challenge(&challenge);
+    let signature = new_key.sign(&formatted_challenge);
+    let serialized_signature = signature.to_bytes().to_vec();
+
+    let result = user_agent_ref
+        .ask(HandleAuthChallengeSolution {
+            solution: auth::AuthChallengeSolution {
+                signature: serialized_signature,
+            },
+        })
+        .await
+        .expect("Shouldn't fail to send message");
+
+    // auth succeeded
+    assert_eq!(
+        result,
+        UserAgentResponse {
+            payload: Some(UserAgentResponsePayload::AuthMessage(
+                arbiter_proto::proto::auth::ServerMessage {
+                    payload: Some(arbiter_proto::proto::auth::server_message::Payload::AuthOk(
+                        AuthOk {},
+                    )),
+                },
+            )),
+        }
+    );
+}

server/crates/arbiter-server/src/actors/user_agent/transport.rs

@@ -2,12 +2,9 @@ use super::UserAgentActor;
 use arbiter_proto::proto::{
     UserAgentRequest, UserAgentResponse,
     auth::{
-        self, AuthChallenge, AuthChallengeRequest, AuthOk, ClientMessage,
-        ServerMessage as AuthServerMessage, client_message::Payload as ClientAuthPayload,
-        server_message::Payload as ServerAuthPayload,
+        ClientMessage as ClientAuthMessage, client_message::Payload as ClientAuthPayload,
     },
     user_agent_request::Payload as UserAgentRequestPayload,
-    user_agent_response::Payload as UserAgentResponsePayload,
 };
 use futures::StreamExt;
 use kameo::{
@@ -19,7 +16,10 @@ use tonic::Status;
 use tracing::error;
 
 use crate::{
-    actors::user_agent::{HandleAuthChallengeRequest, HandleAuthChallengeSolution},
+    actors::user_agent::{
+        HandleAuthChallengeRequest, HandleAuthChallengeSolution, HandleUnsealEncryptedKey,
+        HandleUnsealRequest,
+    },
     context::ServerContext,
 };
 
@@ -59,28 +59,30 @@ async fn process_message(
         Status::invalid_argument("Expected message with payload")
     })?;
 
-    let UserAgentRequestPayload::AuthMessage(ClientMessage {
-        payload: Some(client_message),
-    }) = msg
-    else {
-        error!(
-            actor = "useragent",
-            "Received unexpected message type during authentication"
-        );
-        return Err(Status::invalid_argument(
-            "Expected AuthMessage with ClientMessage payload",
-        ));
-    };
-
-    match client_message {
-        ClientAuthPayload::AuthChallengeRequest(req) => actor
+    match msg {
+        UserAgentRequestPayload::AuthMessage(ClientAuthMessage {
+            payload: Some(ClientAuthPayload::AuthChallengeRequest(req)),
+        }) => actor
             .ask(HandleAuthChallengeRequest { req })
             .await
             .map_err(into_status),
-        ClientAuthPayload::AuthChallengeSolution(solution) => actor
+        UserAgentRequestPayload::AuthMessage(ClientAuthMessage {
+            payload: Some(ClientAuthPayload::AuthChallengeSolution(solution)),
+        }) => actor
             .ask(HandleAuthChallengeSolution { solution })
             .await
             .map_err(into_status),
+        UserAgentRequestPayload::UnsealStart(unseal_start) => actor
+            .ask(HandleUnsealRequest { req: unseal_start })
+            .await
+            .map_err(into_status),
+        UserAgentRequestPayload::UnsealEncryptedKey(unseal_encrypted_key) => actor
+            .ask(HandleUnsealEncryptedKey {
+                req: unseal_encrypted_key,
+            })
+            .await
+            .map_err(into_status),
+        _ => Err(Status::invalid_argument("Expected message with payload")),
     }
 }
 

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

@@ -2,7 +2,6 @@ use std::sync::Arc;
 
 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;
@@ -14,7 +13,7 @@ use crate::{
     actors::bootstrap::{self, BootstrapActor}, context::tls::{TlsDataRaw, TlsManager}, db::{
         self,
         models::ArbiterSetting,
-        schema::{self, arbiter_settings},
+        schema::arbiter_settings,
     }
 };
 

server/crates/arbiter-server/src/db.rs

@@ -1,12 +1,11 @@
-use std::sync::Arc;
 
 use diesel::{
     Connection as _, SqliteConnection,
-    connection::{SimpleConnection as _, TransactionManager},
+    connection::SimpleConnection as _,
 };
 use diesel_async::{
     AsyncConnection, SimpleAsyncConnection,
-    pooled_connection::{AsyncDieselConnectionManager, ManagerConfig, RecyclingMethod},
+    pooled_connection::{AsyncDieselConnectionManager, ManagerConfig},
     sync_connection_wrapper::SyncConnectionWrapper,
 };
 use diesel_migrations::{EmbeddedMigrations, MigrationHarness, embed_migrations};

server/crates/arbiter-server/src/db/models.rs

@@ -1,29 +1,55 @@
 #![allow(unused)]
 #![allow(clippy::all)]
 
-use crate::db::schema::{self, aead_encrypted, arbiter_settings};
+use crate::db::schema::{self, aead_encrypted, arbiter_settings, root_key_history};
 use diesel::{prelude::*, sqlite::Sqlite};
+use restructed::Models;
 
 pub mod types {
     use chrono::{DateTime, Utc};
     pub struct SqliteTimestamp(DateTime<Utc>);
 }
 
-#[derive(Queryable, Debug, Insertable)]
+#[derive(Models, Queryable, Debug, Insertable, Selectable)]
+#[view(
+    NewAeadEncrypted,
+    derive(Insertable),
+    omit(id),
+    attributes_with = "deriveless"
+)]
 #[diesel(table_name = aead_encrypted, check_for_backend(Sqlite))]
 pub struct AeadEncrypted {
     pub id: i32,
     pub ciphertext: Vec<u8>,
     pub tag: Vec<u8>,
-    pub current_nonce: i32,
+    pub current_nonce: Vec<u8>,
     pub schema_version: i32,
+    pub created_at: i32,
+}
+
+#[derive(Models, Queryable, Debug, Insertable, Selectable)]
+#[diesel(table_name = root_key_history, check_for_backend(Sqlite))]
+#[view(
+    NewRootKeyHistory,
+    derive(Insertable),
+    omit(id),
+    attributes_with = "deriveless"
+)]
+pub struct RootKeyHistory {
+    pub id: i32,
+    pub ciphertext: Vec<u8>,
+    pub tag: Vec<u8>,
+    pub root_key_encryption_nonce: Vec<u8>,
+    pub data_encryption_nonce: Vec<u8>,
+    pub schema_version: i32,
+    pub salt: Vec<u8>,
 }
 
 #[derive(Queryable, Debug, Insertable)]
 #[diesel(table_name = arbiter_settings, check_for_backend(Sqlite))]
 pub struct ArbiterSetting {
     pub id: i32,
-    pub root_key_id: Option<i32>, // references aead_encrypted.id
+    pub root_key_id: Option<i32>, // references root_key_history.id
     pub cert_key: Vec<u8>,
     pub cert: Vec<u8>,
 }

server/crates/arbiter-server/src/db/schema.rs

@@ -3,10 +3,11 @@
 diesel::table! {
     aead_encrypted (id) {
         id -> Integer,
-        current_nonce -> Integer,
+        current_nonce -> Binary,
         ciphertext -> Binary,
         tag -> Binary,
         schema_version -> Integer,
+        created_at -> Integer,
     }
 }
 
@@ -29,6 +30,18 @@ diesel::table! {
     }
 }
 
+diesel::table! {
+    root_key_history (id) {
+        id -> Integer,
+        root_key_encryption_nonce -> Binary,
+        data_encryption_nonce -> Binary,
+        ciphertext -> Binary,
+        tag -> Binary,
+        schema_version -> Integer,
+        salt -> Binary,
+    }
+}
+
 diesel::table! {
     useragent_client (id) {
         id -> Integer,
@@ -39,11 +52,12 @@ diesel::table! {
     }
 }
 
-diesel::joinable!(arbiter_settings -> aead_encrypted (root_key_id));
+diesel::joinable!(arbiter_settings -> root_key_history (root_key_id));
 
 diesel::allow_tables_to_appear_in_same_query!(
     aead_encrypted,
     arbiter_settings,
     program_client,
+    root_key_history,
     useragent_client,
 );