feat(unseal): add unseal protocol support for user agents

protobufs/arbiter.proto

@@ -3,29 +3,8 @@ syntax = "proto3";
 package arbiter;
 
 import "auth.proto";
-
+import "client.proto";
-message ClientRequest {
+import "user_agent.proto";
-  oneof payload {
-    arbiter.auth.ClientMessage auth_message = 1;
-  }
-}
-
-message ClientResponse {
-  oneof payload {
-    arbiter.auth.ServerMessage auth_message = 1;
-  }
-}
-
-message UserAgentRequest {
-  oneof payload {
-    arbiter.auth.ClientMessage auth_message = 1;
-  }
-}
-message UserAgentResponse {
-  oneof payload {
-    arbiter.auth.ServerMessage auth_message = 1;
-  }
-}
 
 message ServerInfo {
   string version = 1;

protobufs/client.proto

@@ -0,0 +1,17 @@
+syntax = "proto3";
+
+package arbiter;
+
+import "auth.proto";
+
+message ClientRequest {
+  oneof payload {
+    arbiter.auth.ClientMessage auth_message = 1;
+  }
+}
+
+message ClientResponse {
+  oneof payload {
+    arbiter.auth.ServerMessage auth_message = 1;
+  }
+}

protobufs/unseal.proto

@@ -2,13 +2,20 @@ syntax = "proto3";
 
 package arbiter.unseal;
 
-message UserAgentKeyRequest {}
+import "google/protobuf/empty.proto";
 
-message ServerKeyResponse {
+message UnsealStart {}
+
+message UnsealStartResponse {
   bytes pubkey = 1;
 }
-message UserAgentSealedKey {
+message UnsealEncryptedKey {
-  bytes sealed_key = 1;
+  bytes key = 1;
-  bytes pubkey = 2;
+}
-  bytes nonce = 3;
+
+enum UnsealResult {
+  UNSEAL_RESULT_UNSPECIFIED = 0;
+  UNSEAL_RESULT_SUCCESS = 1;
+  UNSEAL_RESULT_INVALID_KEY = 2;
+  UNSEAL_RESULT_UNBOOTSTRAPPED = 3;
 }

protobufs/user_agent.proto

@@ -0,0 +1,21 @@
+syntax = "proto3";
+
+package arbiter;
+
+import "auth.proto";
+import "unseal.proto";
+
+message UserAgentRequest {
+  oneof payload {
+    arbiter.auth.ClientMessage auth_message = 1;
+    arbiter.unseal.UnsealStart unseal_start = 2;
+    arbiter.unseal.UnsealEncryptedKey unseal_encrypted_key = 3;
+  }
+}
+message UserAgentResponse {
+  oneof payload {
+    arbiter.auth.ServerMessage auth_message = 1;
+    arbiter.unseal.UnsealStartResponse unseal_start_response = 2;
+    arbiter.unseal.UnsealResult unseal_result = 3;
+  }
+}

server/crates/arbiter-proto/src/lib.rs

@@ -6,6 +6,9 @@ pub mod proto {
     pub mod auth {
         tonic::include_proto!("arbiter.auth");
     }
+    pub mod unseal {
+        tonic::include_proto!("arbiter.unseal");
+    }
 }
 
 pub mod transport;

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

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

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

@@ -1,5 +1,5 @@
 use arbiter_proto::proto::{
-    UserAgentRequest, UserAgentResponse,
+    UserAgentResponse,
     auth::{
         self, AuthChallenge, AuthChallengeRequest, AuthOk, ClientMessage,
         ServerMessage as AuthServerMessage, client_message::Payload as ClientAuthPayload,
@@ -11,22 +11,19 @@ use arbiter_proto::proto::{
 use diesel::{ExpressionMethods as _, OptionalExtension as _, QueryDsl, dsl::update};
 use diesel_async::{AsyncConnection, RunQueryDsl};
 use ed25519_dalek::VerifyingKey;
-use futures::StreamExt;
 use kameo::{
     Actor,
     actor::{ActorRef, Spawn},
-    error::SendError,
     messages,
     prelude::Context,
 };
-use tokio::sync::mpsc;
 use tokio::sync::mpsc::Sender;
 use tonic::Status;
 use tracing::{error, info};
 
 use crate::{
     ServerContext,
-    context::bootstrap::{BootstrapActor, ConsumeToken},
+    actors::bootstrap::{BootstrapActor, ConsumeToken},
     db::{self, schema},
     errors::GrpcStatusExt,
 };
@@ -54,12 +51,15 @@ smlang::statemachine!(
     custom_error: false,
     transitions: {
         *Init + AuthRequest(AuthRequestContext) / auth_request_context =  ReceivedAuthRequest(AuthRequestContext),
-        ReceivedAuthRequest(AuthRequestContext) + ReceivedBootstrapToken = Authenticated,
+        ReceivedAuthRequest(AuthRequestContext) + ReceivedBootstrapToken = Idle,
 
         ReceivedAuthRequest(AuthRequestContext) + SentChallenge(ChallengeContext) / move_challenge = WaitingForChallengeSolution(ChallengeContext),
 
-        WaitingForChallengeSolution(ChallengeContext) + ReceivedGoodSolution = Authenticated,
+        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),
     }
 );
 
@@ -69,7 +69,7 @@ impl UserAgentStateMachineContext for DummyContext {
     #[allow(clippy::unused_unit)]
     fn move_challenge(
         &mut self,
-        state_data: &AuthRequestContext,
+        _state_data: &AuthRequestContext,
         event_data: ChallengeContext,
     ) -> Result<ChallengeContext, ()> {
         Ok(event_data)
@@ -83,6 +83,21 @@ impl UserAgentStateMachineContext for DummyContext {
     ) -> 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()))
+    }
 }
 
 #[derive(Actor)]
@@ -90,7 +105,8 @@ pub struct UserAgentActor {
     db: db::DatabasePool,
     bootstapper: ActorRef<BootstrapActor>,
     state: UserAgentStateMachine<DummyContext>,
-    tx: Sender<Result<UserAgentResponse, Status>>,
+    // will be used in future
+    _tx: Sender<Result<UserAgentResponse, Status>>,
 }
 
 impl UserAgentActor {
@@ -102,10 +118,11 @@ impl UserAgentActor {
             db: context.db.clone(),
             bootstapper: context.bootstrapper.clone(),
             state: UserAgentStateMachine::new(DummyContext),
-            tx,
+            _tx: tx,
         }
     }
 
+    #[cfg(test)]
     pub(crate) fn new_manual(
         db: db::DatabasePool,
         bootstapper: ActorRef<BootstrapActor>,
@@ -115,7 +132,7 @@ impl UserAgentActor {
             db,
             bootstapper,
             state: UserAgentStateMachine::new(DummyContext),
-            tx,
+            _tx: tx,
         }
     }
 
@@ -319,8 +336,10 @@ mod tests {
     use kameo::actor::Spawn;
 
     use crate::{
-        actors::user_agent::{HandleAuthChallengeRequest, HandleAuthChallengeSolution},
+        actors::{
-        context::bootstrap::BootstrapActor,
+            bootstrap::BootstrapActor,
+            user_agent::{HandleAuthChallengeRequest, HandleAuthChallengeSolution},
+        },
         db::{self, schema},
     };
 
@@ -449,7 +468,7 @@ mod tests {
         {
             let mut conn = db.get().await.unwrap();
             insert_into(schema::useragent_client::table)
-                .values((schema::useragent_client::public_key.eq(pubkey_bytes.clone())))
+                .values(schema::useragent_client::public_key.eq(pubkey_bytes.clone()))
                 .execute(&mut conn)
                 .await
                 .unwrap();

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

@@ -1,41 +0,0 @@
-use std::sync::Arc;
-
-use dashmap::DashSet;
-
-#[derive(Clone, Default)]
-struct LeaseStorage<T: Eq + std::hash::Hash>(Arc<DashSet<T>>);
-
-// A lease that automatically releases the item when dropped
-pub struct Lease<T: Clone + std::hash::Hash + Eq> {
-    item: T,
-    storage: LeaseStorage<T>,
-}
-impl<T: Clone + std::hash::Hash + Eq> Drop for Lease<T> {
-    fn drop(&mut self) {
-        self.storage.0.remove(&self.item);
-    }
-}
-
-#[derive(Clone, Default)]
-pub struct LeaseHandler<T: Clone + std::hash::Hash + Eq> {
-    storage: LeaseStorage<T>,
-}
-
-impl<T: Clone + std::hash::Hash + Eq> LeaseHandler<T> {
-    pub fn new() -> Self {
-        Self {
-            storage: LeaseStorage(Arc::new(DashSet::new())),
-        }
-    }
-
-    pub fn acquire(&self, item: T) -> Result<Lease<T>, ()> {
-        if self.storage.0.insert(item.clone()) {
-            Ok(Lease {
-                item,
-                storage: self.storage.clone(),
-            })
-        } else {
-            Err(())
-        }
-    }
-}