Merge pull request #202 from logos-blockchain/schouhy/implement-pda-for-public-accounts

Implement PDA for public executions
This commit is contained in:
Sergio Chouhy 2025-12-05 15:15:37 -03:00 committed by GitHub
commit 3e303f2c95
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 384 additions and 63 deletions

View File

@ -12,7 +12,7 @@ chacha20 = { version = "0.9", default-features = false }
k256 = { version = "0.13.3", optional = true }
base58 = { version = "0.2.0", optional = true }
anyhow = { version = "1.0.98", optional = true }
borsh.workspace = true
borsh = "1.5.7"
[features]
default = []

View File

@ -1,6 +1,8 @@
use risc0_zkvm::{DeserializeOwned, guest::env, serde::Deserializer};
use serde::{Deserialize, Serialize};
#[cfg(feature = "host")]
use crate::account::AccountId;
use crate::account::{Account, AccountWithMetadata};
pub type ProgramId = [u32; 8];
@ -12,12 +14,50 @@ pub struct ProgramInput<T> {
pub instruction: T,
}
/// A 32-byte seed used to compute a *Program-Derived AccountId* (PDA).
///
/// Each program can derive up to `2^256` unique account IDs by choosing different
/// seeds. PDAs allow programs to control namespaced account identifiers without
/// collisions between programs.
#[derive(Serialize, Deserialize, Clone)]
#[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))]
pub struct PdaSeed([u8; 32]);
impl PdaSeed {
pub fn new(value: [u8; 32]) -> Self {
Self(value)
}
}
#[cfg(feature = "host")]
impl From<(&ProgramId, &PdaSeed)> for AccountId {
fn from(value: (&ProgramId, &PdaSeed)) -> Self {
use risc0_zkvm::sha::{Impl, Sha256};
const PROGRAM_DERIVED_ACCOUNT_ID_PREFIX: &[u8; 32] =
b"/NSSA/v0.2/AccountId/PDA/\x00\x00\x00\x00\x00\x00\x00";
let mut bytes = [0; 96];
bytes[0..32].copy_from_slice(PROGRAM_DERIVED_ACCOUNT_ID_PREFIX);
let program_id_bytes: &[u8] =
bytemuck::try_cast_slice(value.0).expect("ProgramId should be castable to &[u8]");
bytes[32..64].copy_from_slice(program_id_bytes);
bytes[64..].copy_from_slice(&value.1.0);
AccountId::new(
Impl::hash_bytes(&bytes)
.as_bytes()
.try_into()
.expect("Hash output must be exactly 32 bytes long"),
)
}
}
#[derive(Serialize, Deserialize, Clone)]
#[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))]
pub struct ChainedCall {
pub program_id: ProgramId,
pub instruction_data: InstructionData,
pub pre_states: Vec<AccountWithMetadata>,
pub pda_seeds: Vec<PdaSeed>,
}
/// Represents the final state of an `Account` after a program execution.

View File

@ -0,0 +1,103 @@
use nssa_core::program::{
read_nssa_inputs, write_nssa_outputs_with_chained_call, AccountPostState, ChainedCall, PdaSeed, ProgramInput
};
use risc0_zkvm::serde::to_vec;
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 three accounts: [pinata_definition, pinata_token_holding, winner_token_holding]
let ProgramInput {
pre_states,
instruction: solution,
} = 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().to_vec();
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(
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,
);
}

View File

@ -84,25 +84,25 @@ impl TokenHolding {
fn parse(data: &[u8]) -> Option<Self> {
if data.len() != TOKEN_HOLDING_DATA_SIZE || data[0] != TOKEN_HOLDING_TYPE {
None
} else {
let account_type = data[0];
let definition_id = AccountId::new(
data[1..33]
.try_into()
.expect("Defintion ID must be 32 bytes long"),
);
let balance = u128::from_le_bytes(
data[33..]
.try_into()
.expect("balance must be 16 bytes little-endian"),
);
Some(Self {
definition_id,
balance,
account_type,
})
return None;
}
let account_type = data[0];
let definition_id = AccountId::new(
data[1..33]
.try_into()
.expect("Defintion ID must be 32 bytes long"),
);
let balance = u128::from_le_bytes(
data[33..]
.try_into()
.expect("balance must be 16 bytes little-endian"),
);
Some(Self {
definition_id,
balance,
account_type,
})
}
fn into_data(self) -> Data {
@ -223,7 +223,7 @@ fn initialize_account(pre_states: &[AccountWithMetadata]) -> Vec<AccountPostStat
panic!("Only uninitialized accounts can be initialized");
}
// TODO: We should check that this is an account owned by the token program.
// TODO: #212 We should check that this is an account owned by the token program.
// This check can't be done here since the ID of the program is known only after compiling it
//
// Check definition account is valid
@ -249,7 +249,7 @@ fn main() {
instruction,
} = read_nssa_inputs::<Instruction>();
let (pre_states, post_states) = match instruction[0] {
let post_states = match instruction[0] {
0 => {
// Parse instruction
let total_supply = u128::from_le_bytes(
@ -263,8 +263,7 @@ fn main() {
assert_ne!(name, [0; 6]);
// Execute
let post_states = new_definition(&pre_states, name, total_supply);
(pre_states, post_states)
new_definition(&pre_states, name, total_supply)
}
1 => {
// Parse instruction
@ -279,14 +278,14 @@ fn main() {
assert_eq!(name, [0; 6]);
// Execute
let post_states = transfer(&pre_states, balance_to_move);
(pre_states, post_states)
transfer(&pre_states, balance_to_move)
}
2 => {
// Initialize account
assert_eq!(instruction[1..], [0; 22]);
let post_states = initialize_account(&pre_states);
(pre_states, post_states)
if instruction[1..] != [0; 22] {
panic!("Invalid instruction for initialize account");
}
initialize_account(&pre_states)
}
_ => panic!("Invalid instruction"),
};
@ -655,7 +654,7 @@ mod tests {
AccountWithMetadata {
account: Account {
// Definition ID with
data: vec![0; TOKEN_DEFINITION_DATA_SIZE - 16]
data: [0; TOKEN_DEFINITION_DATA_SIZE - 16]
.into_iter()
.chain(u128::to_le_bytes(1000))
.collect(),

View File

@ -104,6 +104,11 @@ impl Program {
// `program_methods`
Self::new(PINATA_ELF.to_vec()).unwrap()
}
pub fn pinata_token() -> Self {
use crate::program_methods::PINATA_TOKEN_ELF;
Self::new(PINATA_TOKEN_ELF.to_vec()).expect("Piñata program must be a valid R0BF file")
}
}
#[cfg(test)]

View File

@ -3,7 +3,7 @@ use std::collections::{HashMap, HashSet, VecDeque};
use borsh::{BorshDeserialize, BorshSerialize};
use nssa_core::{
account::{Account, AccountId, AccountWithMetadata},
program::{ChainedCall, DEFAULT_PROGRAM_ID, validate_execution},
program::{ChainedCall, DEFAULT_PROGRAM_ID, PdaSeed, ProgramId, validate_execution},
};
use sha2::{Digest, digest::FixedOutput};
@ -107,12 +107,13 @@ impl PublicTransaction {
program_id: message.program_id,
instruction_data: message.instruction_data.clone(),
pre_states: input_pre_states,
pda_seeds: vec![],
};
let mut chained_calls = VecDeque::from_iter([initial_call]);
let mut chained_calls = VecDeque::from_iter([(initial_call, None)]);
let mut chain_calls_counter = 0;
while let Some(chained_call) = chained_calls.pop_front() {
while let Some((chained_call, caller_program_id)) = chained_calls.pop_front() {
if chain_calls_counter > MAX_NUMBER_CHAINED_CALLS {
return Err(NssaError::MaxChainedCallsDepthExceeded);
}
@ -125,6 +126,9 @@ impl PublicTransaction {
let mut program_output =
program.execute(&chained_call.pre_states, &chained_call.instruction_data)?;
let authorized_pdas =
self.compute_authorized_pdas(&caller_program_id, &chained_call.pda_seeds);
for pre in &program_output.pre_states {
let account_id = pre.account_id;
// Check that the program output pre_states coinicide with the values in the public
@ -137,8 +141,11 @@ impl PublicTransaction {
return Err(NssaError::InvalidProgramBehavior);
}
// Check that authorization flags are consistent with the provided ones
if pre.is_authorized && !signer_account_ids.contains(&account_id) {
// Check that authorization flags are consistent with the provided ones or
// authorized by program through the PDA mechanism
let is_authorized = signer_account_ids.contains(&account_id)
|| authorized_pdas.contains(&account_id);
if pre.is_authorized != is_authorized {
return Err(NssaError::InvalidProgramBehavior);
}
}
@ -176,7 +183,7 @@ impl PublicTransaction {
}
for new_call in program_output.chained_calls.into_iter().rev() {
chained_calls.push_front(new_call);
chained_calls.push_front((new_call, Some(chained_call.program_id)));
}
chain_calls_counter += 1;
@ -184,6 +191,21 @@ impl PublicTransaction {
Ok(state_diff)
}
fn compute_authorized_pdas(
&self,
caller_program_id: &Option<ProgramId>,
pda_seeds: &[PdaSeed],
) -> HashSet<AccountId> {
if let Some(caller_program_id) = caller_program_id {
pda_seeds
.iter()
.map(|pda_seed| AccountId::from((caller_program_id, pda_seed)))
.collect()
} else {
HashSet::new()
}
}
}
#[cfg(test)]

View File

@ -239,6 +239,20 @@ impl V02State {
},
);
}
pub fn add_pinata_token_program(&mut self, account_id: AccountId) {
self.insert_program(Program::pinata_token());
self.public_state.insert(
account_id,
Account {
program_owner: Program::pinata_token().id(),
// Difficulty: 3
data: vec![3; 33],
..Account::default()
},
);
}
}
#[cfg(test)]
@ -250,7 +264,7 @@ pub mod tests {
Commitment, Nullifier, NullifierPublicKey, NullifierSecretKey, SharedSecretKey,
account::{Account, AccountId, AccountWithMetadata, Nonce},
encryption::{EphemeralPublicKey, IncomingViewingPublicKey, Scalar},
program::ProgramId,
program::{PdaSeed, ProgramId},
};
use crate::{
@ -2093,14 +2107,18 @@ pub mod tests {
let key = PrivateKey::try_new([1; 32]).unwrap();
let from = AccountId::from(&PublicKey::new_from_private_key(&key));
let to = AccountId::new([2; 32]);
let initial_balance = 100;
let initial_balance = 1000;
let initial_data = [(from, initial_balance), (to, 0)];
let mut state =
V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
let from_key = key;
let amount: u128 = 0;
let instruction: (u128, ProgramId, u32) =
(amount, Program::authenticated_transfer_program().id(), 2);
let amount: u128 = 37;
let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
amount,
Program::authenticated_transfer_program().id(),
2,
None,
);
let expected_to_post = Account {
program_owner: Program::authenticated_transfer_program().id(),
@ -2140,10 +2158,11 @@ pub mod tests {
V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
let from_key = key;
let amount: u128 = 0;
let instruction: (u128, ProgramId, u32) = (
let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
amount,
Program::authenticated_transfer_program().id(),
MAX_NUMBER_CHAINED_CALLS as u32 + 1,
None,
);
let message = public_transaction::Message::try_new(
@ -2164,6 +2183,48 @@ pub mod tests {
));
}
#[test]
fn test_execution_that_requires_authentication_of_a_program_derived_account_id_succeeds() {
let chain_caller = Program::chain_caller();
let pda_seed = PdaSeed::new([37; 32]);
let from = AccountId::from((&chain_caller.id(), &pda_seed));
let to = AccountId::new([2; 32]);
let initial_balance = 1000;
let initial_data = [(from, initial_balance), (to, 0)];
let mut state =
V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
let amount: u128 = 58;
let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
amount,
Program::authenticated_transfer_program().id(),
1,
Some(pda_seed),
);
let expected_to_post = Account {
program_owner: Program::authenticated_transfer_program().id(),
balance: amount, // The `chain_caller` chains the program twice
..Account::default()
};
let message = public_transaction::Message::try_new(
chain_caller.id(),
vec![to, from], // The chain_caller program permutes the account order in the chain
// call
vec![],
instruction,
)
.unwrap();
let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
let tx = PublicTransaction::new(message, witness_set);
state.transition_from_public_transaction(&tx).unwrap();
let from_post = state.get_account_by_id(&from);
let to_post = state.get_account_by_id(&to);
assert_eq!(from_post.balance, initial_balance - amount);
assert_eq!(to_post, expected_to_post);
}
#[test]
fn test_claiming_mechanism_within_chain_call() {
// This test calls the authenticated transfer program through the chain_caller program.
@ -2195,8 +2256,12 @@ pub mod tests {
// The transaction executes the chain_caller program, which internally calls the
// authenticated_transfer program
let instruction: (u128, ProgramId, u32) =
(amount, Program::authenticated_transfer_program().id(), 1);
let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
amount,
Program::authenticated_transfer_program().id(),
1,
None,
);
let message = public_transaction::Message::try_new(
chain_caller.id(),
vec![to, from], // The chain_caller program permutes the account order in the chain
@ -2216,6 +2281,86 @@ pub mod tests {
assert_eq!(to_post, expected_to_post);
}
#[test]
fn test_pda_mechanism_with_pinata_token_program() {
let pinata_token = Program::pinata_token();
let token = Program::token();
let pinata_definition_id = AccountId::new([1; 32]);
let pinata_token_definition_id = AccountId::new([2; 32]);
// Total supply of pinata token will be in an account under a PDA.
let pinata_token_holding_id = AccountId::from((&pinata_token.id(), &PdaSeed::new([0; 32])));
let winner_token_holding_id = AccountId::new([3; 32]);
let mut expected_winner_account_data = [0; 49];
expected_winner_account_data[0] = 1;
expected_winner_account_data[1..33].copy_from_slice(pinata_token_definition_id.value());
expected_winner_account_data[33..].copy_from_slice(&150u128.to_le_bytes());
let expected_winner_token_holding_post = Account {
program_owner: token.id(),
data: expected_winner_account_data.to_vec(),
..Account::default()
};
let mut state = V02State::new_with_genesis_accounts(&[], &[]);
state.add_pinata_token_program(pinata_definition_id);
// Execution of the token program to create new token for the pinata token
// definition and supply accounts
let total_supply: u128 = 10_000_000;
// instruction: [0x00 || total_supply (little-endian 16 bytes) || name (6 bytes)]
let mut instruction: [u8; 23] = [0; 23];
instruction[1..17].copy_from_slice(&total_supply.to_le_bytes());
instruction[17..].copy_from_slice(b"PINATA");
let message = public_transaction::Message::try_new(
token.id(),
vec![pinata_token_definition_id, pinata_token_holding_id],
vec![],
instruction,
)
.unwrap();
let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
let tx = PublicTransaction::new(message, witness_set);
state.transition_from_public_transaction(&tx).unwrap();
// Execution of the token program transfer just to initialize the winner token account
let mut instruction: [u8; 23] = [0; 23];
instruction[0] = 2;
let message = public_transaction::Message::try_new(
token.id(),
vec![pinata_token_definition_id, winner_token_holding_id],
vec![],
instruction,
)
.unwrap();
let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
let tx = PublicTransaction::new(message, witness_set);
state.transition_from_public_transaction(&tx).unwrap();
// Submit a solution to the pinata program to claim the prize
let solution: u128 = 989106;
let message = public_transaction::Message::try_new(
pinata_token.id(),
vec![
pinata_definition_id,
pinata_token_holding_id,
winner_token_holding_id,
],
vec![],
solution,
)
.unwrap();
let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
let tx = PublicTransaction::new(message, witness_set);
state.transition_from_public_transaction(&tx).unwrap();
let winner_token_holding_post = state.get_account_by_id(&winner_token_holding_id);
assert_eq!(
winner_token_holding_post,
expected_winner_token_holding_post
);
}
#[test]
fn test_claiming_mechanism_cannot_claim_initialied_accounts() {
let claimer = Program::claimer();

View File

@ -1,47 +1,54 @@
use nssa_core::program::{
AccountPostState, ChainedCall, ProgramId, ProgramInput, read_nssa_inputs,
AccountPostState, ChainedCall, PdaSeed, ProgramId, ProgramInput, read_nssa_inputs,
write_nssa_outputs_with_chained_call,
};
use risc0_zkvm::serde::to_vec;
type Instruction = (u128, ProgramId, u32);
type Instruction = (u128, ProgramId, u32, Option<PdaSeed>);
/// A program that calls another program `num_chain_calls` times.
/// It permutes the order of the input accounts on the subsequent call
/// The `ProgramId` in the instruction must be the program_id of the authenticated transfers program
fn main() {
let ProgramInput {
pre_states,
instruction: (balance, program_id, num_chain_calls),
instruction: (balance, auth_transfer_id, num_chain_calls, pda_seed),
} = read_nssa_inputs::<Instruction>();
let [sender_pre, receiver_pre] = match pre_states.try_into() {
let [recipient_pre, sender_pre] = match pre_states.try_into() {
Ok(array) => array,
Err(_) => return,
};
let instruction_data = to_vec(&balance).unwrap();
let mut chained_call = vec![
ChainedCall {
program_id,
instruction_data: instruction_data.clone(),
pre_states: vec![receiver_pre.clone(), sender_pre.clone()], // <- Account order permutation here
};
num_chain_calls as usize - 1
];
let mut running_recipient_pre = recipient_pre.clone();
let mut running_sender_pre = sender_pre.clone();
chained_call.push(ChainedCall {
program_id,
instruction_data,
pre_states: vec![receiver_pre.clone(), sender_pre.clone()], // <- Account order permutation here
});
if pda_seed.is_some() {
running_sender_pre.is_authorized = true;
}
let mut chained_calls = Vec::new();
for _i in 0..num_chain_calls {
let new_chained_call = ChainedCall {
program_id: auth_transfer_id,
instruction_data: instruction_data.clone(),
pre_states: vec![running_sender_pre.clone(), running_recipient_pre.clone()], // <- Account order permutation here
pda_seeds: pda_seed.iter().cloned().collect(),
};
chained_calls.push(new_chained_call);
running_sender_pre.account.balance -= balance;
running_recipient_pre.account.balance += balance;
}
write_nssa_outputs_with_chained_call(
vec![sender_pre.clone(), receiver_pre.clone()],
vec![sender_pre.clone(), recipient_pre.clone()],
vec![
AccountPostState::new(sender_pre.account),
AccountPostState::new(receiver_pre.account),
AccountPostState::new(recipient_pre.account),
],
chained_call,
chained_calls,
);
}