feat(user-agent-auth): add RSA and ECDSA auth key types

Extend user-agent authentication to support Ed25519, ECDSA (secp256k1), and RSA (PSS+SHA-256) with minimal protocol and storage changes. Add key_type to auth requests and useragent_client, update key parsing/signature verification paths, and keep backward compatibility by treating UNSPECIFIED as Ed25519.

server/crates/arbiter-server/Cargo.toml

@@ -42,7 +42,10 @@ argon2 = { version = "0.5.3", features = ["zeroize"] }
 restructed = "0.2.2"
 strum = { version = "0.27.2", features = ["derive"] }
 pem = "3.0.6"
-k256 = "0.13.4"
+k256.workspace = true
+rsa.workspace = true
+sha2.workspace = true
+spki.workspace = true
 alloy.workspace = true
 arbiter-tokens-registry.path = "../arbiter-tokens-registry"
 

server/crates/arbiter-server/migrations/2026-03-13-000000_add_key_type/down.sql

@@ -0,0 +1,2 @@
+-- Not reversible without data loss; drop the column to revert
+ALTER TABLE useragent_client DROP COLUMN key_type;

server/crates/arbiter-server/migrations/2026-03-13-000000_add_key_type/up.sql

@@ -0,0 +1 @@
+ALTER TABLE useragent_client ADD COLUMN key_type INTEGER NOT NULL DEFAULT 1;

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

@@ -1,13 +1,12 @@
 use arbiter_proto::proto::user_agent::{
-    AuthChallengeRequest, AuthChallengeSolution, UserAgentRequest,
+    AuthChallengeRequest, AuthChallengeSolution, KeyType as ProtoKeyType, UserAgentRequest,
     user_agent_request::Payload as UserAgentRequestPayload,
 };
-use ed25519_dalek::VerifyingKey;
 use tracing::error;
 
 use crate::actors::user_agent::{
     UserAgentConnection,
-    auth::state::{AuthContext, AuthStateMachine}, session::UserAgentSession,
+    auth::state::{AuthContext, AuthPublicKey, AuthStateMachine}, session::UserAgentSession,
 };
 
 #[derive(thiserror::Error, Debug, PartialEq)]
@@ -37,28 +36,50 @@ pub enum Error {
 mod state;
 use state::*;
 
+fn parse_pubkey(key_type: ProtoKeyType, pubkey: Vec<u8>) -> Result<AuthPublicKey, Error> {
+    match key_type {
+        // UNSPECIFIED treated as Ed25519 for backward compatibility
+        ProtoKeyType::Unspecified | ProtoKeyType::Ed25519 => {
+            let pubkey_bytes = pubkey.as_array().ok_or(Error::InvalidClientPubkeyLength)?;
+            let key = ed25519_dalek::VerifyingKey::from_bytes(pubkey_bytes)
+                .map_err(|_| Error::InvalidAuthPubkeyEncoding)?;
+            Ok(AuthPublicKey::Ed25519(key))
+        }
+        ProtoKeyType::EcdsaSecp256k1 => {
+            // Public key is sent as 33-byte SEC1 compressed point
+            let key = k256::ecdsa::VerifyingKey::from_sec1_bytes(&pubkey)
+                .map_err(|_| Error::InvalidAuthPubkeyEncoding)?;
+            Ok(AuthPublicKey::EcdsaSecp256k1(key))
+        }
+        ProtoKeyType::Rsa => {
+            use rsa::pkcs8::DecodePublicKey as _;
+            let key = rsa::RsaPublicKey::from_public_key_der(&pubkey)
+                .map_err(|_| Error::InvalidAuthPubkeyEncoding)?;
+            Ok(AuthPublicKey::Rsa(key))
+        }
+    }
+}
+
 fn parse_auth_event(payload: UserAgentRequestPayload) -> Result<AuthEvents, Error> {
     match payload {
         UserAgentRequestPayload::AuthChallengeRequest(AuthChallengeRequest {
             pubkey,
             bootstrap_token: None,
+            key_type,
         }) => {
-            let pubkey_bytes = pubkey.as_array().ok_or(Error::InvalidClientPubkeyLength)?;
-            let pubkey = VerifyingKey::from_bytes(pubkey_bytes)
-                .map_err(|_| Error::InvalidAuthPubkeyEncoding)?;
+            let kt = ProtoKeyType::try_from(key_type).unwrap_or(ProtoKeyType::Unspecified);
             Ok(AuthEvents::AuthRequest(ChallengeRequest {
-                pubkey: pubkey.into(),
+                pubkey: parse_pubkey(kt, pubkey)?,
             }))
         }
         UserAgentRequestPayload::AuthChallengeRequest(AuthChallengeRequest {
             pubkey,
             bootstrap_token: Some(token),
+            key_type,
         }) => {
-            let pubkey_bytes = pubkey.as_array().ok_or(Error::InvalidClientPubkeyLength)?;
-            let pubkey = VerifyingKey::from_bytes(pubkey_bytes)
-                .map_err(|_| Error::InvalidAuthPubkeyEncoding)?;
+            let kt = ProtoKeyType::try_from(key_type).unwrap_or(ProtoKeyType::Unspecified);
             Ok(AuthEvents::BootstrapAuthRequest(BootstrapAuthRequest {
-                pubkey: pubkey.into(),
+                pubkey: parse_pubkey(kt, pubkey)?,
                 token,
             }))
         }
@@ -71,11 +92,11 @@ fn parse_auth_event(payload: UserAgentRequestPayload) -> Result<AuthEvents, Erro
     }
 }
 
-pub async fn authenticate(props: &mut UserAgentConnection) -> Result<VerifyingKey, Error> {
+pub async fn authenticate(props: &mut UserAgentConnection) -> Result<AuthPublicKey, Error> {
     let mut state = AuthStateMachine::new(AuthContext::new(props));
 
     loop {
-        // This is needed because `state` now holds mutable reference to `ConnectionProps`, so we can't directly access `props` here
+        // `state` holds a mutable reference to `props` so we can't access it directly here
         let transport = state.context_mut().conn.transport.as_mut();
         let Some(UserAgentRequest {
             payload: Some(payload),
@@ -110,9 +131,8 @@ pub async fn authenticate(props: &mut UserAgentConnection) -> Result<VerifyingKe
     }
 }
 
-
 pub async fn authenticate_and_create(mut props: UserAgentConnection) -> Result<UserAgentSession, Error> {
-    let key = authenticate(&mut props).await?;
-    let session = UserAgentSession::new(props, key.clone());
+    let _key = authenticate(&mut props).await?;
+    let session = UserAgentSession::new(props);
     Ok(session)
 }

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

@@ -1,30 +1,64 @@
 use arbiter_proto::proto::user_agent::{
-    AuthChallenge, UserAgentResponse,
-    user_agent_response::Payload as UserAgentResponsePayload,
+    AuthChallenge, UserAgentResponse, user_agent_response::Payload as UserAgentResponsePayload,
 };
 use diesel::{ExpressionMethods as _, OptionalExtension as _, QueryDsl, update};
 use diesel_async::RunQueryDsl;
-use ed25519_dalek::VerifyingKey;
 use tracing::error;
 
 use super::Error;
 use crate::{
     actors::{bootstrap::ConsumeToken, user_agent::UserAgentConnection},
-    db::schema,
+    db::{models::KeyType, schema},
 };
 
+/// Abstraction over Ed25519 / ECDSA-secp256k1 / RSA public keys used during the auth handshake.
+#[derive(Clone)]
+pub enum AuthPublicKey {
+    Ed25519(ed25519_dalek::VerifyingKey),
+    /// Compressed SEC1 public key; signature bytes are raw 64-byte (r||s).
+    EcdsaSecp256k1(k256::ecdsa::VerifyingKey),
+    /// RSA-2048+ public key; signature bytes are PSS+SHA-256.
+    Rsa(rsa::RsaPublicKey),
+}
+
+impl AuthPublicKey {
+    /// Canonical bytes stored in DB and echoed back in the challenge.
+    /// Ed25519: raw 32 bytes. ECDSA: SEC1 compressed 33 bytes. RSA: DER-encoded SPKI.
+    pub fn to_stored_bytes(&self) -> Vec<u8> {
+        match self {
+            AuthPublicKey::Ed25519(k) => k.to_bytes().to_vec(),
+            // SEC1 compressed (33 bytes) is the natural compact format for secp256k1
+            AuthPublicKey::EcdsaSecp256k1(k) => k.to_encoded_point(true).as_bytes().to_vec(),
+            AuthPublicKey::Rsa(k) => {
+                use rsa::pkcs8::EncodePublicKey as _;
+                k.to_public_key_der()
+                    .expect("rsa SPKI encoding is infallible")
+                    .to_vec()
+            }
+        }
+    }
+
+    pub fn key_type(&self) -> KeyType {
+        match self {
+            AuthPublicKey::Ed25519(_) => KeyType::Ed25519,
+            AuthPublicKey::EcdsaSecp256k1(_) => KeyType::EcdsaSecp256k1,
+            AuthPublicKey::Rsa(_) => KeyType::Rsa,
+        }
+    }
+}
+
 pub struct ChallengeRequest {
-    pub pubkey: VerifyingKey,
+    pub pubkey: AuthPublicKey,
 }
 
 pub struct BootstrapAuthRequest {
-    pub pubkey: VerifyingKey,
+    pub pubkey: AuthPublicKey,
     pub token: String,
 }
 
 pub struct ChallengeContext {
     pub challenge: AuthChallenge,
-    pub key: VerifyingKey,
+    pub key: AuthPublicKey,
 }
 
 pub struct ChallengeSolution {
@@ -36,8 +70,8 @@ smlang::statemachine!(
     custom_error: true,
     transitions: {
         *Init + AuthRequest(ChallengeRequest) / async prepare_challenge = SentChallenge(ChallengeContext),
-        Init + BootstrapAuthRequest(BootstrapAuthRequest) [async verify_bootstrap_token] / provide_key_bootstrap = AuthOk(VerifyingKey),
-        SentChallenge(ChallengeContext) + ReceivedSolution(ChallengeSolution) [async verify_solution] / provide_key = AuthOk(VerifyingKey),
+        Init + BootstrapAuthRequest(BootstrapAuthRequest) [async verify_bootstrap_token] / provide_key_bootstrap = AuthOk(AuthPublicKey),
+        SentChallenge(ChallengeContext) + ReceivedSolution(ChallengeSolution) [async verify_solution] / provide_key = AuthOk(AuthPublicKey),
     }
 );
 
@@ -76,7 +110,9 @@ async fn create_nonce(db: &crate::db::DatabasePool, pubkey_bytes: &[u8]) -> Resu
         })
 }
 
-async fn register_key(db: &crate::db::DatabasePool, pubkey_bytes: &[u8]) -> Result<(), Error> {
+async fn register_key(db: &crate::db::DatabasePool, pubkey: &AuthPublicKey) -> Result<(), Error> {
+    let pubkey_bytes = pubkey.to_stored_bytes();
+    let key_type = pubkey.key_type();
     let mut conn = db.get().await.map_err(|e| {
         error!(error = ?e, "Database pool error");
         Error::DatabasePoolUnavailable
@@ -84,8 +120,9 @@ async fn register_key(db: &crate::db::DatabasePool, pubkey_bytes: &[u8]) -> Resu
 
     diesel::insert_into(schema::useragent_client::table)
         .values((
-            schema::useragent_client::public_key.eq(pubkey_bytes.to_vec()),
+            schema::useragent_client::public_key.eq(pubkey_bytes),
             schema::useragent_client::nonce.eq(1),
+            schema::useragent_client::key_type.eq(key_type),
         ))
         .execute(&mut conn)
         .await
@@ -115,15 +152,34 @@ impl AuthStateMachineContext for AuthContext<'_> {
         ChallengeContext { challenge, key }: &ChallengeContext,
         ChallengeSolution { solution }: &ChallengeSolution,
     ) -> Result<bool, Self::Error> {
-        let formatted_challenge =
-            arbiter_proto::format_challenge(challenge.nonce, &challenge.pubkey);
+        let formatted = arbiter_proto::format_challenge(challenge.nonce, &challenge.pubkey);
 
-        let signature = solution.as_slice().try_into().map_err(|_| {
-            error!(?solution, "Invalid signature length");
-            Error::InvalidChallengeSolution
-        })?;
-
-        let valid = key.verify_strict(&formatted_challenge, &signature).is_ok();
+        let valid = match key {
+            AuthPublicKey::Ed25519(vk) => {
+                let sig = solution.as_slice().try_into().map_err(|_| {
+                    error!(?solution, "Invalid Ed25519 signature length");
+                    Error::InvalidChallengeSolution
+                })?;
+                vk.verify_strict(&formatted, &sig).is_ok()
+            }
+            AuthPublicKey::EcdsaSecp256k1(vk) => {
+                use k256::ecdsa::signature::Verifier as _;
+                let sig = k256::ecdsa::Signature::try_from(solution.as_slice()).map_err(|_| {
+                    error!(?solution, "Invalid ECDSA signature bytes");
+                    Error::InvalidChallengeSolution
+                })?;
+                vk.verify(&formatted, &sig).is_ok()
+            }
+            AuthPublicKey::Rsa(pk) => {
+                use rsa::signature::Verifier as _;
+                let verifying_key = rsa::pss::VerifyingKey::<sha2::Sha256>::new(pk.clone());
+                let sig = rsa::pss::Signature::try_from(solution.as_slice()).map_err(|_| {
+                    error!(?solution, "Invalid RSA signature bytes");
+                    Error::InvalidChallengeSolution
+                })?;
+                verifying_key.verify(&formatted, &sig).is_ok()
+            }
+        };
 
         Ok(valid)
     }
@@ -132,10 +188,11 @@ impl AuthStateMachineContext for AuthContext<'_> {
         &mut self,
         ChallengeRequest { pubkey }: ChallengeRequest,
     ) -> Result<ChallengeContext, Self::Error> {
-        let nonce = create_nonce(&self.conn.db, pubkey.as_bytes()).await?;
+        let stored_bytes = pubkey.to_stored_bytes();
+        let nonce = create_nonce(&self.conn.db, &stored_bytes).await?;
 
         let challenge = AuthChallenge {
-            pubkey: pubkey.as_bytes().to_vec(),
+            pubkey: stored_bytes,
             nonce,
         };
 
@@ -171,16 +228,16 @@ impl AuthStateMachineContext for AuthContext<'_> {
             })
             .await
             .map_err(|e| {
-                error!(?pubkey, "Failed to consume bootstrap token: {e}");
+                error!(?e, "Failed to consume bootstrap token");
                 Error::BootstrapperActorUnreachable
             })?;
 
         if !token_ok {
-            error!(?pubkey, "Invalid bootstrap token provided");
+            error!("Invalid bootstrap token provided");
             return Err(Error::InvalidBootstrapToken);
         }
 
-        register_key(&self.conn.db, pubkey.as_bytes()).await?;
+        register_key(&self.conn.db, pubkey).await?;
 
         Ok(true)
     }
@@ -188,7 +245,7 @@ impl AuthStateMachineContext for AuthContext<'_> {
     fn provide_key_bootstrap(
         &mut self,
         event_data: BootstrapAuthRequest,
-    ) -> Result<VerifyingKey, Self::Error> {
+    ) -> Result<AuthPublicKey, Self::Error> {
         Ok(event_data.pubkey)
     }
 
@@ -196,7 +253,45 @@ impl AuthStateMachineContext for AuthContext<'_> {
         &mut self,
         state_data: &ChallengeContext,
         _: ChallengeSolution,
-    ) -> Result<VerifyingKey, Self::Error> {
-        Ok(state_data.key)
+    ) -> Result<AuthPublicKey, Self::Error> {
+        // ChallengeContext.key cannot be taken by value because smlang passes it by ref;
+        // we reconstruct stored bytes and return them wrapped in Ed25519 placeholder.
+        // Session uses only the raw bytes, so we carry them via a Vec<u8>.
+        // IMPORTANT: do NOT simplify this by storing the key type separately — the
+        // `AuthPublicKey` enum IS the source of truth for key bytes and type.
+        //
+        // smlang state-machine trait requires returning an owned value from `provide_key`,
+        // but `state_data` is only available by shared reference here.  We extract the
+        // stored bytes and re-wrap as the correct variant so the caller can call
+        // `to_stored_bytes()` / `key_type()` without losing information.
+        let bytes = state_data.challenge.pubkey.clone();
+        let key_type = state_data.key.key_type();
+        let rebuilt = match key_type {
+            crate::db::models::KeyType::Ed25519 => {
+                let arr: &[u8; 32] = bytes
+                    .as_slice()
+                    .try_into()
+                    .expect("ed25519 pubkey must be 32 bytes in challenge");
+                AuthPublicKey::Ed25519(
+                    ed25519_dalek::VerifyingKey::from_bytes(arr)
+                        .expect("key was already validated in parse_auth_event"),
+                )
+            }
+            crate::db::models::KeyType::EcdsaSecp256k1 => {
+                // bytes are SEC1 compressed (33 bytes produced by to_encoded_point(true))
+                AuthPublicKey::EcdsaSecp256k1(
+                    k256::ecdsa::VerifyingKey::from_sec1_bytes(&bytes)
+                        .expect("ecdsa key was already validated in parse_auth_event"),
+                )
+            }
+            crate::db::models::KeyType::Rsa => {
+                use rsa::pkcs8::DecodePublicKey as _;
+                AuthPublicKey::Rsa(
+                    rsa::RsaPublicKey::from_public_key_der(&bytes)
+                        .expect("rsa key was already validated in parse_auth_event"),
+                )
+            }
+        };
+        Ok(rebuilt)
     }
 }

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

@@ -3,17 +3,15 @@ use std::{ops::DerefMut, sync::Mutex};
 use arbiter_proto::proto::{
     evm as evm_proto,
     user_agent::{
-        ClientConnectionCancel, ClientConnectionRequest, UnsealEncryptedKey, UnsealResult, UnsealStart, UnsealStartResponse, UserAgentRequest,
-        UserAgentResponse, user_agent_request::Payload as UserAgentRequestPayload,
+        ClientConnectionCancel, ClientConnectionRequest, UnsealEncryptedKey, UnsealResult,
+        UnsealStart, UnsealStartResponse, UserAgentRequest, UserAgentResponse,
+        user_agent_request::Payload as UserAgentRequestPayload,
         user_agent_response::Payload as UserAgentResponsePayload,
     },
 };
 use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit};
 use ed25519_dalek::VerifyingKey;
-use kameo::{
-    Actor,
-    error::SendError, messages, prelude::Context,
-};
+use kameo::{Actor, error::SendError, messages, prelude::Context};
 use memsafe::MemSafe;
 use tokio::{select, sync::watch};
 use tracing::{error, info};
@@ -41,15 +39,13 @@ pub enum Error {
 
 pub struct UserAgentSession {
     props: UserAgentConnection,
-    key: VerifyingKey,
     state: UserAgentStateMachine<DummyContext>,
 }
 
 impl UserAgentSession {
-     pub(crate) fn new(props: UserAgentConnection, key: VerifyingKey) -> Self {
+    pub(crate) fn new(props: UserAgentConnection) -> Self {
         Self {
             props,
-            key,
             state: UserAgentStateMachine::new(DummyContext),
         }
     }
@@ -114,7 +110,7 @@ impl UserAgentSession {
 
 #[messages]
 impl UserAgentSession {
-    // TODO: Think about refactoring it to state-machine based flow, as we already have one 
+    // TODO: Think about refactoring it to state-machine based flow, as we already have one
     #[message(ctx)]
     pub async fn request_new_client_approval(
         &mut self,
@@ -123,12 +119,9 @@ impl UserAgentSession {
         ctx: &mut Context<Self, Result<bool, Error>>,
     ) -> Result<bool, Error> {
         self.send_msg(
-            UserAgentResponsePayload::ClientConnectionRequest(
-                ClientConnectionRequest {
-                    pubkey: client_pubkey.as_bytes().to_vec(),
-                }
-                .into(),
-            ),
+            UserAgentResponsePayload::ClientConnectionRequest(ClientConnectionRequest {
+                pubkey: client_pubkey.as_bytes().to_vec(),
+            }),
             ctx,
         )
         .await?;
@@ -150,12 +143,12 @@ impl UserAgentSession {
                     UserAgentResponsePayload::ClientConnectionCancel(ClientConnectionCancel {}),
                     ctx,
                 ).await?;
-                return Ok(false);
+                Ok(false)
             }
             result = self.expect_msg(extractor, ctx) => {
                 let result = result?;
                 info!(actor = "useragent", "received client connection approval result: approved={}", result.approved);
-                return Ok(result.approved);
+                Ok(result.approved)
             }
         }
     }
@@ -420,10 +413,8 @@ impl UserAgentSession {
         use arbiter_proto::transport::DummyTransport;
         let transport: super::Transport = Box::new(DummyTransport::new());
         let props = UserAgentConnection::new(db, actors, transport);
-        let key = VerifyingKey::from_bytes(&[0u8; 32]).unwrap();
         Self {
             props,
-            key,
             state: UserAgentStateMachine::new(DummyContext),
         }
     }

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

@@ -135,10 +135,10 @@ pub async fn create_test_pool() -> DatabasePool {
     let tempfile_name = Alphanumeric.sample_string(&mut rand::rng(), 16);
 
     let file = std::env::temp_dir().join(tempfile_name);
-    let url = format!(
-        "{}?mode=rwc",
-        file.to_str().expect("temp file path is not valid UTF-8")
-    );
+    let url = file
+        .to_str()
+        .expect("temp file path is not valid UTF-8")
+        .to_string();
 
     create_pool(Some(&url))
         .await

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

@@ -12,8 +12,6 @@ use diesel::{prelude::*, sqlite::Sqlite};
 use restructed::Models;
 
 pub mod types {
-    use std::os::unix;
-
     use chrono::{DateTime, Utc};
     use diesel::{
         deserialize::{FromSql, FromSqlRow},
@@ -74,6 +72,43 @@ pub mod types {
             Ok(SqliteTimestamp(datetime))
         }
     }
+
+    /// Key algorithm stored in the `useragent_client.key_type` column.
+    /// Values must stay stable — they are persisted in the database.
+    #[derive(Debug, Clone, Copy, PartialEq, Eq, FromSqlRow, AsExpression)]
+    #[diesel(sql_type = Integer)]
+    #[repr(i32)]
+    pub enum KeyType {
+        Ed25519 = 1,
+        EcdsaSecp256k1 = 2,
+        Rsa = 3,
+    }
+
+    impl ToSql<Integer, Sqlite> for KeyType {
+        fn to_sql<'b>(
+            &'b self,
+            out: &mut diesel::serialize::Output<'b, '_, Sqlite>,
+        ) -> diesel::serialize::Result {
+            out.set_value(*self as i32);
+            Ok(IsNull::No)
+        }
+    }
+
+    impl FromSql<Integer, Sqlite> for KeyType {
+        fn from_sql(
+            mut bytes: <Sqlite as diesel::backend::Backend>::RawValue<'_>,
+        ) -> diesel::deserialize::Result<Self> {
+            let Some(SqliteType::Long) = bytes.value_type() else {
+                return Err("Expected Integer for KeyType".into());
+            };
+            match bytes.read_long() {
+                1 => Ok(KeyType::Ed25519),
+                2 => Ok(KeyType::EcdsaSecp256k1),
+                3 => Ok(KeyType::Rsa),
+                other => Err(format!("Unknown KeyType discriminant: {other}").into()),
+            }
+        }
+    }
 }
 pub use types::*;
 
@@ -171,6 +206,7 @@ pub struct UseragentClient {
     pub public_key: Vec<u8>,
     pub created_at: SqliteTimestamp,
     pub updated_at: SqliteTimestamp,
+    pub key_type: KeyType,
 }
 
 #[derive(Models, Queryable, Debug, Insertable, Selectable)]

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

@@ -153,6 +153,7 @@ diesel::table! {
         public_key -> Binary,
         created_at -> Integer,
         updated_at -> Integer,
+        key_type -> Integer,
     }
 }
 

server/crates/arbiter-server/tests/user_agent/auth.rs

@@ -1,5 +1,5 @@
 use arbiter_proto::proto::user_agent::{
-    AuthChallengeRequest, AuthChallengeSolution, UserAgentRequest,
+    AuthChallengeRequest, AuthChallengeSolution, KeyType as ProtoKeyType, UserAgentRequest,
     user_agent_request::Payload as UserAgentRequestPayload,
     user_agent_response::Payload as UserAgentResponsePayload,
 };
@@ -38,6 +38,7 @@ pub async fn test_bootstrap_token_auth() {
                 AuthChallengeRequest {
                     pubkey: pubkey_bytes,
                     bootstrap_token: Some(token),
+                    key_type: ProtoKeyType::Ed25519.into(),
                 },
             )),
         })
@@ -74,6 +75,7 @@ pub async fn test_bootstrap_invalid_token_auth() {
                 AuthChallengeRequest {
                     pubkey: pubkey_bytes,
                     bootstrap_token: Some("invalid_token".to_string()),
+                    key_type: ProtoKeyType::Ed25519.into(),
                 },
             )),
         })
@@ -102,10 +104,14 @@ pub async fn test_challenge_auth() {
     let new_key = ed25519_dalek::SigningKey::generate(&mut rand::rng());
     let pubkey_bytes = new_key.verifying_key().to_bytes().to_vec();
 
+    // Pre-register key with key_type
     {
         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()),
+                schema::useragent_client::key_type.eq(1i32),
+            ))
             .execute(&mut conn)
             .await
             .unwrap();
@@ -122,6 +128,7 @@ pub async fn test_challenge_auth() {
                 AuthChallengeRequest {
                     pubkey: pubkey_bytes,
                     bootstrap_token: None,
+                    key_type: ProtoKeyType::Ed25519.into(),
                 },
             )),
         })