Merge branch 'main' into simple_amm

program_methods/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "program_methods"
+version = "0.1.0"
+edition = "2024"
+
+[build-dependencies]
+risc0-build.workspace = true
+
+[package.metadata.risc0]
+methods = ["guest"]

program_methods/build.rs

@@ -0,0 +1,3 @@
+fn main() {
+    risc0_build::embed_methods();
+}

program_methods/guest/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "programs"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+nssa_core.workspace = true
+
+risc0-zkvm.workspace = true
+serde = { workspace = true, default-features = false }

program_methods/guest/src/bin/authenticated_transfer.rs

@@ -0,0 +1,90 @@
+use nssa_core::{
+    account::{Account, AccountWithMetadata},
+    program::{
+        AccountPostState, DEFAULT_PROGRAM_ID, ProgramInput, read_nssa_inputs, write_nssa_outputs,
+    },
+};
+
+/// Initializes a default account under the ownership of this program.
+fn initialize_account(pre_state: AccountWithMetadata) -> AccountPostState {
+    let account_to_claim = AccountPostState::new_claimed(pre_state.account.clone());
+    let is_authorized = pre_state.is_authorized;
+
+    // Continue only if the account to claim has default values
+    if account_to_claim.account() != &Account::default() {
+        panic!("Account must be uninitialized");
+    }
+
+    // Continue only if the owner authorized this operation
+    if !is_authorized {
+        panic!("Account must be authorized");
+    }
+
+    account_to_claim
+}
+
+/// Transfers `balance_to_move` native balance from `sender` to `recipient`.
+fn transfer(
+    sender: AccountWithMetadata,
+    recipient: AccountWithMetadata,
+    balance_to_move: u128,
+) -> Vec<AccountPostState> {
+    // Continue only if the sender has authorized this operation
+    if !sender.is_authorized {
+        panic!("Sender must be authorized");
+    }
+
+    // Continue only if the sender has enough balance
+    if sender.account.balance < balance_to_move {
+        panic!("Sender has insufficient balance");
+    }
+
+    // Create accounts post states, with updated balances
+    let sender_post = {
+        // Modify sender's balance
+        let mut sender_post_account = sender.account.clone();
+        sender_post_account.balance -= balance_to_move;
+        AccountPostState::new(sender_post_account)
+    };
+
+    let recipient_post = {
+        // Modify recipient's balance
+        let mut recipient_post_account = recipient.account.clone();
+        recipient_post_account.balance += balance_to_move;
+
+        // Claim recipient account if it has default program owner
+        if recipient_post_account.program_owner == DEFAULT_PROGRAM_ID {
+            AccountPostState::new_claimed(recipient_post_account)
+        } else {
+            AccountPostState::new(recipient_post_account)
+        }
+    };
+
+    vec![sender_post, recipient_post]
+}
+
+/// A transfer of balance program.
+/// To be used both in public and private contexts.
+fn main() {
+    // Read input accounts.
+    let (
+        ProgramInput {
+            pre_states,
+            instruction: balance_to_move,
+        },
+        instruction_words,
+    ) = read_nssa_inputs();
+
+    let post_states = match (pre_states.as_slice(), balance_to_move) {
+        ([account_to_claim], 0) => {
+            let post = initialize_account(account_to_claim.clone());
+            vec![post]
+        }
+        ([sender, recipient], balance_to_move) => {
+            transfer(sender.clone(), recipient.clone(), balance_to_move)
+        }
+        _ => panic!("invalid params"),
+    };
+
+    write_nssa_outputs(instruction_words, pre_states, post_states);
+}

program_methods/guest/src/bin/pinata.rs

@@ -0,0 +1,84 @@
+use nssa_core::program::{AccountPostState, ProgramInput, read_nssa_inputs, write_nssa_outputs};
+use risc0_zkvm::sha::{Impl, Sha256};
+
+const PRIZE: u128 = 150;
+
+type Instruction = u128;
+
+struct Challenge {
+    difficulty: u8,
+    seed: [u8; 32],
+}
+
+impl Challenge {
+    fn new(bytes: &[u8]) -> Self {
+        assert_eq!(bytes.len(), 33);
+        let difficulty = bytes[0];
+        assert!(difficulty <= 32);
+
+        let mut seed = [0; 32];
+        seed.copy_from_slice(&bytes[1..]);
+        Self { difficulty, seed }
+    }
+
+    // Checks if the leftmost `self.difficulty` number of bytes of SHA256(self.data || solution) are
+    // zero.
+    fn validate_solution(&self, solution: Instruction) -> bool {
+        let mut bytes = [0; 32 + 16];
+        bytes[..32].copy_from_slice(&self.seed);
+        bytes[32..].copy_from_slice(&solution.to_le_bytes());
+        let digest: [u8; 32] = Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap();
+        let difficulty = self.difficulty as usize;
+        digest[..difficulty].iter().all(|&b| b == 0)
+    }
+
+    fn next_data(self) -> [u8; 33] {
+        let mut result = [0; 33];
+        result[0] = self.difficulty;
+        result[1..].copy_from_slice(Impl::hash_bytes(&self.seed).as_bytes());
+        result
+    }
+}
+
+/// A pinata program
+fn main() {
+    // Read input accounts.
+    // It is expected to receive only two accounts: [pinata_account, winner_account]
+    let (
+        ProgramInput {
+            pre_states,
+            instruction: solution,
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
+
+    let [pinata, winner] = match pre_states.try_into() {
+        Ok(array) => array,
+        Err(_) => return,
+    };
+
+    let data = Challenge::new(&pinata.account.data);
+
+    if !data.validate_solution(solution) {
+        return;
+    }
+
+    let mut pinata_post = pinata.account.clone();
+    let mut winner_post = winner.account.clone();
+    pinata_post.balance -= PRIZE;
+    pinata_post.data = data
+        .next_data()
+        .to_vec()
+        .try_into()
+        .expect("33 bytes should fit into Data");
+    winner_post.balance += PRIZE;
+
+    write_nssa_outputs(
+        instruction_words,
+        vec![pinata, winner],
+        vec![
+            AccountPostState::new(pinata_post),
+            AccountPostState::new(winner_post),
+        ],
+    );
+}

program_methods/guest/src/bin/pinata_token.rs

@@ -0,0 +1,117 @@
+use nssa_core::{
+    account::Data,
+    program::{
+        AccountPostState, ChainedCall, PdaSeed, ProgramInput, read_nssa_inputs,
+        write_nssa_outputs_with_chained_call,
+    },
+};
+use risc0_zkvm::{
+    serde::to_vec,
+    sha::{Impl, Sha256},
+};
+
+const PRIZE: u128 = 150;
+
+type Instruction = u128;
+
+struct Challenge {
+    difficulty: u8,
+    seed: [u8; 32],
+}
+
+impl Challenge {
+    fn new(bytes: &[u8]) -> Self {
+        assert_eq!(bytes.len(), 33);
+        let difficulty = bytes[0];
+        assert!(difficulty <= 32);
+
+        let mut seed = [0; 32];
+        seed.copy_from_slice(&bytes[1..]);
+        Self { difficulty, seed }
+    }
+
+    // Checks if the leftmost `self.difficulty` number of bytes of SHA256(self.data || solution) are
+    // zero.
+    fn validate_solution(&self, solution: Instruction) -> bool {
+        let mut bytes = [0; 32 + 16];
+        bytes[..32].copy_from_slice(&self.seed);
+        bytes[32..].copy_from_slice(&solution.to_le_bytes());
+        let digest: [u8; 32] = Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap();
+        let difficulty = self.difficulty as usize;
+        digest[..difficulty].iter().all(|&b| b == 0)
+    }
+
+    fn next_data(self) -> Data {
+        let mut result = [0; 33];
+        result[0] = self.difficulty;
+        result[1..].copy_from_slice(Impl::hash_bytes(&self.seed).as_bytes());
+        result.to_vec().try_into().expect("should fit")
+    }
+}
+
+/// A pinata program
+fn main() {
+    // Read input accounts.
+    // It is expected to receive three accounts: [pinata_definition, pinata_token_holding,
+    // winner_token_holding]
+    let (
+        ProgramInput {
+            pre_states,
+            instruction: solution,
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
+
+    let [
+        pinata_definition,
+        pinata_token_holding,
+        winner_token_holding,
+    ] = match pre_states.try_into() {
+        Ok(array) => array,
+        Err(_) => return,
+    };
+
+    let data = Challenge::new(&pinata_definition.account.data);
+
+    if !data.validate_solution(solution) {
+        return;
+    }
+
+    let mut pinata_definition_post = pinata_definition.account.clone();
+    let pinata_token_holding_post = pinata_token_holding.account.clone();
+    let winner_token_holding_post = winner_token_holding.account.clone();
+    pinata_definition_post.data = data.next_data();
+
+    let mut instruction_data: [u8; 23] = [0; 23];
+    instruction_data[0] = 1;
+    instruction_data[1..17].copy_from_slice(&PRIZE.to_le_bytes());
+
+    // Flip authorization to true for chained call
+    let mut pinata_token_holding_for_chain_call = pinata_token_holding.clone();
+    pinata_token_holding_for_chain_call.is_authorized = true;
+
+    let chained_calls = vec![ChainedCall {
+        program_id: pinata_token_holding_post.program_owner,
+        instruction_data: to_vec(&instruction_data).unwrap(),
+        pre_states: vec![
+            pinata_token_holding_for_chain_call,
+            winner_token_holding.clone(),
+        ],
+        pda_seeds: vec![PdaSeed::new([0; 32])],
+    }];
+
+    write_nssa_outputs_with_chained_call(
+        instruction_words,
+        vec![
+            pinata_definition,
+            pinata_token_holding,
+            winner_token_holding,
+        ],
+        vec![
+            AccountPostState::new(pinata_definition_post),
+            AccountPostState::new(pinata_token_holding_post),
+            AccountPostState::new(winner_token_holding_post),
+        ],
+        chained_calls,
+    );
+}

program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -0,0 +1,277 @@
+use std::collections::HashMap;
+
+use nssa_core::{
+    Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, Nullifier,
+    NullifierPublicKey, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
+    account::{Account, AccountId, AccountWithMetadata},
+    compute_digest_for_path,
+    encryption::Ciphertext,
+    program::{DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, validate_execution},
+};
+use risc0_zkvm::{guest::env, serde::to_vec};
+
+fn main() {
+    let PrivacyPreservingCircuitInput {
+        program_outputs,
+        visibility_mask,
+        private_account_nonces,
+        private_account_keys,
+        private_account_nsks,
+        private_account_membership_proofs,
+        mut program_id,
+    } = env::read();
+
+    let mut pre_states: Vec<AccountWithMetadata> = Vec::new();
+    let mut state_diff: HashMap<AccountId, Account> = HashMap::new();
+
+    let num_calls = program_outputs.len();
+    if num_calls > MAX_NUMBER_CHAINED_CALLS {
+        panic!("Max chained calls depth is exceeded");
+    }
+
+    let Some(last_program_call) = program_outputs.last() else {
+        panic!("Program outputs is empty")
+    };
+
+    if !last_program_call.chained_calls.is_empty() {
+        panic!("Call stack is incomplete");
+    }
+
+    for window in program_outputs.windows(2) {
+        let caller = &window[0];
+        let callee = &window[1];
+
+        if caller.chained_calls.len() > 1 {
+            panic!("Privacy Multi-chained calls are not supported yet");
+        }
+
+        // TODO: Modify when multi-chain calls are supported in the circuit
+        let Some(caller_chained_call) = &caller.chained_calls.first() else {
+            panic!("Expected chained call");
+        };
+
+        // Check that instruction data in caller is the instruction data in callee
+        if caller_chained_call.instruction_data != callee.instruction_data {
+            panic!("Invalid instruction data");
+        }
+
+        // Check that account pre_states in caller are the ones in calle
+        if caller_chained_call.pre_states != callee.pre_states {
+            panic!("Invalid pre states");
+        }
+    }
+
+    for (i, program_output) in program_outputs.iter().enumerate() {
+        let mut program_output = program_output.clone();
+
+        // Check that `program_output` is consistent with the execution of the corresponding
+        // program.
+        let program_output_words =
+            &to_vec(&program_output).expect("program_output must be serializable");
+        env::verify(program_id, program_output_words)
+            .expect("program output must match the program's execution");
+
+        // Check that the program is well behaved.
+        // See the # Programs section for the definition of the `validate_execution` method.
+        if !validate_execution(
+            &program_output.pre_states,
+            &program_output.post_states,
+            program_id,
+        ) {
+            panic!("Bad behaved program");
+        }
+
+        // The invoked program claims the accounts with default program id.
+        for post in program_output
+            .post_states
+            .iter_mut()
+            .filter(|post| post.requires_claim())
+        {
+            // The invoked program can only claim accounts with default program id.
+            if post.account().program_owner == DEFAULT_PROGRAM_ID {
+                post.account_mut().program_owner = program_id;
+            } else {
+                panic!("Cannot claim an initialized account")
+            }
+        }
+
+        for (pre, post) in program_output
+            .pre_states
+            .iter()
+            .zip(&program_output.post_states)
+        {
+            if let Some(account_pre) = state_diff.get(&pre.account_id) {
+                if account_pre != &pre.account {
+                    panic!("Invalid input");
+                }
+            } else {
+                pre_states.push(pre.clone());
+            }
+            state_diff.insert(pre.account_id, post.account().clone());
+        }
+
+        // TODO: Modify when multi-chain calls are supported in the circuit
+        if let Some(next_chained_call) = &program_output.chained_calls.first() {
+            program_id = next_chained_call.program_id;
+        } else if i != program_outputs.len() - 1 {
+            panic!("Inner call without a chained call found")
+        };
+    }
+
+    let n_accounts = pre_states.len();
+    if visibility_mask.len() != n_accounts {
+        panic!("Invalid visibility mask length");
+    }
+
+    // These lists will be the public outputs of this circuit
+    // 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 ciphertexts: Vec<Ciphertext> = Vec::new();
+    let mut new_commitments: Vec<Commitment> = Vec::new();
+    let mut new_nullifiers: Vec<(Nullifier, CommitmentSetDigest)> = Vec::new();
+
+    let mut private_nonces_iter = private_account_nonces.iter();
+    let mut private_keys_iter = private_account_keys.iter();
+    let mut private_nsks_iter = private_account_nsks.iter();
+    let mut private_membership_proofs_iter = private_account_membership_proofs.iter();
+
+    let mut output_index = 0;
+    for i in 0..n_accounts {
+        match visibility_mask[i] {
+            0 => {
+                // Public account
+                public_pre_states.push(pre_states[i].clone());
+
+                let mut post = state_diff.get(&pre_states[i].account_id).unwrap().clone();
+
+                if post.program_owner == DEFAULT_PROGRAM_ID {
+                    // Claim account
+                    post.program_owner = program_id;
+                }
+                public_post_states.push(post);
+            }
+            1 | 2 => {
+                let new_nonce = private_nonces_iter.next().expect("Missing private nonce");
+                let (npk, shared_secret) = private_keys_iter.next().expect("Missing keys");
+
+                if AccountId::from(npk) != pre_states[i].account_id {
+                    panic!("AccountId mismatch");
+                }
+
+                if visibility_mask[i] == 1 {
+                    // Private account with authentication
+                    let nsk = private_nsks_iter.next().expect("Missing nsk");
+
+                    // Verify the nullifier public key
+                    let expected_npk = NullifierPublicKey::from(nsk);
+                    if &expected_npk != npk {
+                        panic!("Nullifier public key mismatch");
+                    }
+
+                    // Check pre_state authorization
+                    if !pre_states[i].is_authorized {
+                        panic!("Pre-state not authorized");
+                    }
+
+                    let membership_proof_opt = private_membership_proofs_iter
+                        .next()
+                        .expect("Missing membership proof");
+                    let (nullifier, set_digest) = membership_proof_opt
+                        .as_ref()
+                        .map(|membership_proof| {
+                            // Compute commitment set digest associated with provided auth path
+                            let commitment_pre = Commitment::new(npk, &pre_states[i].account);
+                            let set_digest =
+                                compute_digest_for_path(&commitment_pre, membership_proof);
+
+                            // Compute update nullifier
+                            let nullifier = Nullifier::for_account_update(&commitment_pre, nsk);
+                            (nullifier, set_digest)
+                        })
+                        .unwrap_or_else(|| {
+                            if pre_states[i].account != Account::default() {
+                                panic!("Found new private account with non default values.");
+                            }
+
+                            // Compute initialization nullifier
+                            let nullifier = Nullifier::for_account_initialization(npk);
+                            (nullifier, DUMMY_COMMITMENT_HASH)
+                        });
+                    new_nullifiers.push((nullifier, set_digest));
+                } else {
+                    // Private account without authentication
+                    if pre_states[i].account != Account::default() {
+                        panic!("Found new private account with non default values.");
+                    }
+
+                    if pre_states[i].is_authorized {
+                        panic!("Found new private account marked as authorized.");
+                    }
+
+                    let membership_proof_opt = private_membership_proofs_iter
+                        .next()
+                        .expect("Missing membership proof");
+                    assert!(
+                        membership_proof_opt.is_none(),
+                        "Membership proof must be None for unauthorized accounts"
+                    );
+                    let nullifier = Nullifier::for_account_initialization(npk);
+                    new_nullifiers.push((nullifier, DUMMY_COMMITMENT_HASH));
+                }
+
+                // Update post-state with new nonce
+                let mut post_with_updated_values =
+                    state_diff.get(&pre_states[i].account_id).unwrap().clone();
+                post_with_updated_values.nonce = *new_nonce;
+
+                if post_with_updated_values.program_owner == DEFAULT_PROGRAM_ID {
+                    // Claim account
+                    post_with_updated_values.program_owner = program_id;
+                }
+
+                // Compute commitment
+                let commitment_post = Commitment::new(npk, &post_with_updated_values);
+
+                // Encrypt and push post state
+                let encrypted_account = EncryptionScheme::encrypt(
+                    &post_with_updated_values,
+                    shared_secret,
+                    &commitment_post,
+                    output_index,
+                );
+
+                new_commitments.push(commitment_post);
+                ciphertexts.push(encrypted_account);
+                output_index += 1;
+            }
+            _ => panic!("Invalid visibility mask value"),
+        }
+    }
+
+    if private_nonces_iter.next().is_some() {
+        panic!("Too many nonces");
+    }
+
+    if private_keys_iter.next().is_some() {
+        panic!("Too many private account keys");
+    }
+
+    if private_nsks_iter.next().is_some() {
+        panic!("Too many private account authentication keys");
+    }
+
+    if private_membership_proofs_iter.next().is_some() {
+        panic!("Too many private account membership proofs");
+    }
+
+    let output = PrivacyPreservingCircuitOutput {
+        public_pre_states,
+        public_post_states,
+        ciphertexts,
+        new_commitments,
+        new_nullifiers,
+    };
+
+    env::commit(&output);
+}

program_methods/src/lib.rs

@@ -0,0 +1 @@
+include!(concat!(env!("OUT_DIR"), "/methods.rs"));