refactor(keyholder): generalize derive_useragent_integrity_key and compute_useragent_pubkey_integrity_tag corespondenly to derive_integrity_key and compute_integrity_tag

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

@@ -172,15 +172,15 @@ pub fn derive_seal_key(mut password: SafeCell<Vec<u8>>, salt: &Salt) -> KeyCell
     key.into()
 }
 
-/// Derives a dedicated key used only for user-agent pubkey integrity tags.
+/// Derives a dedicated key used for integrity tags within a specific domain.
-pub fn derive_useragent_integrity_key(seal_key: &mut KeyCell) -> KeyCell {
+pub fn derive_integrity_key(seal_key: &mut KeyCell, derive_tag: &[u8]) -> KeyCell {
     type HmacSha256 = hmac::Hmac<sha2::Sha256>;
 
     let mut derived = SafeCell::new(Key::default());
     seal_key.0.read_inline(|seal_key_bytes| {
         let mut mac = <HmacSha256 as hmac::Mac>::new_from_slice(seal_key_bytes.as_ref())
             .expect("HMAC key initialization must not fail for 32-byte key");
-        mac.update(USERAGENT_INTEGRITY_DERIVE_TAG);
+        mac.update(derive_tag);
         let output = mac.finalize().into_bytes();
 
         let mut writer = derived.write();
@@ -191,25 +191,29 @@ pub fn derive_useragent_integrity_key(seal_key: &mut KeyCell) -> KeyCell {
     derived.into()
 }
 
-/// Computes an integrity tag for a user-agent pubkey DB entry.
+/// Computes an integrity tag for a specific domain and payload shape.
-pub fn compute_useragent_pubkey_integrity_tag(
+pub fn compute_integrity_tag<'a, I>(
     integrity_key: &mut KeyCell,
-    key_type_discriminant: i32,
+    purpose_tag: &[u8],
-    public_key: &[u8],
+    data_parts: I,
-) -> [u8; 32] {
+) -> [u8; 32]
+where
+    I: IntoIterator<Item = &'a [u8]>,
+{
     type HmacSha256 = hmac::Hmac<sha2::Sha256>;
 
-    let mut tag = [0u8; 32];
+    let mut output_tag = [0u8; 32];
     integrity_key.0.read_inline(|integrity_key_bytes| {
         let mut mac = <HmacSha256 as hmac::Mac>::new_from_slice(integrity_key_bytes.as_ref())
             .expect("HMAC key initialization must not fail for 32-byte key");
-        mac.update(USERAGENT_INTEGRITY_TAG);
+        mac.update(purpose_tag);
-        mac.update(&key_type_discriminant.to_be_bytes());
+        for data_part in data_parts {
-        mac.update(public_key);
+            mac.update(data_part);
-        tag.copy_from_slice(&mac.finalize().into_bytes());
+        }
+        output_tag.copy_from_slice(&mac.finalize().into_bytes());
     });
 
-    tag
+    output_tag
 }
 
 #[cfg(test)]
@@ -285,22 +289,41 @@ mod tests {
     }
 
     #[test]
-    pub fn useragent_integrity_tag_deterministic() {
+    pub fn integrity_tag_deterministic() {
         let salt = generate_salt();
         let mut seal_key = derive_seal_key(SafeCell::new(b"password".to_vec()), &salt);
-        let mut integrity_key = derive_useragent_integrity_key(&mut seal_key);
+        let mut integrity_key = derive_integrity_key(&mut seal_key, USERAGENT_INTEGRITY_DERIVE_TAG);
-        let t1 = compute_useragent_pubkey_integrity_tag(&mut integrity_key, 1, b"pubkey");
+        let key_type = 1i32.to_be_bytes();
-        let t2 = compute_useragent_pubkey_integrity_tag(&mut integrity_key, 1, b"pubkey");
+        let t1 = compute_integrity_tag(
+            &mut integrity_key,
+            USERAGENT_INTEGRITY_TAG,
+            [key_type.as_slice(), b"pubkey".as_ref()],
+        );
+        let t2 = compute_integrity_tag(
+            &mut integrity_key,
+            USERAGENT_INTEGRITY_TAG,
+            [key_type.as_slice(), b"pubkey".as_ref()],
+        );
         assert_eq!(t1, t2);
     }
 
     #[test]
-    pub fn useragent_integrity_tag_changes_with_key_type() {
+    pub fn integrity_tag_changes_with_payload() {
         let salt = generate_salt();
         let mut seal_key = derive_seal_key(SafeCell::new(b"password".to_vec()), &salt);
-        let mut integrity_key = derive_useragent_integrity_key(&mut seal_key);
+        let mut integrity_key = derive_integrity_key(&mut seal_key, USERAGENT_INTEGRITY_DERIVE_TAG);
-        let t1 = compute_useragent_pubkey_integrity_tag(&mut integrity_key, 1, b"pubkey");
+        let key_type_1 = 1i32.to_be_bytes();
-        let t2 = compute_useragent_pubkey_integrity_tag(&mut integrity_key, 2, b"pubkey");
+        let key_type_2 = 2i32.to_be_bytes();
+        let t1 = compute_integrity_tag(
+            &mut integrity_key,
+            USERAGENT_INTEGRITY_TAG,
+            [key_type_1.as_slice(), b"pubkey".as_ref()],
+        );
+        let t2 = compute_integrity_tag(
+            &mut integrity_key,
+            USERAGENT_INTEGRITY_TAG,
+            [key_type_2.as_slice(), b"pubkey".as_ref()],
+        );
         assert_ne!(t1, t2);
     }
 }

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

@@ -32,7 +32,7 @@ enum State {
     Unsealed {
         root_key_history_id: i32,
         root_key: KeyCell,
-        useragent_integrity_key: KeyCell,
+        integrity_key: KeyCell,
     },
 }
 
@@ -146,7 +146,8 @@ impl KeyHolder {
         }
         let salt = v1::generate_salt();
         let mut seal_key = v1::derive_seal_key(seal_key_raw, &salt);
-        let useragent_integrity_key = v1::derive_useragent_integrity_key(&mut seal_key);
+        let integrity_key =
+            v1::derive_integrity_key(&mut seal_key, v1::USERAGENT_INTEGRITY_DERIVE_TAG);
         let mut root_key = KeyCell::new_secure_random();
 
         // Zero nonces are fine because they are one-time
@@ -195,7 +196,7 @@ impl KeyHolder {
         self.state = State::Unsealed {
             root_key,
             root_key_history_id,
-            useragent_integrity_key,
+            integrity_key,
         };
 
         info!("Keyholder bootstrapped successfully");
@@ -229,7 +230,8 @@ impl KeyHolder {
             Error::BrokenDatabase
         })?;
         let mut seal_key = v1::derive_seal_key(seal_key_raw, &salt);
-        let useragent_integrity_key = v1::derive_useragent_integrity_key(&mut seal_key);
+        let integrity_key =
+            v1::derive_integrity_key(&mut seal_key, v1::USERAGENT_INTEGRITY_DERIVE_TAG);
 
         let mut root_key = SafeCell::new(current_key.ciphertext.clone());
 
@@ -253,7 +255,7 @@ impl KeyHolder {
                 error!(?err, "Broken database: invalid encryption key size");
                 Error::BrokenDatabase
             })?,
-            useragent_integrity_key,
+            integrity_key,
         };
 
         info!("Keyholder unsealed successfully");
@@ -263,23 +265,19 @@ impl KeyHolder {
 
     // Decrypts the `aead_encrypted` entry with the given ID and returns the plaintext
     #[message]
-    pub fn sign_useragent_pubkey_integrity_tag(
+    pub fn sign_integrity_tag(
         &mut self,
-        public_key: Vec<u8>,
+        purpose_tag: Vec<u8>,
-        key_type: models::KeyType,
+        data_parts: Vec<Vec<u8>>,
     ) -> Result<Vec<u8>, Error> {
-        let State::Unsealed {
+        let State::Unsealed { integrity_key, .. } = &mut self.state else {
-            useragent_integrity_key,
-            ..
-        } = &mut self.state
-        else {
             return Err(Error::NotBootstrapped);
         };
 
-        let tag = v1::compute_useragent_pubkey_integrity_tag(
+        let tag = v1::compute_integrity_tag(
-            useragent_integrity_key,
+            integrity_key,
-            key_type as i32,
+            &purpose_tag,
-            &public_key,
+            data_parts.iter().map(Vec::as_slice),
         );
         Ok(tag.to_vec())
     }

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

@@ -8,7 +8,7 @@ use super::Error;
 use crate::{
     actors::{
         bootstrap::ConsumeToken,
-        keyholder::{self, SignUseragentPubkeyIntegrityTag},
+        keyholder::{self, SignIntegrityTag},
         user_agent::{AuthPublicKey, UserAgentConnection, auth::Outbound},
     },
     db::schema,
@@ -257,9 +257,12 @@ where
             .conn
             .actors
             .key_holder
-            .ask(SignUseragentPubkeyIntegrityTag {
+            .ask(SignIntegrityTag {
-                public_key: pubkey.to_stored_bytes(),
+                purpose_tag: keyholder::encryption::v1::USERAGENT_INTEGRITY_TAG.to_vec(),
-                key_type: pubkey.key_type(),
+                data_parts: vec![
+                    (pubkey.key_type() as i32).to_be_bytes().to_vec(),
+                    pubkey.to_stored_bytes(),
+                ],
             })
             .await;
 

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

@@ -23,7 +23,7 @@ use crate::{
         evm::{
             Generate, ListWallets, UseragentCreateGrant, UseragentDeleteGrant, UseragentListGrants,
         },
-        keyholder::{self, Bootstrap, SignUseragentPubkeyIntegrityTag, TryUnseal},
+        keyholder::{self, Bootstrap, SignIntegrityTag, TryUnseal},
         user_agent::session::{
             UserAgentSession,
             state::{UnsealContext, UserAgentEvents, UserAgentStates},
@@ -111,14 +111,14 @@ impl UserAgentSession {
                 .props
                 .actors
                 .key_holder
-                .ask(SignUseragentPubkeyIntegrityTag {
+                .ask(SignIntegrityTag {
-                    public_key,
+                    purpose_tag: keyholder::encryption::v1::USERAGENT_INTEGRITY_TAG.to_vec(),
-                    key_type,
+                    data_parts: vec![(key_type as i32).to_be_bytes().to_vec(), public_key],
                 })
                 .await
                 .map_err(|err| {
-                    error!(?err, "Failed to sign user-agent pubkey integrity tag");
+                    error!(?err, "Failed to sign integrity tag");
-                    Error::internal("Failed to sign user-agent pubkey integrity tag")
+                    Error::internal("Failed to sign integrity tag")
                 })?;
             updates.push((id, tag));
         }