From f32e6bb178d7102b01996f0ade5841d1b273bd3f Mon Sep 17 00:00:00 2001
From: Artem Gureev <artem@heliax.dev>
Date: Wed, 1 Jul 2026 14:41:48 +0000
Subject: [PATCH] refactor(key_protocol): use SharedAccountTag seed in
 GroupKeyHolder

---
 .../src/key_management/group_key_holder.rs        | 15 ++++++++++++++-
 lez/wallet/src/lib.rs                             | 10 +---------
 lez/wallet/src/storage/key_chain.rs               | 12 +-----------
 3 files changed, 16 insertions(+), 21 deletions(-)

diff --git a/lee/key_protocol/src/key_management/group_key_holder.rs b/lee/key_protocol/src/key_management/group_key_holder.rs
index 7fb24713..fbfd2550 100644
--- a/lee/key_protocol/src/key_management/group_key_holder.rs
+++ b/lee/key_protocol/src/key_management/group_key_holder.rs
@@ -1,6 +1,6 @@
 use aes_gcm::{Aes256Gcm, KeyInit as _, aead::Aead as _};
 use lee_core::{
-    SharedSecretKey,
+    Identifier, SharedSecretKey,
     encryption::{EphemeralPublicKey, ViewingPublicKey},
     program::{PdaSeed, ProgramId},
 };
@@ -146,6 +146,19 @@ impl GroupKeyHolder {
         SecretSpendingKey(hasher.finalize_fixed().into()).produce_private_key_holder(None)
     }
 
+    #[must_use]
+    pub fn derive_keys_for_regular_shared_account(
+        &self,
+        identifier: Identifier,
+    ) -> PrivateKeyHolder {
+        const PREFIX: &[u8; 32] = b"/LEE/v0.3/SharedAccountTag/\x00\x00\x00\x00\x00";
+        let mut hasher = sha2::Sha256::new();
+        hasher.update(PREFIX);
+        hasher.update(identifier.to_le_bytes());
+        let derivation_seed: [u8; 32] = hasher.finalize().into();
+        self.derive_keys_for_shared_account(&derivation_seed)
+    }
+
     /// Encrypts this holder's GMS under the recipient's [`SealingPublicKey`].
     ///
     /// Uses ML-KEM-768 encapsulation to derive a shared secret, then AES-256-GCM to encrypt
diff --git a/lez/wallet/src/lib.rs b/lez/wallet/src/lib.rs
index acea846d..c4250afd 100644
--- a/lez/wallet/src/lib.rs
+++ b/lez/wallet/src/lib.rs
@@ -395,14 +395,6 @@ impl WalletCore {
         group_name: Label,
     ) -> Result<SharedAccountInfo> {
         let identifier: lee_core::Identifier = rand::random();
-        let derivation_seed = {
-            use sha2::Digest as _;
-            let mut hasher = sha2::Sha256::new();
-            hasher.update(b"/LEE/v0.3/SharedAccountTag/\x00\x00\x00\x00\x00");
-            hasher.update(identifier.to_le_bytes());
-            let result: [u8; 32] = hasher.finalize().into();
-            result
-        };
 
         let holder = self
             .storage
@@ -410,7 +402,7 @@ impl WalletCore {
             .group_key_holder(&group_name)
             .context(format!("Group '{group_name}' not found"))?;
 
-        let keys = holder.derive_keys_for_shared_account(&derivation_seed);
+        let keys = holder.derive_keys_for_regular_shared_account(identifier);
         let npk = keys.generate_nullifier_public_key();
         let vpk = keys.generate_viewing_public_key();
         let account_id = AccountId::from((&npk, identifier));
diff --git a/lez/wallet/src/storage/key_chain.rs b/lez/wallet/src/storage/key_chain.rs
index 1240b0ba..1730d855 100644
--- a/lez/wallet/src/storage/key_chain.rs
+++ b/lez/wallet/src/storage/key_chain.rs
@@ -293,17 +293,7 @@ impl UserKeyChain {
         Some(match (&entry.pda_seed, &entry.authority_program_id) {
             (Some(pda_seed), Some(program_id)) => holder.derive_keys_for_pda(program_id, pda_seed),
             (Some(_), None) => return None,
-            _ => {
-                let derivation_seed = {
-                    use sha2::Digest as _;
-                    let mut hasher = sha2::Sha256::new();
-                    hasher.update(b"/LEE/v0.3/SharedAccountTag/\x00\x00\x00\x00\x00");
-                    hasher.update(entry.identifier.to_le_bytes());
-                    let result: [u8; 32] = hasher.finalize().into();
-                    result
-                };
-                holder.derive_keys_for_shared_account(&derivation_seed)
-            }
+            _ => holder.derive_keys_for_regular_shared_account(entry.identifier),
         })
     }