encrypt privateaccountkind instead of identifier

nssa/core/src/encryption/mod.rs

@@ -33,6 +33,10 @@ pub enum PrivateAccountKind {
 }
 
 impl PrivateAccountKind {
+    ///   Account(ident):                  0x00 || ident (16 LE) || [0u8; 64]
+    ///   Pda { program_id, seed, ident }: 0x01 || program_id (32 LE) || seed (32) || ident (16 LE)
+    pub const HEADER_LEN: usize = 81;
+
     #[must_use]
     pub fn identifier(&self) -> Identifier {
         match self {
@@ -40,6 +44,50 @@ impl PrivateAccountKind {
             Self::Pda { identifier, .. } => *identifier,
         }
     }
+
+    #[must_use]
+    pub fn to_header_bytes(&self) -> [u8; Self::HEADER_LEN] {
+        let mut bytes = [0u8; Self::HEADER_LEN];
+        match self {
+            Self::Account(identifier) => {
+                bytes[0] = 0x00;
+                bytes[1..17].copy_from_slice(&identifier.to_le_bytes());
+                // bytes[17..81] are zero padding
+            }
+            Self::Pda { program_id, seed, identifier } => {
+                bytes[0] = 0x01;
+                for (i, &word) in program_id.iter().enumerate() {
+                    bytes[1 + i * 4..1 + (i + 1) * 4].copy_from_slice(&word.to_le_bytes());
+                }
+                bytes[33..65].copy_from_slice(seed.as_bytes());
+                bytes[65..81].copy_from_slice(&identifier.to_le_bytes());
+            }
+        }
+        bytes
+    }
+
+    #[cfg(feature = "host")]
+    #[must_use]
+    pub fn from_header_bytes(bytes: &[u8; Self::HEADER_LEN]) -> Option<Self> {
+        match bytes[0] {
+            0x00 => {
+                let identifier = Identifier::from_le_bytes(bytes[1..17].try_into().unwrap());
+                Some(Self::Account(identifier))
+            }
+            0x01 => {
+                let mut program_id = [0u32; 8];
+                for (i, word) in program_id.iter_mut().enumerate() {
+                    *word = u32::from_le_bytes(
+                        bytes[1 + i * 4..1 + (i + 1) * 4].try_into().unwrap(),
+                    );
+                }
+                let seed = PdaSeed::new(bytes[33..65].try_into().unwrap());
+                let identifier = Identifier::from_le_bytes(bytes[65..81].try_into().unwrap());
+                Some(Self::Pda { program_id, seed, identifier })
+            }
+            _ => None,
+        }
+    }
 }
 
 #[derive(Serialize, Deserialize, Clone, Copy)]
@@ -68,13 +116,14 @@ impl EncryptionScheme {
     #[must_use]
     pub fn encrypt(
         account: &Account,
-        identifier: Identifier,
+        kind: &PrivateAccountKind,
         shared_secret: &SharedSecretKey,
         commitment: &Commitment,
         output_index: u32,
     ) -> Ciphertext {
-        // Plaintext: identifier (16 bytes, little-endian) || account bytes
-        let mut buffer = identifier.to_le_bytes().to_vec();
+        // Plaintext: PrivateAccountKind::HEADER_LEN bytes header || account bytes.
+        // Both variants produce the same header length — see PrivateAccountKind::to_header_bytes.
+        let mut buffer = kind.to_header_bytes().to_vec();
         buffer.extend_from_slice(&account.to_bytes());
         Self::symmetric_transform(&mut buffer, shared_secret, commitment, output_index);
         Ciphertext(buffer)
@@ -117,17 +166,19 @@ impl EncryptionScheme {
         shared_secret: &SharedSecretKey,
         commitment: &Commitment,
         output_index: u32,
-    ) -> Option<(Identifier, Account)> {
+    ) -> Option<(PrivateAccountKind, Account)> {
         use std::io::Cursor;
         let mut buffer = ciphertext.0.clone();
         Self::symmetric_transform(&mut buffer, shared_secret, commitment, output_index);
 
-        if buffer.len() < 16 {
+        if buffer.len() < PrivateAccountKind::HEADER_LEN {
             return None;
         }
-        let identifier = Identifier::from_le_bytes(buffer[..16].try_into().unwrap());
+        let header: &[u8; PrivateAccountKind::HEADER_LEN] =
+            buffer[..PrivateAccountKind::HEADER_LEN].try_into().unwrap();
+        let kind = PrivateAccountKind::from_header_bytes(header)?;
 
-        let mut cursor = Cursor::new(&buffer[16..]);
+        let mut cursor = Cursor::new(&buffer[PrivateAccountKind::HEADER_LEN..]);
         Account::from_cursor(&mut cursor)
             .inspect_err(|err| {
                 println!(
@@ -140,6 +191,40 @@ impl EncryptionScheme {
                 );
             })
             .ok()
-            .map(|account| (identifier, account))
+            .map(|account| (kind, account))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::account::{Account, AccountId};
+
+    #[test]
+    fn encrypt_same_length_for_account_and_pda() {
+        let account = Account::default();
+        let secret = SharedSecretKey([0u8; 32]);
+        let commitment = crate::Commitment::new(&AccountId::new([0u8; 32]), &Account::default());
+
+        let account_ct = EncryptionScheme::encrypt(
+            &account,
+            &PrivateAccountKind::Account(42),
+            &secret,
+            &commitment,
+            0,
+        );
+        let pda_ct = EncryptionScheme::encrypt(
+            &account,
+            &PrivateAccountKind::Pda {
+                program_id: [1u32; 8],
+                seed: PdaSeed::new([2u8; 32]),
+                identifier: 42,
+            },
+            &secret,
+            &commitment,
+            0,
+        );
+
+        assert_eq!(account_ct.0.len(), pda_ct.0.len());
     }
 }

nssa/core/src/program.rs

@@ -35,6 +35,11 @@ impl PdaSeed {
     pub const fn new(value: [u8; 32]) -> Self {
         Self(value)
     }
+
+    #[must_use]
+    pub const fn as_bytes(&self) -> &[u8; 32] {
+        &self.0
+    }
 }
 
 impl AccountId {

program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -5,6 +5,7 @@ use std::{
 
 use nssa_core::{
     Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, Identifier,
+    PrivateAccountKind,
     MembershipProof, Nullifier, NullifierPublicKey, NullifierSecretKey,
     PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput, SharedSecretKey,
     account::{Account, AccountId, AccountWithMetadata, Nonce},
@@ -606,7 +607,7 @@ fn compute_circuit_output(
                 // Encrypt and push post state
                 let encrypted_account = EncryptionScheme::encrypt(
                     &post_with_updated_nonce,
-                    *identifier,
+                    &PrivateAccountKind::Account(*identifier),
                     shared_secret,
                     &commitment_post,
                     output_index,
@@ -690,7 +691,7 @@ fn compute_circuit_output(
 
                 let encrypted_account = EncryptionScheme::encrypt(
                     &post_with_updated_nonce,
-                    *identifier,
+                    &PrivateAccountKind::Account(*identifier),
                     shared_secret,
                     &commitment_post,
                     output_index,

wallet/src/lib.rs

@@ -357,7 +357,7 @@ impl WalletCore {
                     let acc_ead = tx.message.encrypted_private_post_states[output_index].clone();
                     let acc_comm = tx.message.new_commitments[output_index].clone();
 
-                    let (identifier, res_acc) = nssa_core::EncryptionScheme::decrypt(
+                    let (kind, res_acc) = nssa_core::EncryptionScheme::decrypt(
                         &acc_ead.ciphertext,
                         secret,
                         &acc_comm,
@@ -370,7 +370,7 @@ impl WalletCore {
                     println!("Received new acc {res_acc:#?}");
 
                     self.storage
-                        .insert_private_account_data(*acc_account_id, identifier, res_acc);
+                        .insert_private_account_data(*acc_account_id, kind.identifier(), res_acc);
                 }
                 AccDecodeData::Skip => {}
             }
@@ -544,7 +544,8 @@ impl WalletCore {
                                 .try_into()
                                 .expect("Ciphertext ID is expected to fit in u32"),
                         )
-                        .map(|(identifier, res_acc)| {
+                        .map(|(kind, res_acc)| {
+                            let identifier = kind.identifier();
                             let account_id = nssa::AccountId::from((
                                 &key_chain.nullifier_public_key,
                                 identifier,