Merge branch 'schouhy/implement-nssa-v0.1-private-state' into Pravdyvy/sequencer-update

ci_scripts/test-ubuntu.sh

@@ -4,7 +4,7 @@ curl -L https://risczero.com/install | bash
 /home/runner/.risc0/bin/rzup install 
 source env.sh
 
-cargo test --release
+RISC0_DEV_MODE=1 cargo test --release
 cd integration_tests
 export NSSA_WALLET_HOME_DIR=$(pwd)/configs/debug/wallet/
-cargo run $(pwd)/configs/debug all
+cargo run $(pwd)/configs/debug all

nssa/Cargo.toml

@@ -5,8 +5,8 @@ edition = "2024"
 
 [dependencies]
 thiserror = "2.0.12"
-risc0-zkvm = "3.0.1"
-nssa-core = { path = "core", features=["host"]}
+risc0-zkvm = "3.0.3"
+nssa-core = { path = "core", features = ["host"] }
 program-methods = { path = "program_methods" }
 serde = "1.0.219"
 sha2 = "0.10.9"

nssa/core/Cargo.toml

@@ -4,13 +4,13 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
-risc0-zkvm = "3.0.1"
+risc0-zkvm = { version = "3.0.3" }
 serde = { version = "1.0", default-features = false }
 thiserror = { version = "2.0.12", optional = true }
 bytemuck = { version = "1.13", optional = true }
 chacha20 = { version = "0.9", default-features = false }
-k256 = "0.13.3"
+k256 = { version = "0.13.3", optional = true }
 
 [features]
 default = []
-host = ["thiserror", "bytemuck"]
+host = ["thiserror", "bytemuck", "k256"]

nssa/core/src/account.rs

@@ -1,16 +1,7 @@
+use crate::program::ProgramId;
 use serde::{Deserialize, Serialize};
 
-use crate::program::ProgramId;
-
-mod commitment;
-mod encoding;
-mod nullifier;
-
-pub use commitment::Commitment;
-pub use nullifier::{Nullifier, NullifierPublicKey, NullifierSecretKey};
-
 pub type Nonce = u128;
-// TODO: Consider changing `Data` to `Vec<u32>` for r0 friendlinenss
 type Data = Vec<u8>;
 
 /// Account to be used both in public and private contexts

nssa/core/src/account/commitment.rs

@@ -1,34 +0,0 @@
-use risc0_zkvm::{
-    serde::to_vec,
-    sha::{Impl, Sha256},
-};
-use serde::{Deserialize, Serialize};
-
-use crate::account::{Account, NullifierPublicKey};
-
-#[derive(Serialize, Deserialize)]
-#[cfg_attr(any(feature = "host", test), derive(Debug, Clone, PartialEq, Eq, Hash))]
-pub struct Commitment(pub(super) [u8; 32]);
-
-impl Commitment {
-    pub fn new(Npk: &NullifierPublicKey, account: &Account) -> Self {
-        let mut bytes = Vec::new();
-        bytes.extend_from_slice(&Npk.to_byte_array());
-        let account_bytes_with_hashed_data = {
-            let mut this = Vec::new();
-            for word in &account.program_owner {
-                this.extend_from_slice(&word.to_le_bytes());
-            }
-            this.extend_from_slice(&account.balance.to_le_bytes());
-            this.extend_from_slice(&account.nonce.to_le_bytes());
-            let hashed_data: [u8; 32] = Impl::hash_bytes(&account.data)
-                .as_bytes()
-                .try_into()
-                .unwrap();
-            this.extend_from_slice(&hashed_data);
-            this
-        };
-        bytes.extend_from_slice(&account_bytes_with_hashed_data);
-        Self(Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap())
-    }
-}

nssa/core/src/circuit_io.rs

@@ -0,0 +1,94 @@
+use serde::{Deserialize, Serialize};
+
+use crate::{
+    Commitment, CommitmentSetDigest, MembershipProof, Nullifier, NullifierPublicKey,
+    NullifierSecretKey, SharedSecretKey,
+    account::{Account, AccountWithMetadata, Nonce},
+    encryption::Ciphertext,
+    program::{ProgramId, ProgramOutput},
+};
+
+#[derive(Serialize, Deserialize)]
+pub struct PrivacyPreservingCircuitInput {
+    pub program_output: ProgramOutput,
+    pub visibility_mask: Vec<u8>,
+    pub private_account_nonces: Vec<Nonce>,
+    pub private_account_keys: Vec<(NullifierPublicKey, SharedSecretKey)>,
+    pub private_account_auth: Vec<(NullifierSecretKey, MembershipProof)>,
+    pub program_id: ProgramId,
+}
+
+#[derive(Serialize, Deserialize)]
+#[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))]
+pub struct PrivacyPreservingCircuitOutput {
+    pub public_pre_states: Vec<AccountWithMetadata>,
+    pub public_post_states: Vec<Account>,
+    pub ciphertexts: Vec<Ciphertext>,
+    pub new_commitments: Vec<Commitment>,
+    pub new_nullifiers: Vec<(Nullifier, CommitmentSetDigest)>,
+}
+
+#[cfg(feature = "host")]
+impl PrivacyPreservingCircuitOutput {
+    pub fn to_bytes(&self) -> Vec<u8> {
+        bytemuck::cast_slice(&risc0_zkvm::serde::to_vec(&self).unwrap()).to_vec()
+    }
+}
+
+#[cfg(feature = "host")]
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        Commitment, Nullifier, NullifierPublicKey,
+        account::{Account, AccountWithMetadata},
+    };
+    use risc0_zkvm::serde::from_slice;
+
+    #[test]
+    fn test_privacy_preserving_circuit_output_to_bytes_is_compatible_with_from_slice() {
+        let output = PrivacyPreservingCircuitOutput {
+            public_pre_states: vec![
+                AccountWithMetadata {
+                    account: Account {
+                        program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
+                        balance: 12345678901234567890,
+                        data: b"test data".to_vec(),
+                        nonce: 18446744073709551614,
+                    },
+                    is_authorized: true,
+                },
+                AccountWithMetadata {
+                    account: Account {
+                        program_owner: [9, 9, 9, 8, 8, 8, 7, 7],
+                        balance: 123123123456456567112,
+                        data: b"test data".to_vec(),
+                        nonce: 9999999999999999999999,
+                    },
+                    is_authorized: false,
+                },
+            ],
+            public_post_states: vec![Account {
+                program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
+                balance: 100,
+                data: b"post state data".to_vec(),
+                nonce: 18446744073709551615,
+            }],
+            ciphertexts: vec![Ciphertext(vec![255, 255, 1, 1, 2, 2])],
+            new_commitments: vec![Commitment::new(
+                &NullifierPublicKey::from(&[1; 32]),
+                &Account::default(),
+            )],
+            new_nullifiers: vec![(
+                Nullifier::new(
+                    &Commitment::new(&NullifierPublicKey::from(&[2; 32]), &Account::default()),
+                    &[1; 32],
+                ),
+                [0xab; 32],
+            )],
+        };
+        let bytes = output.to_bytes();
+        let output_from_slice: PrivacyPreservingCircuitOutput = from_slice(&bytes).unwrap();
+        assert_eq!(output, output_from_slice);
+    }
+}

nssa/core/src/commitment.rs

@@ -0,0 +1,63 @@
+use risc0_zkvm::sha::{Impl, Sha256};
+use serde::{Deserialize, Serialize};
+
+use crate::{NullifierPublicKey, account::Account};
+
+#[derive(Serialize, Deserialize)]
+#[cfg_attr(any(feature = "host", test), derive(Debug, Clone, PartialEq, Eq, Hash))]
+pub struct Commitment(pub(super) [u8; 32]);
+
+impl Commitment {
+    pub fn new(npk: &NullifierPublicKey, account: &Account) -> Self {
+        let mut bytes = Vec::new();
+        bytes.extend_from_slice(&npk.to_byte_array());
+        let account_bytes_with_hashed_data = {
+            let mut this = Vec::new();
+            for word in &account.program_owner {
+                this.extend_from_slice(&word.to_le_bytes());
+            }
+            this.extend_from_slice(&account.balance.to_le_bytes());
+            this.extend_from_slice(&account.nonce.to_le_bytes());
+            let hashed_data: [u8; 32] = Impl::hash_bytes(&account.data)
+                .as_bytes()
+                .try_into()
+                .unwrap();
+            this.extend_from_slice(&hashed_data);
+            this
+        };
+        bytes.extend_from_slice(&account_bytes_with_hashed_data);
+        Self(Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap())
+    }
+}
+
+pub type CommitmentSetDigest = [u8; 32];
+
+pub type MembershipProof = (usize, Vec<[u8; 32]>);
+
+pub fn compute_digest_for_path(
+    commitment: &Commitment,
+    proof: &MembershipProof,
+) -> CommitmentSetDigest {
+    let value_bytes = commitment.to_byte_array();
+    let mut result: [u8; 32] = Impl::hash_bytes(&value_bytes)
+        .as_bytes()
+        .try_into()
+        .unwrap();
+    let mut level_index = proof.0;
+    for node in &proof.1 {
+        let is_left_child = level_index & 1 == 0;
+        if is_left_child {
+            let mut bytes = [0u8; 64];
+            bytes[..32].copy_from_slice(&result);
+            bytes[32..].copy_from_slice(node);
+            result = Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap();
+        } else {
+            let mut bytes = [0u8; 64];
+            bytes[..32].copy_from_slice(node);
+            bytes[32..].copy_from_slice(&result);
+            result = Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap();
+        }
+        level_index >>= 1;
+    }
+    result
+}

nssa/core/src/encoding.rs

@@ -1,17 +1,25 @@
 // TODO: Consider switching to deriving Borsh
-
-use risc0_zkvm::sha::{Impl, Sha256};
-
 #[cfg(feature = "host")]
 use std::io::Cursor;
 
 #[cfg(feature = "host")]
 use std::io::Read;
 
-use crate::account::{Account, Commitment, Nullifier, NullifierPublicKey};
+use crate::account::Account;
+
+#[cfg(feature = "host")]
+use crate::encryption::shared_key_derivation::Secp256k1Point;
+
+use crate::encryption::Ciphertext;
+
 #[cfg(feature = "host")]
 use crate::error::NssaCoreError;
 
+use crate::Commitment;
+#[cfg(feature = "host")]
+use crate::Nullifier;
+use crate::NullifierPublicKey;
+
 impl Account {
     pub fn to_bytes(&self) -> Vec<u8> {
         let mut bytes = Vec::new();
@@ -93,6 +101,42 @@ impl Nullifier {
     }
 }
 
+impl Ciphertext {
+    pub fn to_bytes(&self) -> Vec<u8> {
+        let mut bytes = Vec::new();
+        let ciphertext_length: u32 = self.0.len() as u32;
+        bytes.extend_from_slice(&ciphertext_length.to_le_bytes());
+        bytes.extend_from_slice(&self.0);
+
+        bytes
+    }
+
+    #[cfg(feature = "host")]
+    pub fn from_cursor(cursor: &mut Cursor<&[u8]>) -> Result<Self, NssaCoreError> {
+        let mut u32_bytes = [0; 4];
+
+        cursor.read_exact(&mut u32_bytes)?;
+        let ciphertext_lenght = u32::from_le_bytes(u32_bytes);
+        let mut ciphertext = vec![0; ciphertext_lenght as usize];
+        cursor.read_exact(&mut ciphertext)?;
+
+        Ok(Self(ciphertext))
+    }
+}
+
+#[cfg(feature = "host")]
+impl Secp256k1Point {
+    pub fn to_bytes(&self) -> [u8; 33] {
+        self.0.clone().try_into().unwrap()
+    }
+
+    pub fn from_cursor(cursor: &mut Cursor<&[u8]>) -> Result<Self, NssaCoreError> {
+        let mut value = vec![0; 33];
+        cursor.read_exact(&mut value)?;
+        Ok(Self(value))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -106,7 +150,7 @@ mod tests {
             data: b"hola mundo".to_vec(),
         };
 
-        // program owner || balance || nonce || hash(data)
+        // program owner || balance || nonce || data_len || data
         let expected_bytes = [
             1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 4, 0, 0, 0, 5, 0, 0, 0, 6, 0, 0, 0, 7, 0, 0, 0, 8,
             0, 0, 0, 192, 186, 220, 114, 113, 65, 236, 234, 222, 15, 215, 191, 227, 198, 23, 0, 42,
@@ -127,6 +171,7 @@ mod tests {
         assert_eq!(expected_bytes, bytes);
     }
 
+    #[cfg(feature = "host")]
     #[test]
     fn test_nullifier_to_bytes() {
         let nullifier = Nullifier((0..32).collect::<Vec<u8>>().try_into().unwrap());

nssa/core/src/encryption/mod.rs

@@ -0,0 +1,77 @@
+use chacha20::{
+    ChaCha20,
+    cipher::{KeyIvInit, StreamCipher},
+};
+use risc0_zkvm::sha::{Impl, Sha256};
+use serde::{Deserialize, Serialize};
+
+#[cfg(feature = "host")]
+pub(crate) mod shared_key_derivation;
+
+#[cfg(feature = "host")]
+pub use shared_key_derivation::{EphemeralPublicKey, EphemeralSecretKey, IncomingViewingPublicKey};
+
+use crate::{Commitment, account::Account};
+
+#[derive(Serialize, Deserialize, Clone)]
+pub struct SharedSecretKey([u8; 32]);
+
+pub struct EncryptionScheme;
+
+#[derive(Serialize, Deserialize)]
+#[cfg_attr(any(feature = "host", test), derive(Debug, Clone, PartialEq, Eq))]
+pub struct Ciphertext(pub(crate) Vec<u8>);
+
+impl EncryptionScheme {
+    pub fn encrypt(
+        account: &Account,
+        shared_secret: &SharedSecretKey,
+        commitment: &Commitment,
+        output_index: u32,
+    ) -> Ciphertext {
+        let mut buffer = account.to_bytes().to_vec();
+        Self::symmetric_transform(&mut buffer, shared_secret, commitment, output_index);
+        Ciphertext(buffer)
+    }
+
+    fn symmetric_transform(
+        buffer: &mut [u8],
+        shared_secret: &SharedSecretKey,
+        commitment: &Commitment,
+        output_index: u32,
+    ) {
+        let key = Self::kdf(shared_secret, commitment, output_index);
+        let mut cipher = ChaCha20::new(&key.into(), &[0; 12].into());
+        cipher.apply_keystream(buffer);
+    }
+
+    fn kdf(
+        shared_secret: &SharedSecretKey,
+        commitment: &Commitment,
+        output_index: u32,
+    ) -> [u8; 32] {
+        let mut bytes = Vec::new();
+
+        bytes.extend_from_slice(b"NSSA/v0.1/KDF-SHA256");
+        bytes.extend_from_slice(&shared_secret.0);
+        bytes.extend_from_slice(&commitment.to_byte_array());
+        bytes.extend_from_slice(&output_index.to_le_bytes());
+
+        Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap()
+    }
+
+    #[cfg(feature = "host")]
+    pub fn decrypt(
+        ciphertext: &Ciphertext,
+        shared_secret: &SharedSecretKey,
+        commitment: &Commitment,
+        output_index: u32,
+    ) -> Option<Account> {
+        use std::io::Cursor;
+        let mut buffer = ciphertext.0.to_owned();
+        Self::symmetric_transform(&mut buffer, shared_secret, commitment, output_index);
+
+        let mut cursor = Cursor::new(buffer.as_slice());
+        Account::from_cursor(&mut cursor).ok()
+    }
+}

nssa/core/src/encryption/shared_key_derivation.rs

@@ -0,0 +1,56 @@
+use serde::{Deserialize, Serialize};
+
+use k256::{
+    AffinePoint, EncodedPoint, FieldBytes, ProjectivePoint, Scalar,
+    elliptic_curve::{
+        PrimeField,
+        sec1::{FromEncodedPoint, ToEncodedPoint},
+    },
+};
+
+use crate::SharedSecretKey;
+
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
+pub struct Secp256k1Point(pub(crate) Vec<u8>);
+
+impl Secp256k1Point {
+    pub fn from_scalar(value: [u8; 32]) -> Secp256k1Point {
+        let x_bytes: FieldBytes = value.into();
+        let x = Scalar::from_repr(x_bytes).unwrap();
+
+        let p = ProjectivePoint::GENERATOR * x;
+        let q = AffinePoint::from(p);
+        let enc = q.to_encoded_point(true);
+
+        Self(enc.as_bytes().to_vec())
+    }
+}
+
+pub type EphemeralSecretKey = [u8; 32];
+pub type EphemeralPublicKey = Secp256k1Point;
+pub type IncomingViewingPublicKey = Secp256k1Point;
+impl From<&EphemeralSecretKey> for EphemeralPublicKey {
+    fn from(value: &EphemeralSecretKey) -> Self {
+        Secp256k1Point::from_scalar(*value)
+    }
+}
+
+impl SharedSecretKey {
+    pub fn new(scalar: &[u8; 32], point: &Secp256k1Point) -> Self {
+        let scalar = Scalar::from_repr((*scalar).into()).unwrap();
+        let point: [u8; 33] = point.0.clone().try_into().unwrap();
+
+        let encoded = EncodedPoint::from_bytes(point).unwrap();
+        let pubkey_affine = AffinePoint::from_encoded_point(&encoded).unwrap();
+
+        let shared = ProjectivePoint::from(pubkey_affine) * scalar;
+        let shared_affine = shared.to_affine();
+
+        let encoded = shared_affine.to_encoded_point(false);
+        let x_bytes_slice = encoded.x().unwrap();
+        let mut x_bytes = [0u8; 32];
+        x_bytes.copy_from_slice(x_bytes_slice);
+
+        Self(x_bytes)
+    }
+}

nssa/core/src/lib.rs

@@ -1,338 +1,15 @@
-use chacha20::{
-    ChaCha20,
-    cipher::{KeyIvInit, StreamCipher},
-};
-use risc0_zkvm::{
-    serde::to_vec,
-    sha::{Impl, Sha256},
-};
-use serde::{Deserialize, Serialize};
-
-#[cfg(feature = "host")]
-use crate::error::NssaCoreError;
-
-use crate::{
-    account::{
-        Account, AccountWithMetadata, Commitment, Nonce, Nullifier, NullifierPublicKey,
-        NullifierSecretKey,
-    },
-    program::{ProgramId, ProgramOutput},
-};
-
-#[cfg(feature = "host")]
-use std::io::{Cursor, Read};
-
 pub mod account;
+mod circuit_io;
+mod commitment;
+mod encoding;
+pub mod encryption;
+mod nullifier;
 pub mod program;
 
-use k256::{
-    AffinePoint, EncodedPoint, FieldBytes, ProjectivePoint, PublicKey, Scalar,
-    elliptic_curve::{
-        PrimeField,
-        sec1::{FromEncodedPoint, ToEncodedPoint},
-    },
-};
+pub use circuit_io::{PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput};
+pub use commitment::{Commitment, CommitmentSetDigest, MembershipProof, compute_digest_for_path};
+pub use encryption::{EncryptionScheme, SharedSecretKey};
+pub use nullifier::{Nullifier, NullifierPublicKey, NullifierSecretKey};
 
 #[cfg(feature = "host")]
 pub mod error;
-
-pub type CommitmentSetDigest = [u8; 32];
-pub type MembershipProof = (usize, Vec<[u8; 32]>);
-pub fn compute_root_associated_to_path(
-    commitment: &Commitment,
-    proof: &MembershipProof,
-) -> CommitmentSetDigest {
-    let value_bytes = commitment.to_byte_array();
-    let mut result: [u8; 32] = Impl::hash_bytes(&value_bytes)
-        .as_bytes()
-        .try_into()
-        .unwrap();
-    let mut level_index = proof.0;
-    for node in &proof.1 {
-        let is_left_child = level_index & 1 == 0;
-        if is_left_child {
-            let mut bytes = [0u8; 64];
-            bytes[..32].copy_from_slice(&result);
-            bytes[32..].copy_from_slice(node);
-            result = Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap();
-        } else {
-            let mut bytes = [0u8; 64];
-            bytes[..32].copy_from_slice(node);
-            bytes[32..].copy_from_slice(&result);
-            result = Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap();
-        }
-        level_index >>= 1;
-    }
-    result
-}
-
-pub type EphemeralPublicKey = Secp256k1Point;
-pub type IncomingViewingPublicKey = Secp256k1Point;
-
-pub type EphemeralSecretKey = [u8; 32];
-
-impl From<&EphemeralSecretKey> for EphemeralPublicKey {
-    fn from(value: &EphemeralSecretKey) -> Self {
-        Secp256k1Point::from_scalar(*value)
-    }
-}
-
-#[derive(Serialize, Deserialize, Clone)]
-#[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))]
-pub struct Secp256k1Point(pub Vec<u8>);
-impl Secp256k1Point {
-    pub fn from_scalar(value: [u8; 32]) -> Secp256k1Point {
-        let x_bytes: FieldBytes = value.into();
-        let x = Scalar::from_repr(x_bytes).unwrap();
-
-        let p = ProjectivePoint::GENERATOR * x;
-        let q = AffinePoint::from(p);
-        let enc = q.to_encoded_point(true);
-
-        Self(enc.as_bytes().to_vec())
-    }
-}
-
-#[derive(Serialize, Deserialize)]
-#[cfg_attr(any(feature = "host", test), derive(Debug, Clone, PartialEq, Eq))]
-pub struct EncryptedAccountData {
-    ciphertext: Vec<u8>,
-    epk: EphemeralPublicKey,
-    view_tag: u8,
-}
-
-impl EncryptedAccountData {
-    #[cfg(feature = "host")]
-    pub fn decrypt(self, isk: &[u8; 32], output_index: u32) -> Option<Account> {
-        let ss_bytes = Self::ecdh(isk, &self.epk.0.clone().try_into().unwrap());
-        let ipk = IncomingViewingPublicKey::from_scalar(*isk);
-
-        let key = Self::kdf(
-            ss_bytes,
-            &self.epk,
-            &ipk,
-            // &commitment.to_byte_array(),
-            output_index,
-        );
-        let mut cipher = ChaCha20::new(&key.into(), &[0; 12].into());
-        let mut buffer = self.ciphertext;
-
-        cipher.apply_keystream(&mut buffer);
-        let mut cursor = Cursor::new(buffer.as_slice());
-        Account::from_cursor(&mut cursor).ok()
-    }
-
-    pub fn new(
-        account: &Account,
-        // commitment: &Commitment,
-        esk: &EphemeralSecretKey,
-        npk: &NullifierPublicKey,
-        ipk: &IncomingViewingPublicKey,
-        output_index: u32,
-    ) -> Self {
-        let mut buffer = account.to_bytes().to_vec();
-
-        let ss_bytes = Self::ecdh(esk, &ipk.0.clone().try_into().unwrap());
-        let epk = EphemeralPublicKey::from(esk);
-
-        let key = Self::kdf(
-            ss_bytes,
-            &epk,
-            ipk,
-            // &commitment.to_byte_array(),
-            output_index,
-        );
-        let mut cipher = ChaCha20::new(&key.into(), &[0; 12].into());
-        cipher.apply_keystream(&mut buffer);
-
-        let view_tag = Self::view_tag(&npk, &ipk);
-        Self {
-            ciphertext: buffer,
-            epk,
-            view_tag,
-        }
-    }
-
-    pub fn kdf(
-        ss_bytes: [u8; 32],
-        epk: &EphemeralPublicKey,
-        ipk: &IncomingViewingPublicKey,
-        // commitment: &[u8; 32],
-        output_index: u32,
-    ) -> [u8; 32] {
-        let mut bytes = Vec::new();
-
-        bytes.extend_from_slice(b"NSSA/v0.1/KDF-SHA256");
-        bytes.extend_from_slice(&ss_bytes);
-        bytes.extend_from_slice(&epk.0[..]);
-        bytes.extend_from_slice(&ipk.0[..]);
-        // bytes.extend_from_slice(&commitment[..]);
-        bytes.extend_from_slice(&output_index.to_le_bytes());
-
-        Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap()
-    }
-
-    pub fn ecdh(scalar: &[u8; 32], point: &[u8; 33]) -> [u8; 32] {
-        let scalar = Scalar::from_repr((*scalar).into()).unwrap();
-
-        let encoded = EncodedPoint::from_bytes(point).unwrap();
-        let pubkey_affine = AffinePoint::from_encoded_point(&encoded).unwrap();
-
-        let shared = ProjectivePoint::from(pubkey_affine) * scalar;
-        let shared_affine = shared.to_affine();
-
-        let encoded = shared_affine.to_encoded_point(false);
-        let x_bytes_slice = encoded.x().unwrap();
-        let mut x_bytes = [0u8; 32];
-        x_bytes.copy_from_slice(x_bytes_slice);
-
-        x_bytes
-    }
-
-    #[cfg(feature = "host")]
-    pub fn from_cursor(cursor: &mut Cursor<&[u8]>) -> Result<Self, NssaCoreError> {
-        let mut u32_bytes = [0; 4];
-
-        cursor.read_exact(&mut u32_bytes)?;
-        let ciphertext_lenght = u32::from_le_bytes(u32_bytes);
-        let mut ciphertext = vec![0; ciphertext_lenght as usize];
-        cursor.read_exact(&mut ciphertext)?;
-
-        let mut epk_bytes = vec![0; 33];
-        cursor.read_exact(&mut epk_bytes)?;
-
-        let mut tag_bytes = [0; 1];
-        cursor.read_exact(&mut tag_bytes)?;
-
-        Ok(Self {
-            ciphertext,
-            epk: Secp256k1Point(epk_bytes),
-            view_tag: tag_bytes[0],
-        })
-    }
-
-    fn view_tag(npk: &NullifierPublicKey, ipk: &&IncomingViewingPublicKey) -> u8 {
-        // TODO: implement
-        0
-    }
-}
-
-impl EncryptedAccountData {
-    pub fn to_bytes(&self) -> Vec<u8> {
-        let mut bytes = Vec::new();
-        let ciphertext_length: u32 = self.ciphertext.len() as u32;
-        bytes.extend_from_slice(&ciphertext_length.to_le_bytes());
-        bytes.extend_from_slice(&self.ciphertext);
-        bytes.extend_from_slice(&self.epk.0);
-        bytes.push(self.view_tag);
-
-        bytes
-    }
-}
-
-#[derive(Serialize, Deserialize)]
-pub struct PrivacyPreservingCircuitInput {
-    pub program_output: ProgramOutput,
-    pub visibility_mask: Vec<u8>,
-    pub private_account_nonces: Vec<Nonce>,
-    pub private_account_keys: Vec<(
-        NullifierPublicKey,
-        IncomingViewingPublicKey,
-        EphemeralSecretKey,
-    )>,
-    pub private_account_auth: Vec<(NullifierSecretKey, MembershipProof)>,
-    pub program_id: ProgramId,
-}
-
-#[derive(Serialize, Deserialize)]
-#[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))]
-pub struct PrivacyPreservingCircuitOutput {
-    pub public_pre_states: Vec<AccountWithMetadata>,
-    pub public_post_states: Vec<Account>,
-    pub encrypted_private_post_states: Vec<EncryptedAccountData>,
-    pub new_commitments: Vec<Commitment>,
-    pub new_nullifiers: Vec<(Nullifier, CommitmentSetDigest)>,
-}
-
-#[cfg(feature = "host")]
-impl PrivacyPreservingCircuitOutput {
-    pub fn to_bytes(&self) -> Vec<u8> {
-        bytemuck::cast_slice(&to_vec(&self).unwrap()).to_vec()
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use std::io::Cursor;
-
-    use risc0_zkvm::serde::from_slice;
-
-    use crate::{
-        EncryptedAccountData, EphemeralPublicKey, PrivacyPreservingCircuitOutput, Secp256k1Point,
-        account::{Account, AccountWithMetadata, Commitment, Nullifier, NullifierPublicKey},
-    };
-
-    #[test]
-    fn test_privacy_preserving_circuit_output_to_bytes_is_compatible_with_from_slice() {
-        let output = PrivacyPreservingCircuitOutput {
-            public_pre_states: vec![
-                AccountWithMetadata {
-                    account: Account {
-                        program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
-                        balance: 12345678901234567890,
-                        data: b"test data".to_vec(),
-                        nonce: 18446744073709551614,
-                    },
-                    is_authorized: true,
-                },
-                AccountWithMetadata {
-                    account: Account {
-                        program_owner: [9, 9, 9, 8, 8, 8, 7, 7],
-                        balance: 123123123456456567112,
-                        data: b"test data".to_vec(),
-                        nonce: 9999999999999999999999,
-                    },
-                    is_authorized: false,
-                },
-            ],
-            public_post_states: vec![Account {
-                program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
-                balance: 100,
-                data: b"post state data".to_vec(),
-                nonce: 18446744073709551615,
-            }],
-            encrypted_private_post_states: vec![EncryptedAccountData {
-                ciphertext: vec![255, 255, 1, 1, 2, 2],
-                epk: EphemeralPublicKey::from_scalar([123; 32]),
-                view_tag: 1,
-            }],
-            new_commitments: vec![Commitment::new(
-                &NullifierPublicKey::from(&[1; 32]),
-                &Account::default(),
-            )],
-            new_nullifiers: vec![(
-                Nullifier::new(
-                    &Commitment::new(&NullifierPublicKey::from(&[2; 32]), &Account::default()),
-                    &[1; 32],
-                ),
-                [0xab; 32],
-            )],
-        };
-        let bytes = output.to_bytes();
-        let output_from_slice: PrivacyPreservingCircuitOutput = from_slice(&bytes).unwrap();
-        assert_eq!(output, output_from_slice);
-    }
-
-    #[test]
-    fn test_encrypted_account_data_to_bytes_roundtrip() {
-        let data = EncryptedAccountData {
-            ciphertext: vec![255, 255, 1, 1, 2, 2],
-            epk: EphemeralPublicKey::from_scalar([123; 32]),
-            view_tag: 95,
-        };
-        let bytes = data.to_bytes();
-        let mut cursor = Cursor::new(bytes.as_slice());
-        let data_from_cursor = EncryptedAccountData::from_cursor(&mut cursor).unwrap();
-        assert_eq!(data, data_from_cursor);
-    }
-}

nssa/core/src/nullifier.rs

@@ -1,7 +1,7 @@
 use risc0_zkvm::sha::{Impl, Sha256};
 use serde::{Deserialize, Serialize};
 
-use crate::account::Commitment;
+use crate::Commitment;
 
 #[derive(Serialize, Deserialize, PartialEq, Eq)]
 #[cfg_attr(any(feature = "host", test), derive(Debug, Clone, Hash))]
@@ -58,11 +58,11 @@ mod tests {
             57, 5, 64, 115, 153, 56, 184, 51, 207, 238, 99, 165, 147, 214, 213, 151, 30, 251, 30,
             196, 134, 22, 224, 211, 237, 120, 136, 225, 188, 220, 249, 28,
         ];
-        let expected_Npk = NullifierPublicKey([
+        let expected_npk = NullifierPublicKey([
             202, 120, 42, 189, 194, 218, 78, 244, 31, 6, 108, 169, 29, 61, 22, 221, 69, 138, 197,
             161, 241, 39, 142, 242, 242, 50, 188, 201, 99, 28, 176, 238,
         ]);
-        let Npk = NullifierPublicKey::from(&nsk);
-        assert_eq!(Npk, expected_Npk);
+        let npk = NullifierPublicKey::from(&nsk);
+        assert_eq!(npk, expected_npk);
     }
 }

nssa/core/src/program.rs

@@ -3,9 +3,6 @@ use risc0_zkvm::serde::Deserializer;
 use risc0_zkvm::{DeserializeOwned, guest::env};
 use serde::{Deserialize, Serialize};
 
-#[cfg(feature = "host")]
-use crate::error::NssaCoreError;
-
 pub type ProgramId = [u32; 8];
 pub type InstructionData = Vec<u32>;
 pub const DEFAULT_PROGRAM_ID: ProgramId = [0; 8];

nssa/program_methods/Cargo.toml

@@ -4,8 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [build-dependencies]
-risc0-build = { version = "3.0.1" }
+risc0-build = { version = "3.0.3" }
 
 [package.metadata.risc0]
 methods = ["guest"]
-

nssa/program_methods/guest/Cargo.lock

@@ -312,12 +312,6 @@ dependencies = [
  "windows-targets 0.52.6",
 ]
 
-[[package]]
-name = "base16ct"
-version = "0.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf"
-
 [[package]]
 name = "base64"
 version = "0.22.1"
@@ -532,6 +526,19 @@ dependencies = [
  "cpufeatures",
 ]
 
+[[package]]
+name = "chrono"
+version = "0.4.41"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c469d952047f47f91b68d1cba3f10d63c11d73e4636f24f08daf0278abf01c4d"
+dependencies = [
+ "android-tzdata",
+ "iana-time-zone",
+ "num-traits",
+ "serde",
+ "windows-link",
+]
+
 [[package]]
 name = "cipher"
 version = "0.4.4"
@@ -593,18 +600,6 @@ dependencies = [
  "libc",
 ]
 
-[[package]]
-name = "crypto-bigint"
-version = "0.5.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0dc92fb57ca44df6db8059111ab3af99a63d5d0f8375d9972e319a379c6bab76"
-dependencies = [
- "generic-array",
- "rand_core 0.6.4",
- "subtle",
- "zeroize",
-]
-
 [[package]]
 name = "crypto-common"
 version = "0.1.6"
@@ -806,18 +801,6 @@ name = "dyn-clone"
 version = "1.0.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555"
-name = "ecdsa"
-version = "0.16.9"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca"
-dependencies = [
- "der",
- "digest",
- "elliptic-curve",
- "rfc6979",
- "signature",
- "spki",
-]
 
 [[package]]
 name = "educe"
@@ -843,25 +826,6 @@ version = "0.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4445909572dbd556c457c849c4ca58623d84b27c8fff1e74b0b4227d8b90d17b"
 
-[[package]]
-name = "elliptic-curve"
-version = "0.13.8"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b5e6043086bf7973472e0c7dff2142ea0b680d30e18d9cc40f267efbf222bd47"
-dependencies = [
- "base16ct",
- "crypto-bigint",
- "digest",
- "ff",
- "generic-array",
- "group",
- "pkcs8",
- "rand_core 0.6.4",
- "sec1",
- "subtle",
- "zeroize",
-]
-
 [[package]]
 name = "embedded-io"
 version = "0.4.0"
@@ -925,16 +889,6 @@ version = "2.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
 
-[[package]]
-name = "ff"
-version = "0.13.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c0b50bfb653653f9ca9095b427bed08ab8d75a137839d9ad64eb11810d5b6393"
-dependencies = [
- "rand_core 0.6.4",
- "subtle",
-]
-
 [[package]]
 name = "fnv"
 version = "1.0.7"
@@ -1053,7 +1007,6 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a"
 dependencies = [
  "typenum",
  "version_check",
- "zeroize",
 ]
 
 [[package]]
@@ -1089,17 +1042,6 @@ version = "0.31.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f"
 
-[[package]]
-name = "group"
-version = "0.13.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63"
-dependencies = [
- "ff",
- "rand_core 0.6.4",
- "subtle",
-]
-
 [[package]]
 name = "hashbrown"
 version = "0.12.3"
@@ -1143,15 +1085,6 @@ version = "0.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6fe2267d4ed49bc07b63801559be28c718ea06c4738b7a03c94df7386d2cde46"
 
-[[package]]
-name = "hmac"
-version = "0.12.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
-dependencies = [
- "digest",
-]
-
 [[package]]
 name = "http"
 version = "1.3.1"
@@ -1478,20 +1411,6 @@ dependencies = [
  "wasm-bindgen",
 ]
 
-[[package]]
-name = "k256"
-version = "0.13.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f6e3919bbaa2945715f0bb6d3934a173d1e9a59ac23767fbaaef277265a7411b"
-dependencies = [
- "cfg-if",
- "ecdsa",
- "elliptic-curve",
- "once_cell",
- "sha2",
- "signature",
-]
-
 [[package]]
 name = "keccak"
 version = "0.1.5"
@@ -1663,7 +1582,6 @@ name = "nssa-core"
 version = "0.1.0"
 dependencies = [
  "chacha20",
- "k256",
  "risc0-zkvm",
  "serde",
 ]
@@ -1731,6 +1649,27 @@ dependencies = [
  "libm",
 ]
 
+[[package]]
+name = "num_enum"
+version = "0.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a973b4e44ce6cad84ce69d797acf9a044532e4184c4f267913d1b546a0727b7a"
+dependencies = [
+ "num_enum_derive",
+ "rustversion",
+]
+
+[[package]]
+name = "num_enum_derive"
+version = "0.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77e878c846a8abae00dd069496dbe8751b16ac1c3d6bd2a7283a938e8228f90d"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.104",
+]
+
 [[package]]
 name = "objc"
 version = "0.2.7"
@@ -2166,16 +2105,6 @@ dependencies = [
  "webpki-roots",
 ]
 
-[[package]]
-name = "rfc6979"
-version = "0.4.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2"
-dependencies = [
- "hmac",
- "subtle",
-]
-
 [[package]]
 name = "ring"
 version = "0.17.14"
@@ -2192,9 +2121,9 @@ dependencies = [
 
 [[package]]
 name = "risc0-binfmt"
-version = "3.0.1"
+version = "3.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e2af322c052ae9973054f67434bc953eae44dbac68054d304b46848634d2a45d"
+checksum = "1c8f97f81bcdead4101bca06469ecef481a2695cd04e7e877b49dea56a7f6f2a"
 dependencies = [
  "anyhow",
  "borsh",
@@ -2214,9 +2143,9 @@ dependencies = [
 
 [[package]]
 name = "risc0-build"
-version = "3.0.1"
+version = "3.0.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b042a8a6fdd02793c0acb526fbd3f0d92abe664f1158be38d45a92def57385b2"
+checksum = "1bbb512d728e011d03ce0958ca7954624ee13a215bcafd859623b3c63b2a3f60"
 dependencies = [
  "anyhow",
  "cargo_metadata",
@@ -2238,9 +2167,9 @@ dependencies = [
 
 [[package]]
 name = "risc0-circuit-keccak"
-version = "4.0.1"
+version = "4.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "966ba960d718123e16c603d6919a0c5d7cd2f872d428639ab5650108520f133a"
+checksum = "5f195f865ac1afdc21a172d7756fdcc21be18e13eb01d78d3d7f2b128fa881ba"
 dependencies = [
  "anyhow",
  "bytemuck",
@@ -2254,9 +2183,9 @@ dependencies = [
 
 [[package]]
 name = "risc0-circuit-recursion"
-version = "4.0.1"
+version = "4.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "30f6afac63cfd291aaa7f40edf976db5452698e66cd79f16ccbf7c0eb5a5d94e"
+checksum = "dca8f15c8abc0fd8c097aa7459879110334d191c63dd51d4c28881c4a497279e"
 dependencies = [
  "anyhow",
  "bytemuck",
@@ -2269,9 +2198,9 @@ dependencies = [
 
 [[package]]
 name = "risc0-circuit-rv32im"
-version = "4.0.1"
+version = "4.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8fdd0865593941a6800c65a7c3b48be8700f73eb597681f6f594c7465839ce8a"
+checksum = "ae1b0689f4a270a2f247b04397ebb431b8f64fe5170e98ee4f9d71bd04825205"
 dependencies = [
  "anyhow",
  "bit-vec",
@@ -2297,9 +2226,9 @@ dependencies = [
 
 [[package]]
 name = "risc0-groth16"
-version = "3.0.1"
+version = "3.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b68a622c69d0b97f511ee43db1d7f7f00b4dacead9c8aceae03fc5383f4764c1"
+checksum = "724285dc79604abfb2d40feaefe3e335420a6b293511661f77d6af62f1f5fae9"
 dependencies = [
  "anyhow",
  "ark-bn254",
@@ -2318,19 +2247,20 @@ dependencies = [
 
 [[package]]
 name = "risc0-zkos-v1compat"
-version = "2.0.2"
+version = "2.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "328c9f0ec5f6eb8b7624347b5dcf82729f304adbc364399825f3ab6f8588189c"
+checksum = "840c2228803557a8b7dc035a8f196516b6fd68c9dc6ac092f0c86241b5b1bafb"
 dependencies = [
  "include_bytes_aligned",
  "no_std_strings",
+ "risc0-zkvm-platform",
 ]
 
 [[package]]
 name = "risc0-zkp"
-version = "3.0.1"
+version = "3.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2c246f34b86a165308e37a064afa86e66ee7b8525f02bcf03f2124aaeedba04f"
+checksum = "ffb6bf356f469bb8744f72a07a37134c5812c1d55d6271bba80e87bdb7a58c8e"
 dependencies = [
  "anyhow",
  "blake2",
@@ -2353,9 +2283,9 @@ dependencies = [
 
 [[package]]
 name = "risc0-zkvm"
-version = "3.0.1"
+version = "3.0.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0fdf8d11f9e61cfd3e577fb6e6e11cc34ca247831c2555ee0a6a53deba9702d2"
+checksum = "3fcce11648a9ff60b8e7af2f0ce7fbf8d25275ab6d414cc91b9da69ee75bc978"
 dependencies = [
  "anyhow",
  "bincode",
@@ -2389,15 +2319,17 @@ dependencies = [
 
 [[package]]
 name = "risc0-zkvm-platform"
-version = "2.0.4"
+version = "2.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "06fc0e464f4ac44c3f1fd17b479e09e3ccbd1c40219837d750580b03030dca60"
+checksum = "5c826f09626ab2ae76671673e5a232548ddd95a34eece2ea4ced5f010383f95b"
 dependencies = [
  "bytemuck",
  "cfg-if",
  "getrandom 0.2.16",
  "getrandom 0.3.3",
  "libm",
+ "num_enum",
+ "paste",
  "stability",
 ]
 
@@ -2527,9 +2459,9 @@ checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f"
 
 [[package]]
 name = "rzup"
-version = "0.5.0"
+version = "0.5.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "76c6dfdbd72b2b0a537ad1e6256a3ce493ac87b9e89815465e393ee6595968e2"
+checksum = "5d2aed296f203fa64bcb4b52069356dd86d6ec578593985b919b6995bee1f0ae"
 dependencies = [
  "hex",
  "rsa",
@@ -2566,17 +2498,6 @@ dependencies = [
  "ref-cast",
  "serde",
  "serde_json",
-name = "sec1"
-version = "0.7.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc"
-dependencies = [
- "base16ct",
- "der",
- "generic-array",
- "pkcs8",
- "subtle",
- "zeroize",
 ]
 
 [[package]]
@@ -3671,8 +3592,3 @@ dependencies = [
  "quote",
  "syn 2.0.104",
 ]
-
-[[patch.unused]]
-name = "k256"
-version = "0.13.3"
-source = "git+https://github.com/risc0/RustCrypto-elliptic-curves?tag=k256%2Fv0.13.3-risczero.1#ff5d67b095cfcc2569b7789f2079ed87ef2c7756"

nssa/program_methods/guest/Cargo.toml

@@ -6,9 +6,5 @@ edition = "2021"
 [workspace]
 
 [dependencies]
-risc0-zkvm = { version = "3.0.1", default-features = false, features = ['std'] }
+risc0-zkvm = { version = "3.0.3", default-features = false, features = ['std'] }
 nssa-core = { path = "../../core" }
-
-[patch.crates-io]
-k256 = { git = "https://github.com/risc0/RustCrypto-elliptic-curves", tag = "k256/v0.13.3-risczero.1" }
-

nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -1,11 +1,12 @@
 use risc0_zkvm::{guest::env, serde::to_vec};
 
 use nssa_core::{
-    account::{Account, AccountWithMetadata, Commitment, Nullifier, NullifierPublicKey},
-    compute_root_associated_to_path,
+    account::{Account, AccountWithMetadata},
+    compute_digest_for_path,
+    encryption::Ciphertext,
     program::{validate_execution, ProgramOutput, DEFAULT_PROGRAM_ID},
-    CommitmentSetDigest, EncryptedAccountData, EphemeralPublicKey, EphemeralSecretKey,
-    IncomingViewingPublicKey, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
+    Commitment, CommitmentSetDigest, EncryptionScheme, Nullifier, NullifierPublicKey,
+    PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
 };
 
 fn main() {
@@ -19,7 +20,7 @@ fn main() {
     } = env::read();
 
     // TODO: Check that `program_execution_proof` is one of the allowed built-in programs
-    // assert!(BUILTIN_PROGRAM_IDS.contains(executing_program_id));
+    // assert_eq!(program_id, AUTHENTICATED_TRANSFER_PROGRAM_ID);
 
     // Check that `program_output` is consistent with the execution of the corresponding program.
     env::verify(program_id, &to_vec(&program_output).unwrap()).unwrap();
@@ -42,7 +43,7 @@ fn main() {
     // and will be populated next.
     let mut public_pre_states: Vec<AccountWithMetadata> = Vec::new();
     let mut public_post_states: Vec<Account> = Vec::new();
-    let mut encrypted_private_post_states: Vec<EncryptedAccountData> = Vec::new();
+    let mut ciphertexts: Vec<Ciphertext> = Vec::new();
     let mut new_commitments: Vec<Commitment> = Vec::new();
     let mut new_nullifiers: Vec<(Nullifier, CommitmentSetDigest)> = Vec::new();
 
@@ -67,23 +68,22 @@ fn main() {
             }
             1 | 2 => {
                 let new_nonce = private_nonces_iter.next().expect("Missing private nonce");
-                let (Npk, Ipk, esk) = private_keys_iter.next().expect("Missing private keys");
+                let (npk, shared_secret) = private_keys_iter.next().expect("Missing keys");
 
                 if visibility_mask[i] == 1 {
                     // Private account with authentication
                     let (nsk, membership_proof) =
                         private_auth_iter.next().expect("Missing private auth");
 
-                    // Verify Npk
-                    let expected_Npk = NullifierPublicKey::from(nsk);
-                    if &expected_Npk != Npk {
-                        panic!("Npk mismatch");
+                    // Verify the nullifier public key
+                    let expected_npk = NullifierPublicKey::from(nsk);
+                    if &expected_npk != npk {
+                        panic!("Nullifier public key mismatch");
                     }
 
                     // Compute commitment set digest associated with provided auth path
-                    let commitment_pre = Commitment::new(Npk, &pre_states[i].account);
-                    let set_digest =
-                        compute_root_associated_to_path(&commitment_pre, membership_proof);
+                    let commitment_pre = Commitment::new(npk, &pre_states[i].account);
+                    let set_digest = compute_digest_for_path(&commitment_pre, membership_proof);
 
                     // Check pre_state authorization
                     if !pre_states[i].is_authorized {
@@ -113,20 +113,18 @@ fn main() {
                 }
 
                 // Compute commitment
-                let commitment_post = Commitment::new(Npk, &post_with_updated_values);
+                let commitment_post = Commitment::new(npk, &post_with_updated_values);
 
                 // Encrypt and push post state
-                let encrypted_account = EncryptedAccountData::new(
+                let encrypted_account = EncryptionScheme::encrypt(
                     &post_with_updated_values,
-                    // &commitment_post,
-                    esk,
-                    Npk,
-                    Ipk,
+                    shared_secret,
+                    &commitment_post,
                     output_index,
                 );
 
                 new_commitments.push(commitment_post);
-                encrypted_private_post_states.push(encrypted_account);
+                ciphertexts.push(encrypted_account);
                 output_index += 1;
             }
             _ => panic!("Invalid visibility mask value"),
@@ -148,7 +146,7 @@ fn main() {
     let output = PrivacyPreservingCircuitOutput {
         public_pre_states,
         public_post_states,
-        encrypted_private_post_states,
+        ciphertexts,
         new_commitments,
         new_nullifiers,
     };

nssa/src/lib.rs

@@ -1,16 +1,18 @@
 pub mod address;
 pub mod error;
+mod merkle_tree;
 mod privacy_preserving_transaction;
 pub mod program;
 pub mod public_transaction;
 mod signature;
 mod state;
-mod merkle_tree;
 
 pub use address::Address;
+pub use privacy_preserving_transaction::{
+    PrivacyPreservingTransaction, circuit::execute_and_prove,
+};
 pub use public_transaction::PublicTransaction;
 pub use signature::PrivateKey;
 pub use signature::PublicKey;
 pub use signature::Signature;
 pub use state::V01State;
-

nssa/src/merkle_tree/mod.rs

@@ -1,5 +1,3 @@
-use std::collections::{HashMap, HashSet};
-
 use sha2::{Digest, Sha256};
 
 mod default_values;
@@ -47,8 +45,7 @@ impl MerkleTree {
 
     /// Number of levels required to hold all values
     fn depth(&self) -> usize {
-        let result = self.length.next_power_of_two().trailing_zeros() as usize;
-        result
+        self.length.next_power_of_two().trailing_zeros() as usize
     }
 
     fn get_node(&self, index: usize) -> &Node {
@@ -108,7 +105,6 @@ impl MerkleTree {
         self.set_node(node_index, node_hash);
         self.length += 1;
 
-        let root_index = self.root_index();
         for _ in 0..self.depth() {
             let parent_index = (node_index - 1) >> 1;
             let left_child = self.get_node((parent_index << 1) + 1);
@@ -121,14 +117,6 @@ impl MerkleTree {
         new_index
     }
 
-    pub fn new(values: &[Value]) -> Self {
-        let mut this = Self::with_capacity(values.len());
-        for value in values.iter().cloned() {
-            this.insert(value);
-        }
-        this
-    }
-
     pub fn get_authentication_path_for(&self, index: usize) -> Option<Vec<Node>> {
         if index >= self.length {
             return None;
@@ -155,22 +143,6 @@ impl MerkleTree {
     }
 }
 
-// Reference implementation
-fn verify_authentication_path(value: &Value, index: usize, path: &[Node], root: &Node) -> bool {
-    let mut result = hash_value(value);
-    let mut level_index = index;
-    for node in path {
-        let is_left_child = level_index & 1 == 0;
-        if is_left_child {
-            result = hash_two(&result, node);
-        } else {
-            result = hash_two(node, &result);
-        }
-        level_index >>= 1;
-    }
-    &result == root
-}
-
 fn prev_power_of_two(x: usize) -> usize {
     if x == 0 {
         return 0;
@@ -180,11 +152,19 @@ fn prev_power_of_two(x: usize) -> usize {
 
 #[cfg(test)]
 mod tests {
+    impl MerkleTree {
+        pub fn new(values: &[Value]) -> Self {
+            let mut this = Self::with_capacity(values.len());
+            for value in values.iter().cloned() {
+                this.insert(value);
+            }
+            this
+        }
+    }
+
     use hex_literal::hex;
-    use nssa_core::account::{Account, NullifierPublicKey};
 
     use super::*;
-
     #[test]
     fn test_empty_merkle_tree() {
         let tree = MerkleTree::with_capacity(4);
@@ -399,6 +379,22 @@ mod tests {
         assert_eq!(expected_tree, tree);
     }
 
+    // Reference implementation
+    fn verify_authentication_path(value: &Value, index: usize, path: &[Node], root: &Node) -> bool {
+        let mut result = hash_value(value);
+        let mut level_index = index;
+        for node in path {
+            let is_left_child = level_index & 1 == 0;
+            if is_left_child {
+                result = hash_two(&result, node);
+            } else {
+                result = hash_two(node, &result);
+            }
+            level_index >>= 1;
+        }
+        &result == root
+    }
+
     #[test]
     fn test_authentication_path_1() {
         let values = [[1; 32], [2; 32], [3; 32], [4; 32]];

nssa/src/privacy_preserving_transaction/circuit.rs

@@ -1,8 +1,8 @@
 use nssa_core::{
-    CommitmentSetDigest, EphemeralSecretKey, IncomingViewingPublicKey, MembershipProof,
-    PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
-    account::{Account, AccountWithMetadata, Nonce, NullifierPublicKey, NullifierSecretKey},
-    program::{InstructionData, ProgramId, ProgramOutput},
+    MembershipProof, NullifierPublicKey, NullifierSecretKey, PrivacyPreservingCircuitInput,
+    PrivacyPreservingCircuitOutput, SharedSecretKey,
+    account::AccountWithMetadata,
+    program::{InstructionData, ProgramOutput},
 };
 use risc0_zkvm::{ExecutorEnv, InnerReceipt, Receipt, default_prover};
 
@@ -10,6 +10,7 @@ use crate::{error::NssaError, program::Program};
 
 use program_methods::{PRIVACY_PRESERVING_CIRCUIT_ELF, PRIVACY_PRESERVING_CIRCUIT_ID};
 
+/// Proof of the privacy preserving execution circuit
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Proof(Vec<u8>);
 
@@ -21,16 +22,14 @@ impl Proof {
     }
 }
 
+/// Generates a proof of the execution of a NSSA program inside the privacy preserving execution
+/// circuit
 pub fn execute_and_prove(
     pre_states: &[AccountWithMetadata],
     instruction_data: &InstructionData,
     visibility_mask: &[u8],
     private_account_nonces: &[u128],
-    private_account_keys: &[(
-        NullifierPublicKey,
-        IncomingViewingPublicKey,
-        EphemeralSecretKey,
-    )],
+    private_account_keys: &[(NullifierPublicKey, SharedSecretKey)],
     private_account_auth: &[(NullifierSecretKey, MembershipProof)],
     program: &Program,
 ) -> Result<(PrivacyPreservingCircuitOutput, Proof), NssaError> {
@@ -90,18 +89,12 @@ fn execute_and_prove_program(
 #[cfg(test)]
 mod tests {
     use nssa_core::{
-        EncryptedAccountData,
-        account::{
-            Account, AccountWithMetadata, Commitment, Nullifier, NullifierPublicKey,
-            NullifierSecretKey,
-        },
+        Commitment, EncryptionScheme, Nullifier,
+        account::{Account, AccountWithMetadata},
     };
-    use risc0_zkvm::{InnerReceipt, Journal, Receipt};
 
     use crate::{
-        Address, V01State,
-        merkle_tree::MerkleTree,
-        privacy_preserving_transaction::circuit::{Proof, execute_and_prove},
+        privacy_preserving_transaction::circuit::execute_and_prove,
         program::Program,
         state::{
             CommitmentSet,
@@ -109,8 +102,6 @@ mod tests {
         },
     };
 
-    use rand::{Rng, RngCore, rngs::OsRng};
-
     use super::*;
 
     #[test]
@@ -147,12 +138,16 @@ mod tests {
 
         let expected_sender_pre = sender.clone();
         let recipient_keys = test_private_account_keys_1();
+
+        let esk = [3; 32];
+        let shared_secret = SharedSecretKey::new(&esk, &recipient_keys.ivk());
+
         let (output, proof) = execute_and_prove(
             &[sender, recipient],
             &Program::serialize_instruction(balance_to_move).unwrap(),
             &[0, 2],
             &[0xdeadbeef],
-            &[(recipient_keys.npk(), recipient_keys.ivk(), [3; 32])],
+            &[(recipient_keys.npk(), shared_secret.clone())],
             &[],
             &Program::authenticated_transfer_program(),
         )
@@ -166,12 +161,15 @@ mod tests {
         assert_eq!(sender_post, expected_sender_post);
         assert_eq!(output.new_commitments.len(), 1);
         assert_eq!(output.new_nullifiers.len(), 0);
-        assert_eq!(output.encrypted_private_post_states.len(), 1);
+        assert_eq!(output.ciphertexts.len(), 1);
 
-        let recipient_post = output.encrypted_private_post_states[0]
-            .clone()
-            .decrypt(&recipient_keys.isk, 0)
-            .unwrap();
+        let recipient_post = EncryptionScheme::decrypt(
+            &output.ciphertexts[0],
+            &shared_secret,
+            &output.new_commitments[0],
+            0,
+        )
+        .unwrap();
         assert_eq!(recipient_post, expected_recipient_post);
     }
 
@@ -196,7 +194,7 @@ mod tests {
         let balance_to_move: u128 = 37;
 
         let mut commitment_set = CommitmentSet::with_capacity(2);
-        commitment_set.extend(&[commitment_sender.clone()]);
+        commitment_set.extend(std::slice::from_ref(&commitment_sender));
 
         let expected_new_nullifiers = vec![(
             Nullifier::new(&commitment_sender, &sender_keys.nsk),
@@ -222,14 +220,20 @@ mod tests {
             Commitment::new(&recipient_keys.npk(), &expected_private_account_2),
         ];
 
+        let esk_1 = [3; 32];
+        let shared_secret_1 = SharedSecretKey::new(&esk_1, &sender_keys.ivk());
+
+        let esk_2 = [5; 32];
+        let shared_secret_2 = SharedSecretKey::new(&esk_2, &recipient_keys.ivk());
+
         let (output, proof) = execute_and_prove(
             &[sender_pre.clone(), recipient],
             &Program::serialize_instruction(balance_to_move).unwrap(),
             &[1, 2],
             &[0xdeadbeef1, 0xdeadbeef2],
             &[
-                (sender_keys.npk(), sender_keys.ivk(), [3; 32]),
-                (recipient_keys.npk(), recipient_keys.ivk(), [5; 32]),
+                (sender_keys.npk(), shared_secret_1.clone()),
+                (recipient_keys.npk(), shared_secret_2.clone()),
             ],
             &[(
                 sender_keys.nsk,
@@ -244,18 +248,24 @@ mod tests {
         assert!(output.public_post_states.is_empty());
         assert_eq!(output.new_commitments, expected_new_commitments);
         assert_eq!(output.new_nullifiers, expected_new_nullifiers);
-        assert_eq!(output.encrypted_private_post_states.len(), 2);
+        assert_eq!(output.ciphertexts.len(), 2);
 
-        let recipient_post_1 = output.encrypted_private_post_states[0]
-            .clone()
-            .decrypt(&sender_keys.isk, 0)
-            .unwrap();
-        assert_eq!(recipient_post_1, expected_private_account_1);
+        let sender_post = EncryptionScheme::decrypt(
+            &output.ciphertexts[0],
+            &shared_secret_1,
+            &expected_new_commitments[0],
+            0,
+        )
+        .unwrap();
+        assert_eq!(sender_post, expected_private_account_1);
 
-        let recipient_post_2 = output.encrypted_private_post_states[1]
-            .clone()
-            .decrypt(&recipient_keys.isk, 1)
-            .unwrap();
-        assert_eq!(recipient_post_2, expected_private_account_2);
+        let recipient_post = EncryptionScheme::decrypt(
+            &output.ciphertexts[1],
+            &shared_secret_2,
+            &expected_new_commitments[1],
+            1,
+        )
+        .unwrap();
+        assert_eq!(recipient_post, expected_private_account_2);
     }
 }

nssa/src/privacy_preserving_transaction/encoding.rs

@@ -1,19 +1,44 @@
-// TODO: Consider switching to deriving Borsh
-
 use std::io::{Cursor, Read};
 
 use nssa_core::{
-    EncryptedAccountData,
-    account::{Account, Commitment, Nullifier},
+    Commitment, Nullifier,
+    account::Account,
+    encryption::{Ciphertext, EphemeralPublicKey},
 };
 
-use crate::{Address, error::NssaError};
+use crate::{
+    Address, error::NssaError, privacy_preserving_transaction::message::EncryptedAccountData,
+};
 
 use super::message::Message;
 
 const MESSAGE_ENCODING_PREFIX_LEN: usize = 22;
 const MESSAGE_ENCODING_PREFIX: &[u8; MESSAGE_ENCODING_PREFIX_LEN] = b"\x01/NSSA/v0.1/TxMessage/";
 
+impl EncryptedAccountData {
+    pub(crate) fn to_bytes(&self) -> Vec<u8> {
+        let mut bytes = self.ciphertext.to_bytes();
+        bytes.extend_from_slice(&self.epk.to_bytes());
+        bytes.push(self.view_tag);
+        bytes
+    }
+
+    pub fn from_cursor(cursor: &mut Cursor<&[u8]>) -> Result<Self, NssaError> {
+        let ciphertext = Ciphertext::from_cursor(cursor)?;
+        let epk = EphemeralPublicKey::from_cursor(cursor)?;
+
+        let mut tag_bytes = [0; 1];
+        cursor.read_exact(&mut tag_bytes)?;
+        let view_tag = tag_bytes[0];
+
+        Ok(Self {
+            ciphertext,
+            epk,
+            view_tag,
+        })
+    }
+}
+
 impl Message {
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
         let mut bytes = MESSAGE_ENCODING_PREFIX.to_vec();
@@ -63,6 +88,7 @@ impl Message {
         bytes
     }
 
+    #[allow(unused)]
     pub(crate) fn from_cursor(cursor: &mut Cursor<&[u8]>) -> Result<Self, NssaError> {
         let prefix = {
             let mut this = [0u8; MESSAGE_ENCODING_PREFIX_LEN];
@@ -90,7 +116,7 @@ impl Message {
         // Nonces
         cursor.read_exact(&mut len_bytes)?;
         let nonces_len = u32::from_le_bytes(len_bytes) as usize;
-        let mut nonces = Vec::with_capacity(nonces_len as usize);
+        let mut nonces = Vec::with_capacity(nonces_len);
         for _ in 0..nonces_len {
             let mut buf = [0u8; 16];
             cursor.read_exact(&mut buf)?;
@@ -128,7 +154,7 @@ impl Message {
         for _ in 0..new_nullifiers_len {
             let nullifier = Nullifier::from_cursor(cursor)?;
             let mut commitment_set_digest = [0; 32];
-            cursor.read_exact(&mut commitment_set_digest);
+            cursor.read_exact(&mut commitment_set_digest)?;
             new_nullifiers.push((nullifier, commitment_set_digest));
         }
 

nssa/src/privacy_preserving_transaction/message.rs

@@ -1,9 +1,17 @@
 use nssa_core::{
-    CommitmentSetDigest, EncryptedAccountData,
-    account::{Account, Commitment, Nonce, Nullifier},
+    Commitment, CommitmentSetDigest, Nullifier, PrivacyPreservingCircuitOutput,
+    account::{Account, Nonce},
+    encryption::{Ciphertext, EphemeralPublicKey},
 };
 
-use crate::Address;
+use crate::{Address, error::NssaError};
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct EncryptedAccountData {
+    pub(crate) ciphertext: Ciphertext,
+    pub(crate) epk: EphemeralPublicKey,
+    pub(crate) view_tag: u8,
+}
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Message {
@@ -16,22 +24,36 @@ pub struct Message {
 }
 
 impl Message {
-    pub fn new(
+    pub fn try_from_circuit_output(
         public_addresses: Vec<Address>,
         nonces: Vec<Nonce>,
-        public_post_states: Vec<Account>,
-        encrypted_private_post_states: Vec<EncryptedAccountData>,
-        new_commitments: Vec<Commitment>,
-        new_nullifiers: Vec<(Nullifier, CommitmentSetDigest)>,
-    ) -> Self {
-        Self {
+        ephemeral_public_keys: Vec<EphemeralPublicKey>,
+        output: PrivacyPreservingCircuitOutput,
+    ) -> Result<Self, NssaError> {
+        if ephemeral_public_keys.len() != output.ciphertexts.len() {
+            return Err(NssaError::InvalidInput(
+                "Ephemeral public keys and ciphertexts length mismatch".into(),
+            ));
+        }
+
+        let encrypted_private_post_states = output
+            .ciphertexts
+            .into_iter()
+            .zip(ephemeral_public_keys)
+            .map(|(ciphertext, epk)| EncryptedAccountData {
+                ciphertext,
+                epk,
+                view_tag: 0, // TODO: implement
+            })
+            .collect();
+        Ok(Self {
             public_addresses,
             nonces,
-            public_post_states,
+            public_post_states: output.public_post_states,
             encrypted_private_post_states,
-            new_commitments,
-            new_nullifiers,
-        }
+            new_commitments: output.new_commitments,
+            new_nullifiers: output.new_nullifiers,
+        })
     }
 }
 
@@ -39,9 +61,7 @@ impl Message {
 pub mod tests {
     use std::io::Cursor;
 
-    use nssa_core::account::{
-        Account, Commitment, Nullifier, NullifierPublicKey, NullifierSecretKey,
-    };
+    use nssa_core::{Commitment, Nullifier, NullifierPublicKey, account::Account};
 
     use crate::{Address, privacy_preserving_transaction::message::Message};
 
@@ -52,8 +72,8 @@ pub mod tests {
         let nsk1 = [11; 32];
         let nsk2 = [12; 32];
 
-        let Npk1 = NullifierPublicKey::from(&nsk1);
-        let Npk2 = NullifierPublicKey::from(&nsk2);
+        let npk1 = NullifierPublicKey::from(&nsk1);
+        let npk2 = NullifierPublicKey::from(&nsk2);
 
         let public_addresses = vec![Address::new([1; 32])];
 
@@ -63,9 +83,9 @@ pub mod tests {
 
         let encrypted_private_post_states = Vec::new();
 
-        let new_commitments = vec![Commitment::new(&Npk2, &account2)];
+        let new_commitments = vec![Commitment::new(&npk2, &account2)];
 
-        let old_commitment = Commitment::new(&Npk1, &account1);
+        let old_commitment = Commitment::new(&npk1, &account1);
         let new_nullifiers = vec![(Nullifier::new(&old_commitment, &nsk1), [0; 32])];
 
         Message {
@@ -78,23 +98,6 @@ pub mod tests {
         }
     }
 
-    #[test]
-    fn test_constructor() {
-        let message = message_for_tests();
-        let expected_message = message.clone();
-
-        let message = Message::new(
-            message.public_addresses,
-            message.nonces,
-            message.public_post_states,
-            message.encrypted_private_post_states,
-            message.new_commitments,
-            message.new_nullifiers,
-        );
-
-        assert_eq!(message, expected_message);
-    }
-
     #[test]
     fn test_message_serialization_roundtrip() {
         let message = message_for_tests();

nssa/src/privacy_preserving_transaction/mod.rs

@@ -1,10 +1,8 @@
 mod encoding;
-mod message;
-mod transaction;
-mod witness_set;
+pub mod message;
+pub mod transaction;
+pub mod witness_set;
 
 pub mod circuit;
 
-pub use message::Message;
 pub use transaction::PrivacyPreservingTransaction;
-pub use witness_set::WitnessSet;

nssa/src/privacy_preserving_transaction/transaction.rs

@@ -1,10 +1,13 @@
 use std::collections::{HashMap, HashSet};
 
-use nssa_core::account::{Account, AccountWithMetadata, Commitment, Nullifier};
-use nssa_core::{CommitmentSetDigest, EncryptedAccountData, PrivacyPreservingCircuitOutput};
+use nssa_core::{
+    Commitment, CommitmentSetDigest, Nullifier, PrivacyPreservingCircuitOutput,
+    account::{Account, AccountWithMetadata},
+};
 
 use crate::error::NssaError;
 use crate::privacy_preserving_transaction::circuit::Proof;
+use crate::privacy_preserving_transaction::message::EncryptedAccountData;
 use crate::{Address, V01State};
 
 use super::message::Message;
@@ -145,7 +148,11 @@ fn check_privacy_preserving_circuit_proof_is_valid(
     let output = PrivacyPreservingCircuitOutput {
         public_pre_states: public_pre_states.to_vec(),
         public_post_states: public_post_states.to_vec(),
-        encrypted_private_post_states: encrypted_private_post_states.to_vec(),
+        ciphertexts: encrypted_private_post_states
+            .iter()
+            .cloned()
+            .map(|value| value.ciphertext)
+            .collect(),
         new_commitments: new_commitments.to_vec(),
         new_nullifiers: new_nullifiers.to_vec(),
     };

nssa/src/program.rs

@@ -1,13 +1,9 @@
-use borsh::{BorshDeserialize, BorshSerialize};
 use nssa_core::{
     account::{Account, AccountWithMetadata},
-    program::{DEFAULT_PROGRAM_ID, InstructionData, ProgramId, ProgramOutput},
+    program::{InstructionData, ProgramId, ProgramOutput},
 };
 use program_methods::{AUTHENTICATED_TRANSFER_ELF, AUTHENTICATED_TRANSFER_ID};
-use risc0_zkvm::{
-    ExecutorEnv, ExecutorEnvBuilder, Journal, Receipt, default_executor, default_prover,
-    serde::to_vec,
-};
+use risc0_zkvm::{ExecutorEnv, ExecutorEnvBuilder, default_executor, serde::to_vec};
 use serde::Serialize;
 
 use crate::error::NssaError;
@@ -50,9 +46,7 @@ impl Program {
             .map_err(|e| NssaError::ProgramExecutionFailed(e.to_string()))?;
 
         // Get outputs
-        let ProgramOutput {
-            mut post_states, ..
-        } = session_info
+        let ProgramOutput { post_states, .. } = session_info
             .journal
             .decode()
             .map_err(|e| NssaError::ProgramExecutionFailed(e.to_string()))?;
@@ -83,11 +77,7 @@ impl Program {
 
 #[cfg(test)]
 mod tests {
-    use nssa_core::{
-        account::{Account, AccountWithMetadata},
-        program::ProgramOutput,
-    };
-    use risc0_zkvm::{InnerReceipt, Receipt, serde::to_vec};
+    use nssa_core::account::{Account, AccountWithMetadata};
 
     use crate::program::Program;
 

nssa/src/public_transaction/encoding.rs

@@ -11,8 +11,7 @@ use crate::{
 };
 
 const MESSAGE_ENCODING_PREFIX_LEN: usize = 22;
-const MESSAGE_ENCODING_PREFIX: &[u8; MESSAGE_ENCODING_PREFIX_LEN] =
-    b"\x00/NSSA/v0.1/TxMessage/";
+const MESSAGE_ENCODING_PREFIX: &[u8; MESSAGE_ENCODING_PREFIX_LEN] = b"\x00/NSSA/v0.1/TxMessage/";
 
 impl Message {
     /// Serializes a `Message` into bytes in the following layout:

nssa/src/state.rs

@@ -4,8 +4,8 @@ use crate::{
     public_transaction::PublicTransaction,
 };
 use nssa_core::{
-    CommitmentSetDigest, MembershipProof,
-    account::{Account, Commitment, Nullifier},
+    Commitment, CommitmentSetDigest, MembershipProof, Nullifier,
+    account::Account,
     program::{DEFAULT_PROGRAM_ID, ProgramId},
 };
 use std::collections::{HashMap, HashSet};
@@ -13,7 +13,7 @@ use std::collections::{HashMap, HashSet};
 pub(crate) struct CommitmentSet {
     merkle_tree: MerkleTree,
     commitments: HashMap<Commitment, usize>,
-    pub root_history: HashSet<CommitmentSetDigest>,
+    root_history: HashSet<CommitmentSetDigest>,
 }
 
 impl CommitmentSet {
@@ -21,13 +21,12 @@ impl CommitmentSet {
         self.merkle_tree.root()
     }
 
-    pub(crate) fn get_proof_for(&self, commitment: &Commitment) -> Option<MembershipProof> {
+    pub fn get_proof_for(&self, commitment: &Commitment) -> Option<MembershipProof> {
         let index = *self.commitments.get(commitment)?;
-        let proof = self
-            .merkle_tree
+
+        self.merkle_tree
             .get_authentication_path_for(index)
-            .map(|path| (index, path));
-        proof
+            .map(|path| (index, path))
     }
 
     pub(crate) fn extend(&mut self, commitments: &[Commitment]) {
@@ -55,7 +54,7 @@ type NullifierSet = HashSet<Nullifier>;
 
 pub struct V01State {
     public_state: HashMap<Address, Account>,
-    pub private_state: (CommitmentSet, NullifierSet),
+    private_state: (CommitmentSet, NullifierSet),
     builtin_programs: HashMap<ProgramId, Program>,
 }
 
@@ -161,6 +160,10 @@ impl V01State {
             .unwrap_or(Account::default())
     }
 
+    pub fn get_proof_for_commitment(&self, commitment: &Commitment) -> Option<MembershipProof> {
+        self.private_state.0.get_proof_for(commitment)
+    }
+
     pub(crate) fn builtin_programs(&self) -> &HashMap<ProgramId, Program> {
         &self.builtin_programs
     }
@@ -201,16 +204,6 @@ impl V01State {
         }
         Ok(())
     }
-
-    pub(crate) fn check_commitment_set_digest_is_valid(
-        &self,
-        commitment_set_digest: &CommitmentSetDigest,
-    ) -> bool {
-        self.private_state
-            .0
-            .root_history
-            .contains(commitment_set_digest)
-    }
 }
 
 #[cfg(test)]
@@ -222,20 +215,18 @@ pub mod tests {
         Address, PublicKey, PublicTransaction, V01State,
         error::NssaError,
         privacy_preserving_transaction::{
-            Message, PrivacyPreservingTransaction, WitnessSet, circuit,
+            PrivacyPreservingTransaction, circuit, message::Message, witness_set::WitnessSet,
         },
         program::Program,
         public_transaction,
         signature::PrivateKey,
     };
+
     use nssa_core::{
-        IncomingViewingPublicKey,
-        account::{
-            Account, AccountWithMetadata, Commitment, Nonce, Nullifier, NullifierPublicKey,
-            NullifierSecretKey,
-        },
+        Commitment, Nullifier, NullifierPublicKey, NullifierSecretKey, SharedSecretKey,
+        account::{Account, AccountWithMetadata, Nonce},
+        encryption::{EphemeralPublicKey, IncomingViewingPublicKey},
     };
-    use program_methods::AUTHENTICATED_TRANSFER_ID;
 
     fn transfer_transaction(
         from: Address,
@@ -485,7 +476,7 @@ pub mod tests {
         }
 
         pub fn with_private_account(mut self, keys: &TestPrivateKeys, account: &Account) -> Self {
-            let commitment = Commitment::new(&keys.npk(), &account);
+            let commitment = Commitment::new(&keys.npk(), account);
             self.private_state.0.extend(&[commitment]);
             self
         }
@@ -780,11 +771,11 @@ pub mod tests {
         balance_to_move: u128,
         state: &V01State,
     ) -> PrivacyPreservingTransaction {
-        let esk = [3; 32];
         let sender = AccountWithMetadata {
             account: state.get_account_by_address(&sender_keys.address()),
             is_authorized: true,
         };
+
         let sender_nonce = sender.account.nonce;
 
         let recipient = AccountWithMetadata {
@@ -792,25 +783,28 @@ pub mod tests {
             is_authorized: false,
         };
 
+        let esk = [3; 32];
+        let shared_secret = SharedSecretKey::new(&esk, &recipient_keys.ivk());
+        let epk = EphemeralPublicKey::from_scalar(esk);
+
         let (output, proof) = circuit::execute_and_prove(
             &[sender, recipient],
             &Program::serialize_instruction(balance_to_move).unwrap(),
             &[0, 2],
             &[0xdeadbeef],
-            &[(recipient_keys.npk(), recipient_keys.ivk(), esk)],
+            &[(recipient_keys.npk(), shared_secret)],
             &[],
             &Program::authenticated_transfer_program(),
         )
         .unwrap();
 
-        let message = Message::new(
+        let message = Message::try_from_circuit_output(
             vec![sender_keys.address()],
             vec![sender_nonce],
-            output.public_post_states,
-            output.encrypted_private_post_states,
-            output.new_commitments.clone(),
-            output.new_nullifiers,
-        );
+            vec![epk],
+            output,
+        )
+        .unwrap();
 
         let witness_set = WitnessSet::for_message(&message, proof, &[&sender_keys.signing_key]);
         PrivacyPreservingTransaction::new(message, witness_set)
@@ -825,7 +819,7 @@ pub mod tests {
         state: &V01State,
     ) -> PrivacyPreservingTransaction {
         let program = Program::authenticated_transfer_program();
-        let sender_commitment = Commitment::new(&sender_keys.npk(), &sender_private_account);
+        let sender_commitment = Commitment::new(&sender_keys.npk(), sender_private_account);
         let sender_pre = AccountWithMetadata {
             account: sender_private_account.clone(),
             is_authorized: true,
@@ -835,35 +829,33 @@ pub mod tests {
             is_authorized: false,
         };
 
+        let esk_1 = [3; 32];
+        let shared_secret_1 = SharedSecretKey::new(&esk_1, &sender_keys.ivk());
+        let epk_1 = EphemeralPublicKey::from_scalar(esk_1);
+
+        let esk_2 = [3; 32];
+        let shared_secret_2 = SharedSecretKey::new(&esk_2, &recipient_keys.ivk());
+        let epk_2 = EphemeralPublicKey::from_scalar(esk_2);
+
         let (output, proof) = circuit::execute_and_prove(
             &[sender_pre, recipient_pre],
             &Program::serialize_instruction(balance_to_move).unwrap(),
             &[1, 2],
             &new_nonces,
             &[
-                (sender_keys.npk(), sender_keys.ivk(), [3; 32]),
-                (recipient_keys.npk(), recipient_keys.ivk(), [4; 32]),
+                (sender_keys.npk(), shared_secret_1),
+                (recipient_keys.npk(), shared_secret_2),
             ],
             &[(
                 sender_keys.nsk,
-                state
-                    .private_state
-                    .0
-                    .get_proof_for(&sender_commitment)
-                    .unwrap(),
+                state.get_proof_for_commitment(&sender_commitment).unwrap(),
             )],
             &program,
         )
         .unwrap();
 
-        let message = Message::new(
-            vec![],
-            vec![],
-            output.public_post_states,
-            output.encrypted_private_post_states,
-            output.new_commitments.clone(),
-            output.new_nullifiers,
-        );
+        let message =
+            Message::try_from_circuit_output(vec![], vec![], vec![epk_1, epk_2], output).unwrap();
 
         let witness_set = WitnessSet::for_message(&message, proof, &[]);
 
@@ -879,7 +871,7 @@ pub mod tests {
         state: &V01State,
     ) -> PrivacyPreservingTransaction {
         let program = Program::authenticated_transfer_program();
-        let sender_commitment = Commitment::new(&sender_keys.npk(), &sender_private_account);
+        let sender_commitment = Commitment::new(&sender_keys.npk(), sender_private_account);
         let sender_pre = AccountWithMetadata {
             account: sender_private_account.clone(),
             is_authorized: true,
@@ -889,32 +881,27 @@ pub mod tests {
             is_authorized: false,
         };
 
+        let esk = [3; 32];
+        let shared_secret = SharedSecretKey::new(&esk, &sender_keys.ivk());
+        let epk = EphemeralPublicKey::from_scalar(esk);
+
         let (output, proof) = circuit::execute_and_prove(
             &[sender_pre, recipient_pre],
             &Program::serialize_instruction(balance_to_move).unwrap(),
             &[1, 0],
             &[new_nonce],
-            &[(sender_keys.npk(), sender_keys.ivk(), [3; 32])],
+            &[(sender_keys.npk(), shared_secret)],
             &[(
                 sender_keys.nsk,
-                state
-                    .private_state
-                    .0
-                    .get_proof_for(&sender_commitment)
-                    .unwrap(),
+                state.get_proof_for_commitment(&sender_commitment).unwrap(),
             )],
             &program,
         )
         .unwrap();
 
-        let message = Message::new(
-            vec![recipient_address.clone()],
-            vec![],
-            output.public_post_states,
-            output.encrypted_private_post_states,
-            output.new_commitments.clone(),
-            output.new_nullifiers,
-        );
+        let message =
+            Message::try_from_circuit_output(vec![*recipient_address], vec![], vec![epk], output)
+                .unwrap();
 
         let witness_set = WitnessSet::for_message(&message, proof, &[]);
 

nssa/test_program_methods/Cargo.toml

@@ -4,7 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [build-dependencies]
-risc0-build = { version = "3.0.1" }
+risc0-build = { version = "3.0.3" }
 
 [package.metadata.risc0]
 methods = ["guest"]

nssa/test_program_methods/guest/Cargo.toml

@@ -6,5 +6,5 @@ edition = "2021"
 [workspace]
 
 [dependencies]
-risc0-zkvm = { version = "3.0.1", default-features = false, features = ['std'] }
+risc0-zkvm = { version = "3.0.3", default-features = false, features = ['std'] }
 nssa-core = { path = "../../core" }

wallet/Cargo.toml

@@ -18,7 +18,7 @@ reqwest.workspace = true
 thiserror.workspace = true
 tokio.workspace = true
 tempfile.workspace = true
-risc0-zkvm = "3.0.1"
+risc0-zkvm = "3.0.3"
 hex.workspace = true
 actix-rt.workspace = true
 clap.workspace = true