diff --git a/Cargo.lock b/Cargo.lock index cee0396813b73f..ed2f8adec1152a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7082,6 +7082,17 @@ dependencies = [ "solana-zk-token-sdk 1.16.0", ] +[[package]] +name = "solana-zk-token-proof-program-tests" +version = "1.16.0" +dependencies = [ + "bytemuck", + "solana-program-runtime", + "solana-program-test", + "solana-sdk 1.16.0", + "solana-zk-token-sdk 1.16.0", +] + [[package]] name = "solana-zk-token-sdk" version = "1.15.1" @@ -7236,9 +7247,9 @@ dependencies = [ [[package]] name = "spl-token-2022" -version = "0.6.0" +version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67fcd758e8d22c5fce17315015f5ff319604d1a6e57a73c72795639dba898890" +checksum = "0043b590232c400bad5ee9eb983ced003d15163c4c5d56b090ac6d9a57457b47" dependencies = [ "arrayref", "bytemuck", diff --git a/Cargo.toml b/Cargo.toml index 5fac23fd023095..6ac6978c7d43ed 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -62,6 +62,7 @@ members = [ "programs/stake", "programs/vote", "programs/zk-token-proof", + "programs/zk-token-proof-tests", "pubsub-client", "quic-client", "rayon-threadlimit", @@ -360,7 +361,7 @@ spl-associated-token-account = "=1.1.3" spl-instruction-padding = "0.1" spl-memo = "=3.0.1" spl-token = "=3.5.0" -spl-token-2022 = "=0.6.0" +spl-token-2022 = "=0.6.1" static_assertions = "1.1.0" stream-cancel = "0.8.1" strum = "0.24" diff --git a/programs/sbf/Cargo.lock b/programs/sbf/Cargo.lock index 8d33f7c359bad5..cf00794a9c3805 100644 --- a/programs/sbf/Cargo.lock +++ b/programs/sbf/Cargo.lock @@ -6406,9 +6406,9 @@ dependencies = [ [[package]] name = "spl-token-2022" -version = "0.6.0" +version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67fcd758e8d22c5fce17315015f5ff319604d1a6e57a73c72795639dba898890" +checksum = "0043b590232c400bad5ee9eb983ced003d15163c4c5d56b090ac6d9a57457b47" dependencies = [ "arrayref", "bytemuck", diff --git a/programs/zk-token-proof-tests/Cargo.toml b/programs/zk-token-proof-tests/Cargo.toml new file mode 100644 index 00000000000000..8ddb5656a73b87 --- /dev/null +++ b/programs/zk-token-proof-tests/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "solana-zk-token-proof-program-tests" +authors = ["Solana Labs Maintainers "] +repository = "https://github.com/solana-labs/solana" +version = "1.16.0" +license = "Apache-2.0" +edition = "2021" +publish = false + +[dev-dependencies] +bytemuck = { version = "1.11.0", features = ["derive"] } +solana-program-runtime = { path = "../../program-runtime", version = "=1.16.0" } +solana-program-test = { path = "../../program-test", version = "=1.16.0" } +solana-sdk = { path = "../../sdk", version = "=1.16.0" } +solana-zk-token-sdk = { path = "../../zk-token-sdk", version = "=1.16.0" } diff --git a/programs/zk-token-proof-tests/tests/process_transaction.rs b/programs/zk-token-proof-tests/tests/process_transaction.rs new file mode 100644 index 00000000000000..555295e4e9a6ed --- /dev/null +++ b/programs/zk-token-proof-tests/tests/process_transaction.rs @@ -0,0 +1,788 @@ +use { + bytemuck::Pod, + solana_program_test::*, + solana_sdk::{ + instruction::InstructionError, + signature::Signer, + signer::keypair::Keypair, + system_instruction, + transaction::{Transaction, TransactionError}, + }, + solana_zk_token_sdk::{ + encryption::elgamal::ElGamalKeypair, instruction::*, zk_token_proof_instruction::*, + zk_token_proof_program, zk_token_proof_state::ProofContextState, + }, + std::mem::size_of, +}; + +const VERIFY_INSTRUCTION_TYPES: [ProofInstruction; 6] = [ + ProofInstruction::VerifyCloseAccount, + ProofInstruction::VerifyWithdraw, + ProofInstruction::VerifyWithdrawWithheldTokens, + ProofInstruction::VerifyTransfer, + ProofInstruction::VerifyTransferWithFee, + ProofInstruction::VerifyPubkeyValidity, +]; + +#[tokio::test] +async fn test_close_account() { + let elgamal_keypair = ElGamalKeypair::new_rand(); + + let zero_ciphertext = elgamal_keypair.public.encrypt(0_u64); + let success_proof_data = CloseAccountData::new(&elgamal_keypair, &zero_ciphertext).unwrap(); + + let incorrect_keypair = ElGamalKeypair { + public: ElGamalKeypair::new_rand().public, + secret: ElGamalKeypair::new_rand().secret, + }; + let fail_proof_data = CloseAccountData::new(&incorrect_keypair, &zero_ciphertext).unwrap(); + + test_verify_proof_without_context( + ProofInstruction::VerifyCloseAccount, + &success_proof_data, + &fail_proof_data, + ) + .await; + + test_verify_proof_with_context( + ProofInstruction::VerifyCloseAccount, + size_of::>(), + &success_proof_data, + &fail_proof_data, + ) + .await; + + test_close_context_state( + ProofInstruction::VerifyCloseAccount, + size_of::>(), + &success_proof_data, + ) + .await; +} + +#[tokio::test] +async fn test_withdraw_withheld_tokens() { + let elgamal_keypair = ElGamalKeypair::new_rand(); + let destination_keypair = ElGamalKeypair::new_rand(); + + let amount: u64 = 0; + let withdraw_withheld_authority_ciphertext = elgamal_keypair.public.encrypt(amount); + + let success_proof_data = WithdrawWithheldTokensData::new( + &elgamal_keypair, + &destination_keypair.public, + &withdraw_withheld_authority_ciphertext, + amount, + ) + .unwrap(); + + let incorrect_keypair = ElGamalKeypair { + public: ElGamalKeypair::new_rand().public, + secret: ElGamalKeypair::new_rand().secret, + }; + let fail_proof_data = WithdrawWithheldTokensData::new( + &incorrect_keypair, + &destination_keypair.public, + &withdraw_withheld_authority_ciphertext, + amount, + ) + .unwrap(); + + test_verify_proof_without_context( + ProofInstruction::VerifyWithdrawWithheldTokens, + &success_proof_data, + &fail_proof_data, + ) + .await; + + test_verify_proof_with_context( + ProofInstruction::VerifyWithdrawWithheldTokens, + size_of::>(), + &success_proof_data, + &fail_proof_data, + ) + .await; + + test_close_context_state( + ProofInstruction::VerifyWithdrawWithheldTokens, + size_of::>(), + &success_proof_data, + ) + .await; +} + +#[tokio::test] +async fn test_transfer() { + let source_keypair = ElGamalKeypair::new_rand(); + let dest_pubkey = ElGamalKeypair::new_rand().public; + let auditor_pubkey = ElGamalKeypair::new_rand().public; + + let spendable_balance: u64 = 0; + let spendable_ciphertext = source_keypair.public.encrypt(spendable_balance); + + let transfer_amount: u64 = 0; + + let success_proof_data = TransferData::new( + transfer_amount, + (spendable_balance, &spendable_ciphertext), + &source_keypair, + (&dest_pubkey, &auditor_pubkey), + ) + .unwrap(); + + let incorrect_keypair = ElGamalKeypair { + public: ElGamalKeypair::new_rand().public, + secret: ElGamalKeypair::new_rand().secret, + }; + + let fail_proof_data = TransferData::new( + transfer_amount, + (spendable_balance, &spendable_ciphertext), + &incorrect_keypair, + (&dest_pubkey, &auditor_pubkey), + ) + .unwrap(); + + test_verify_proof_without_context( + ProofInstruction::VerifyTransfer, + &success_proof_data, + &fail_proof_data, + ) + .await; + + test_verify_proof_with_context( + ProofInstruction::VerifyTransfer, + size_of::>(), + &success_proof_data, + &fail_proof_data, + ) + .await; + + test_close_context_state( + ProofInstruction::VerifyTransfer, + size_of::>(), + &success_proof_data, + ) + .await; +} + +#[tokio::test] +async fn test_transfer_with_fee() { + let source_keypair = ElGamalKeypair::new_rand(); + let destination_pubkey = ElGamalKeypair::new_rand().public; + let auditor_pubkey = ElGamalKeypair::new_rand().public; + let withdraw_withheld_authority_pubkey = ElGamalKeypair::new_rand().public; + + let spendable_balance: u64 = 120; + let spendable_ciphertext = source_keypair.public.encrypt(spendable_balance); + + let transfer_amount: u64 = 0; + + let fee_parameters = FeeParameters { + fee_rate_basis_points: 400, + maximum_fee: 3, + }; + + let success_proof_data = TransferWithFeeData::new( + transfer_amount, + (spendable_balance, &spendable_ciphertext), + &source_keypair, + (&destination_pubkey, &auditor_pubkey), + fee_parameters, + &withdraw_withheld_authority_pubkey, + ) + .unwrap(); + + let incorrect_keypair = ElGamalKeypair { + public: ElGamalKeypair::new_rand().public, + secret: ElGamalKeypair::new_rand().secret, + }; + + let fail_proof_data = TransferWithFeeData::new( + transfer_amount, + (spendable_balance, &spendable_ciphertext), + &incorrect_keypair, + (&destination_pubkey, &auditor_pubkey), + fee_parameters, + &withdraw_withheld_authority_pubkey, + ) + .unwrap(); + + test_verify_proof_without_context( + ProofInstruction::VerifyTransferWithFee, + &success_proof_data, + &fail_proof_data, + ) + .await; + + test_verify_proof_with_context( + ProofInstruction::VerifyTransferWithFee, + size_of::>(), + &success_proof_data, + &fail_proof_data, + ) + .await; + + test_close_context_state( + ProofInstruction::VerifyTransferWithFee, + size_of::>(), + &success_proof_data, + ) + .await; +} + +#[tokio::test] +async fn test_withdraw() { + let elgamal_keypair = ElGamalKeypair::new_rand(); + + let current_balance: u64 = 77; + let current_ciphertext = elgamal_keypair.public.encrypt(current_balance); + let withdraw_amount: u64 = 55; + + let success_proof_data = WithdrawData::new( + withdraw_amount, + &elgamal_keypair, + current_balance, + ¤t_ciphertext, + ) + .unwrap(); + + let incorrect_keypair = ElGamalKeypair { + public: ElGamalKeypair::new_rand().public, + secret: ElGamalKeypair::new_rand().secret, + }; + let fail_proof_data = WithdrawData::new( + withdraw_amount, + &incorrect_keypair, + current_balance, + ¤t_ciphertext, + ) + .unwrap(); + + test_verify_proof_without_context( + ProofInstruction::VerifyWithdraw, + &success_proof_data, + &fail_proof_data, + ) + .await; + + test_verify_proof_with_context( + ProofInstruction::VerifyWithdraw, + size_of::>(), + &success_proof_data, + &fail_proof_data, + ) + .await; + + test_close_context_state( + ProofInstruction::VerifyWithdraw, + size_of::>(), + &success_proof_data, + ) + .await; +} + +#[tokio::test] +async fn test_pubkey_validity() { + let elgamal_keypair = ElGamalKeypair::new_rand(); + + let success_proof_data = PubkeyValidityData::new(&elgamal_keypair).unwrap(); + + let incorrect_keypair = ElGamalKeypair { + public: ElGamalKeypair::new_rand().public, + secret: ElGamalKeypair::new_rand().secret, + }; + + let fail_proof_data = PubkeyValidityData::new(&incorrect_keypair).unwrap(); + + test_verify_proof_without_context( + ProofInstruction::VerifyPubkeyValidity, + &success_proof_data, + &fail_proof_data, + ) + .await; + + test_verify_proof_with_context( + ProofInstruction::VerifyPubkeyValidity, + size_of::>(), + &success_proof_data, + &fail_proof_data, + ) + .await; + + test_close_context_state( + ProofInstruction::VerifyPubkeyValidity, + size_of::>(), + &success_proof_data, + ) + .await; +} + +async fn test_verify_proof_without_context( + proof_instruction: ProofInstruction, + success_proof_data: &T, + fail_proof_data: &T, +) where + T: Pod + ZkProofData, + U: Pod, +{ + let mut context = ProgramTest::default().start_with_context().await; + + let client = &mut context.banks_client; + let payer = &context.payer; + let recent_blockhash = context.last_blockhash; + + // verify a valid proof (wihtout creating a context account) + let instructions = vec![proof_instruction.encode_verify_proof(None, success_proof_data)]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer], + recent_blockhash, + ); + client.process_transaction(transaction).await.unwrap(); + + // try to verify an invalid proof (without creating a context account) + let instructions = vec![proof_instruction.encode_verify_proof(None, fail_proof_data)]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer], + recent_blockhash, + ); + let err = client + .process_transaction(transaction) + .await + .unwrap_err() + .unwrap(); + assert_eq!( + err, + TransactionError::InstructionError(0, InstructionError::InvalidInstructionData) + ); + + // try to verify a valid proof, but with a wrong proof type + for wrong_instruction_type in VERIFY_INSTRUCTION_TYPES { + if proof_instruction == wrong_instruction_type { + continue; + } + + let instruction = + vec![wrong_instruction_type.encode_verify_proof(None, success_proof_data)]; + let transaction = Transaction::new_signed_with_payer( + &instruction, + Some(&payer.pubkey()), + &[payer], + recent_blockhash, + ); + let err = client + .process_transaction(transaction) + .await + .unwrap_err() + .unwrap(); + assert_eq!( + err, + TransactionError::InstructionError(0, InstructionError::InvalidInstructionData) + ); + } +} + +async fn test_verify_proof_with_context( + instruction_type: ProofInstruction, + space: usize, + success_proof_data: &T, + fail_proof_data: &T, +) where + T: Pod + ZkProofData, + U: Pod, +{ + let mut context = ProgramTest::default().start_with_context().await; + let rent = context.banks_client.get_rent().await.unwrap(); + + let client = &mut context.banks_client; + let payer = &context.payer; + let recent_blockhash = context.last_blockhash; + + let context_state_account = Keypair::new(); + let context_state_authority = Keypair::new(); + + let context_state_info = ContextStateInfo { + context_state_account: &context_state_account.pubkey(), + context_state_authority: &context_state_authority.pubkey(), + }; + + // try to create proof context state with an invalid proof + let instructions = vec![ + system_instruction::create_account( + &payer.pubkey(), + &context_state_account.pubkey(), + rent.minimum_balance(space), + space as u64, + &zk_token_proof_program::id(), + ), + instruction_type.encode_verify_proof(Some(context_state_info), fail_proof_data), + ]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer, &context_state_account], + recent_blockhash, + ); + let err = client + .process_transaction(transaction) + .await + .unwrap_err() + .unwrap(); + assert_eq!( + err, + TransactionError::InstructionError(1, InstructionError::InvalidInstructionData) + ); + + // try to create proof context state with incorrect account data length + let instructions = vec![ + system_instruction::create_account( + &payer.pubkey(), + &context_state_account.pubkey(), + rent.minimum_balance(space), + (space.checked_sub(1).unwrap()) as u64, + &zk_token_proof_program::id(), + ), + instruction_type.encode_verify_proof(Some(context_state_info), success_proof_data), + ]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer, &context_state_account], + recent_blockhash, + ); + let err = client + .process_transaction(transaction) + .await + .unwrap_err() + .unwrap(); + assert_eq!( + err, + TransactionError::InstructionError(1, InstructionError::InvalidAccountData) + ); + + // try to create proof context state with insufficient rent + let instructions = vec![ + system_instruction::create_account( + &payer.pubkey(), + &context_state_account.pubkey(), + rent.minimum_balance(space).checked_sub(1).unwrap(), + space as u64, + &zk_token_proof_program::id(), + ), + instruction_type.encode_verify_proof(Some(context_state_info), success_proof_data), + ]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer, &context_state_account], + recent_blockhash, + ); + let err = client + .process_transaction(transaction) + .await + .unwrap_err() + .unwrap(); + assert_eq!( + err, + TransactionError::InsufficientFundsForRent { account_index: 1 }, + ); + + // try to create proof context state with an invalid `ProofType` + for wrong_instruction_type in VERIFY_INSTRUCTION_TYPES { + if instruction_type == wrong_instruction_type { + continue; + } + + let instructions = vec![ + system_instruction::create_account( + &payer.pubkey(), + &context_state_account.pubkey(), + rent.minimum_balance(space), + space as u64, + &zk_token_proof_program::id(), + ), + wrong_instruction_type + .encode_verify_proof(Some(context_state_info), success_proof_data), + ]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer, &context_state_account], + recent_blockhash, + ); + let err = client + .process_transaction(transaction) + .await + .unwrap_err() + .unwrap(); + assert_eq!( + err, + TransactionError::InstructionError(1, InstructionError::InvalidInstructionData) + ); + } + + // successfully create a proof context state + let instructions = vec![ + system_instruction::create_account( + &payer.pubkey(), + &context_state_account.pubkey(), + rent.minimum_balance(space), + space as u64, + &zk_token_proof_program::id(), + ), + instruction_type.encode_verify_proof(Some(context_state_info), success_proof_data), + ]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer, &context_state_account], + recent_blockhash, + ); + client.process_transaction(transaction).await.unwrap(); + + // try overwriting the context state + let instructions = + vec![instruction_type.encode_verify_proof(Some(context_state_info), success_proof_data)]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer], + recent_blockhash, + ); + let err = client + .process_transaction(transaction) + .await + .unwrap_err() + .unwrap(); + assert_eq!( + err, + TransactionError::InstructionError(0, InstructionError::AccountAlreadyInitialized) + ); + + // self-owned context state account + let context_state_account_and_authority = Keypair::new(); + let context_state_info = ContextStateInfo { + context_state_account: &context_state_account_and_authority.pubkey(), + context_state_authority: &context_state_account_and_authority.pubkey(), + }; + + let instructions = vec![ + system_instruction::create_account( + &payer.pubkey(), + &context_state_account_and_authority.pubkey(), + rent.minimum_balance(space), + space as u64, + &zk_token_proof_program::id(), + ), + instruction_type.encode_verify_proof(Some(context_state_info), success_proof_data), + ]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer, &context_state_account_and_authority], + recent_blockhash, + ); + client.process_transaction(transaction).await.unwrap(); +} + +async fn test_close_context_state( + instruction_type: ProofInstruction, + space: usize, + success_proof_data: &T, +) where + T: Pod + ZkProofData, + U: Pod, +{ + let mut context = ProgramTest::default().start_with_context().await; + let rent = context.banks_client.get_rent().await.unwrap(); + + let client = &mut context.banks_client; + let payer = &context.payer; + let recent_blockhash = context.last_blockhash; + + let context_state_account = Keypair::new(); + let context_state_authority = Keypair::new(); + + let context_state_info = ContextStateInfo { + context_state_account: &context_state_account.pubkey(), + context_state_authority: &context_state_authority.pubkey(), + }; + + let destination_account = Keypair::new(); + + // create a proof context state + let instructions = vec![ + system_instruction::create_account( + &payer.pubkey(), + &context_state_account.pubkey(), + rent.minimum_balance(space), + space as u64, + &zk_token_proof_program::id(), + ), + instruction_type.encode_verify_proof(Some(context_state_info), success_proof_data), + ]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer, &context_state_account], + recent_blockhash, + ); + client.process_transaction(transaction).await.unwrap(); + + // try to close context state with incorrect authority + let incorrect_authority = Keypair::new(); + let instruction = close_context_state( + ContextStateInfo { + context_state_account: &context_state_account.pubkey(), + context_state_authority: &incorrect_authority.pubkey(), + }, + &destination_account.pubkey(), + ); + let transaction = Transaction::new_signed_with_payer( + &[instruction], + Some(&payer.pubkey()), + &[payer, &incorrect_authority], + recent_blockhash, + ); + let err = client + .process_transaction(transaction) + .await + .unwrap_err() + .unwrap(); + assert_eq!( + err, + TransactionError::InstructionError(0, InstructionError::InvalidAccountOwner) + ); + + // successfully close proof context state + let instruction = close_context_state( + ContextStateInfo { + context_state_account: &context_state_account.pubkey(), + context_state_authority: &context_state_authority.pubkey(), + }, + &destination_account.pubkey(), + ); + let transaction = Transaction::new_signed_with_payer( + &[instruction.clone()], + Some(&payer.pubkey()), + &[payer, &context_state_authority], + recent_blockhash, + ); + client.process_transaction(transaction).await.unwrap(); + + // create and close proof context in a single transaction + let instructions = vec![ + system_instruction::create_account( + &payer.pubkey(), + &context_state_account.pubkey(), + 0_u64, // do not deposit rent + space as u64, + &zk_token_proof_program::id(), + ), + instruction_type.encode_verify_proof(Some(context_state_info), success_proof_data), + close_context_state( + ContextStateInfo { + context_state_account: &context_state_account.pubkey(), + context_state_authority: &context_state_authority.pubkey(), + }, + &destination_account.pubkey(), + ), + ]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer, &context_state_account, &context_state_authority], + recent_blockhash, + ); + client.process_transaction(transaction).await.unwrap(); + + // close proof context state with owner as destination + let instructions = vec![ + system_instruction::create_account( + &payer.pubkey(), + &context_state_account.pubkey(), + 0_u64, + space as u64, + &zk_token_proof_program::id(), + ), + instruction_type.encode_verify_proof(Some(context_state_info), success_proof_data), + close_context_state( + ContextStateInfo { + context_state_account: &context_state_account.pubkey(), + context_state_authority: &context_state_authority.pubkey(), + }, + &context_state_authority.pubkey(), + ), + ]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer, &context_state_account, &context_state_authority], + recent_blockhash, + ); + client.process_transaction(transaction).await.unwrap(); + + // try close account with itself as destination + let instructions = vec![ + system_instruction::create_account( + &payer.pubkey(), + &context_state_account.pubkey(), + 0_u64, + space as u64, + &zk_token_proof_program::id(), + ), + instruction_type.encode_verify_proof(Some(context_state_info), success_proof_data), + close_context_state( + ContextStateInfo { + context_state_account: &context_state_account.pubkey(), + context_state_authority: &context_state_authority.pubkey(), + }, + &context_state_account.pubkey(), + ), + ]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer, &context_state_account, &context_state_authority], + recent_blockhash, + ); + let err = client + .process_transaction(transaction) + .await + .unwrap_err() + .unwrap(); + assert_eq!( + err, + TransactionError::InstructionError(2, InstructionError::InvalidInstructionData) + ); + + // close self-owned proof context accounts + let context_state_account_and_authority = Keypair::new(); + let context_state_info = ContextStateInfo { + context_state_account: &context_state_account_and_authority.pubkey(), + context_state_authority: &context_state_account_and_authority.pubkey(), + }; + + let instructions = vec![ + system_instruction::create_account( + &payer.pubkey(), + &context_state_account_and_authority.pubkey(), + 0_u64, + space as u64, + &zk_token_proof_program::id(), + ), + instruction_type.encode_verify_proof(Some(context_state_info), success_proof_data), + close_context_state(context_state_info, &context_state_account.pubkey()), + ]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&payer.pubkey()), + &[payer, &context_state_account_and_authority], + recent_blockhash, + ); + client.process_transaction(transaction).await.unwrap(); +} diff --git a/programs/zk-token-proof/src/lib.rs b/programs/zk-token-proof/src/lib.rs index 2b086d40926d9c..799d5a21146ade 100644 --- a/programs/zk-token-proof/src/lib.rs +++ b/programs/zk-token-proof/src/lib.rs @@ -3,26 +3,114 @@ use { bytemuck::Pod, solana_program_runtime::{ic_msg, invoke_context::InvokeContext}, - solana_sdk::instruction::{InstructionError, TRANSACTION_LEVEL_STACK_HEIGHT}, - solana_zk_token_sdk::zk_token_proof_instruction::*, + solana_sdk::{ + instruction::{InstructionError, TRANSACTION_LEVEL_STACK_HEIGHT}, + system_program, + }, + solana_zk_token_sdk::{ + zk_token_proof_instruction::*, + zk_token_proof_program::id, + zk_token_proof_state::{ProofContextState, ProofContextStateMeta}, + }, std::result::Result, }; -fn verify(invoke_context: &mut InvokeContext) -> Result<(), InstructionError> { +fn process_verify_proof(invoke_context: &mut InvokeContext) -> Result<(), InstructionError> +where + T: Pod + ZkProofData, + U: Pod, +{ let transaction_context = &invoke_context.transaction_context; let instruction_context = transaction_context.get_current_instruction_context()?; let instruction_data = instruction_context.get_instruction_data(); - let instruction = ProofInstruction::decode_data::(instruction_data); - - let proof = instruction.ok_or_else(|| { + let proof_data = ProofInstruction::proof_data::(instruction_data).ok_or_else(|| { ic_msg!(invoke_context, "invalid proof data"); InstructionError::InvalidInstructionData })?; - proof.verify().map_err(|err| { - ic_msg!(invoke_context, "proof verification failed: {:?}", err); + proof_data.verify_proof().map_err(|err| { + ic_msg!(invoke_context, "proof_verification failed: {:?}", err); InstructionError::InvalidInstructionData - }) + })?; + + // create context state if accounts are provided with the instruction + if instruction_context.get_number_of_instruction_accounts() > 0 { + let context_state_authority = *instruction_context + .try_borrow_instruction_account(transaction_context, 1)? + .get_key(); + + let mut proof_context_account = + instruction_context.try_borrow_instruction_account(transaction_context, 0)?; + + if *proof_context_account.get_owner() != id() { + return Err(InstructionError::InvalidAccountOwner); + } + + let proof_context_state_meta = + ProofContextStateMeta::try_from_bytes(proof_context_account.get_data())?; + + if proof_context_state_meta.proof_type != ProofType::Uninitialized.into() { + return Err(InstructionError::AccountAlreadyInitialized); + } + + let context_state_data = ProofContextState::encode( + &context_state_authority, + T::PROOF_TYPE, + proof_data.context_data(), + ); + + if proof_context_account.get_data().len() != context_state_data.len() { + return Err(InstructionError::InvalidAccountData); + } + + proof_context_account.set_data(context_state_data)?; + } + + Ok(()) +} + +fn process_close_proof_context(invoke_context: &mut InvokeContext) -> Result<(), InstructionError> { + let transaction_context = &invoke_context.transaction_context; + let instruction_context = transaction_context.get_current_instruction_context()?; + + let owner_pubkey = { + let owner_account = + instruction_context.try_borrow_instruction_account(transaction_context, 2)?; + + if !owner_account.is_signer() { + return Err(InstructionError::MissingRequiredSignature); + } + *owner_account.get_key() + }; // done with `owner_account`, so drop it to prevent a potential double borrow + + let proof_context_account_pubkey = *instruction_context + .try_borrow_instruction_account(transaction_context, 0)? + .get_key(); + let destination_account_pubkey = *instruction_context + .try_borrow_instruction_account(transaction_context, 1)? + .get_key(); + if proof_context_account_pubkey == destination_account_pubkey { + return Err(InstructionError::InvalidInstructionData); + } + + let mut proof_context_account = + instruction_context.try_borrow_instruction_account(transaction_context, 0)?; + let proof_context_state_meta = + ProofContextStateMeta::try_from_bytes(proof_context_account.get_data())?; + let expected_owner_pubkey = proof_context_state_meta.context_state_authority; + + if owner_pubkey != expected_owner_pubkey { + return Err(InstructionError::InvalidAccountOwner); + } + + let mut destination_account = + instruction_context.try_borrow_instruction_account(transaction_context, 1)?; + destination_account.checked_add_lamports(proof_context_account.get_lamports())?; + proof_context_account.set_lamports(0)?; + proof_context_account.set_data_length(0)?; + proof_context_account.set_owner(system_program::id().as_ref())?; + + Ok(()) } pub fn process_instruction(invoke_context: &mut InvokeContext) -> Result<(), InstructionError> { @@ -40,32 +128,39 @@ pub fn process_instruction(invoke_context: &mut InvokeContext) -> Result<(), Ins let transaction_context = &invoke_context.transaction_context; let instruction_context = transaction_context.get_current_instruction_context()?; let instruction_data = instruction_context.get_instruction_data(); - let instruction = ProofInstruction::decode_type(instruction_data); + let instruction = ProofInstruction::instruction_type(instruction_data) + .ok_or(InstructionError::InvalidInstructionData)?; - match instruction.ok_or(InstructionError::InvalidInstructionData)? { + match instruction { + ProofInstruction::CloseContextState => { + ic_msg!(invoke_context, "CloseContextState"); + process_close_proof_context(invoke_context) + } ProofInstruction::VerifyCloseAccount => { ic_msg!(invoke_context, "VerifyCloseAccount"); - verify::(invoke_context) + process_verify_proof::(invoke_context) } ProofInstruction::VerifyWithdraw => { ic_msg!(invoke_context, "VerifyWithdraw"); - verify::(invoke_context) + process_verify_proof::(invoke_context) } ProofInstruction::VerifyWithdrawWithheldTokens => { ic_msg!(invoke_context, "VerifyWithdrawWithheldTokens"); - verify::(invoke_context) + process_verify_proof::( + invoke_context, + ) } ProofInstruction::VerifyTransfer => { ic_msg!(invoke_context, "VerifyTransfer"); - verify::(invoke_context) + process_verify_proof::(invoke_context) } ProofInstruction::VerifyTransferWithFee => { ic_msg!(invoke_context, "VerifyTransferWithFee"); - verify::(invoke_context) + process_verify_proof::(invoke_context) } ProofInstruction::VerifyPubkeyValidity => { ic_msg!(invoke_context, "VerifyPubkeyValidity"); - verify::(invoke_context) + process_verify_proof::(invoke_context) } } } diff --git a/zk-token-sdk/src/instruction/close_account.rs b/zk-token-sdk/src/instruction/close_account.rs index 8bf3947ec28ca9..561fb1a093dd12 100644 --- a/zk-token-sdk/src/instruction/close_account.rs +++ b/zk-token-sdk/src/instruction/close_account.rs @@ -1,19 +1,21 @@ -use { - crate::zk_token_elgamal::pod, - bytemuck::{Pod, Zeroable}, -}; #[cfg(not(target_os = "solana"))] use { crate::{ encryption::elgamal::{ElGamalCiphertext, ElGamalKeypair, ElGamalPubkey}, errors::ProofError, - instruction::Verifiable, sigma_proofs::zero_balance_proof::ZeroBalanceProof, transcript::TranscriptProtocol, }, merlin::Transcript, std::convert::TryInto, }; +use { + crate::{ + instruction::{ProofType, ZkProofData}, + zk_token_elgamal::pod, + }, + bytemuck::{Pod, Zeroable}, +}; /// This struct includes the cryptographic proof *and* the account data information needed to verify /// the proof @@ -25,14 +27,21 @@ use { #[derive(Clone, Copy, Pod, Zeroable)] #[repr(C)] pub struct CloseAccountData { + /// The context data for the close account proof + pub context: CloseAccountProofContext, + + /// Proof that the source account available balance is zero + pub proof: CloseAccountProof, // 96 bytes +} + +#[derive(Clone, Copy, Pod, Zeroable)] +#[repr(C)] +pub struct CloseAccountProofContext { /// The source account ElGamal pubkey pub pubkey: pod::ElGamalPubkey, // 32 bytes /// The source account available balance in encrypted form pub ciphertext: pod::ElGamalCiphertext, // 64 bytes - - /// Proof that the source account available balance is zero - pub proof: CloseAccountProof, // 96 bytes } #[cfg(not(target_os = "solana"))] @@ -44,25 +53,32 @@ impl CloseAccountData { let pod_pubkey = pod::ElGamalPubkey(keypair.public.to_bytes()); let pod_ciphertext = pod::ElGamalCiphertext(ciphertext.to_bytes()); - let mut transcript = CloseAccountProof::transcript_new(&pod_pubkey, &pod_ciphertext); + let context = CloseAccountProofContext { + pubkey: pod_pubkey, + ciphertext: pod_ciphertext, + }; + let mut transcript = CloseAccountProof::transcript_new(&pod_pubkey, &pod_ciphertext); let proof = CloseAccountProof::new(keypair, ciphertext, &mut transcript); - Ok(CloseAccountData { - pubkey: pod_pubkey, - ciphertext: pod_ciphertext, - proof, - }) + Ok(CloseAccountData { context, proof }) } } -#[cfg(not(target_os = "solana"))] -impl Verifiable for CloseAccountData { - fn verify(&self) -> Result<(), ProofError> { - let mut transcript = CloseAccountProof::transcript_new(&self.pubkey, &self.ciphertext); +impl ZkProofData for CloseAccountData { + const PROOF_TYPE: ProofType = ProofType::CloseAccount; + + fn context_data(&self) -> &CloseAccountProofContext { + &self.context + } + + #[cfg(not(target_os = "solana"))] + fn verify_proof(&self) -> Result<(), ProofError> { + let mut transcript = + CloseAccountProof::transcript_new(&self.context.pubkey, &self.context.ciphertext); - let pubkey = self.pubkey.try_into()?; - let ciphertext = self.ciphertext.try_into()?; + let pubkey = self.context.pubkey.try_into()?; + let ciphertext = self.context.ciphertext.try_into()?; self.proof.verify(&pubkey, &ciphertext, &mut transcript) } } @@ -127,11 +143,11 @@ mod test { // general case: encryption of 0 let ciphertext = keypair.public.encrypt(0_u64); let close_account_data = CloseAccountData::new(&keypair, &ciphertext).unwrap(); - assert!(close_account_data.verify().is_ok()); + assert!(close_account_data.verify_proof().is_ok()); // general case: encryption of > 0 let ciphertext = keypair.public.encrypt(1_u64); let close_account_data = CloseAccountData::new(&keypair, &ciphertext).unwrap(); - assert!(close_account_data.verify().is_err()); + assert!(close_account_data.verify_proof().is_err()); } } diff --git a/zk-token-sdk/src/instruction/mod.rs b/zk-token-sdk/src/instruction/mod.rs index 3d5e44a5e01ee9..f257ee961a1a4e 100644 --- a/zk-token-sdk/src/instruction/mod.rs +++ b/zk-token-sdk/src/instruction/mod.rs @@ -5,6 +5,7 @@ pub mod transfer_with_fee; pub mod withdraw; pub mod withdraw_withheld; +use num_derive::{FromPrimitive, ToPrimitive}; #[cfg(not(target_os = "solana"))] use { crate::{ @@ -17,14 +18,35 @@ use { curve25519_dalek::scalar::Scalar, }; pub use { - close_account::CloseAccountData, pubkey_validity::PubkeyValidityData, transfer::TransferData, - transfer_with_fee::TransferWithFeeData, withdraw::WithdrawData, - withdraw_withheld::WithdrawWithheldTokensData, + bytemuck::Pod, + close_account::{CloseAccountData, CloseAccountProofContext}, + pubkey_validity::{PubkeyValidityData, PubkeyValidityProofContext}, + transfer::{TransferData, TransferProofContext}, + transfer_with_fee::{FeeParameters, TransferWithFeeData, TransferWithFeeProofContext}, + withdraw::{WithdrawData, WithdrawProofContext}, + withdraw_withheld::{WithdrawWithheldTokensData, WithdrawWithheldTokensProofContext}, }; -#[cfg(not(target_os = "solana"))] -pub trait Verifiable { - fn verify(&self) -> Result<(), ProofError>; +#[derive(Clone, Copy, Debug, FromPrimitive, ToPrimitive, PartialEq, Eq)] +#[repr(u8)] +pub enum ProofType { + /// Empty proof type used to distinguish if a proof context account is initialized + Uninitialized, + CloseAccount, + Withdraw, + WithdrawWithheldTokens, + Transfer, + TransferWithFee, + PubkeyValidity, +} + +pub trait ZkProofData { + const PROOF_TYPE: ProofType; + + fn context_data(&self) -> &T; + + #[cfg(not(target_os = "solana"))] + fn verify_proof(&self) -> Result<(), ProofError>; } #[cfg(not(target_os = "solana"))] diff --git a/zk-token-sdk/src/instruction/pubkey_validity.rs b/zk-token-sdk/src/instruction/pubkey_validity.rs index 3fc5a6288fb446..83f0461a1733f3 100644 --- a/zk-token-sdk/src/instruction/pubkey_validity.rs +++ b/zk-token-sdk/src/instruction/pubkey_validity.rs @@ -1,19 +1,21 @@ -use { - crate::zk_token_elgamal::pod, - bytemuck::{Pod, Zeroable}, -}; #[cfg(not(target_os = "solana"))] use { crate::{ encryption::elgamal::{ElGamalKeypair, ElGamalPubkey}, errors::ProofError, - instruction::Verifiable, sigma_proofs::pubkey_proof::PubkeySigmaProof, transcript::TranscriptProtocol, }, merlin::Transcript, std::convert::TryInto, }; +use { + crate::{ + instruction::{ProofType, ZkProofData}, + zk_token_elgamal::pod, + }, + bytemuck::{Pod, Zeroable}, +}; /// This struct includes the cryptographic proof *and* the account data information needed to /// verify the proof @@ -24,34 +26,45 @@ use { #[derive(Clone, Copy, Pod, Zeroable)] #[repr(C)] pub struct PubkeyValidityData { - /// The public key to be proved - pub pubkey: pod::ElGamalPubkey, + /// The context data for the public key validity proof + pub context: PubkeyValidityProofContext, /// Proof that the public key is well-formed pub proof: PubkeyValidityProof, // 64 bytes } +#[derive(Clone, Copy, Pod, Zeroable)] +#[repr(C)] +pub struct PubkeyValidityProofContext { + /// The public key to be proved + pub pubkey: pod::ElGamalPubkey, // 32 bytes +} + #[cfg(not(target_os = "solana"))] impl PubkeyValidityData { pub fn new(keypair: &ElGamalKeypair) -> Result { let pod_pubkey = pod::ElGamalPubkey(keypair.public.to_bytes()); - let mut transcript = PubkeyValidityProof::transcript_new(&pod_pubkey); + let context = PubkeyValidityProofContext { pubkey: pod_pubkey }; + let mut transcript = PubkeyValidityProof::transcript_new(&pod_pubkey); let proof = PubkeyValidityProof::new(keypair, &mut transcript); - Ok(PubkeyValidityData { - pubkey: pod_pubkey, - proof, - }) + Ok(PubkeyValidityData { context, proof }) } } -#[cfg(not(target_os = "solana"))] -impl Verifiable for PubkeyValidityData { - fn verify(&self) -> Result<(), ProofError> { - let mut transcript = PubkeyValidityProof::transcript_new(&self.pubkey); - let pubkey = self.pubkey.try_into()?; +impl ZkProofData for PubkeyValidityData { + const PROOF_TYPE: ProofType = ProofType::PubkeyValidity; + + fn context_data(&self) -> &PubkeyValidityProofContext { + &self.context + } + + #[cfg(not(target_os = "solana"))] + fn verify_proof(&self) -> Result<(), ProofError> { + let mut transcript = PubkeyValidityProof::transcript_new(&self.context.pubkey); + let pubkey = self.context.pubkey.try_into()?; self.proof.verify(&pubkey, &mut transcript) } } @@ -100,6 +113,6 @@ mod test { let keypair = ElGamalKeypair::new_rand(); let pubkey_validity_data = PubkeyValidityData::new(&keypair).unwrap(); - assert!(pubkey_validity_data.verify().is_ok()); + assert!(pubkey_validity_data.verify_proof().is_ok()); } } diff --git a/zk-token-sdk/src/instruction/transfer.rs b/zk-token-sdk/src/instruction/transfer.rs index 362ed46557cd1f..276993225211f4 100644 --- a/zk-token-sdk/src/instruction/transfer.rs +++ b/zk-token-sdk/src/instruction/transfer.rs @@ -1,7 +1,3 @@ -use { - crate::zk_token_elgamal::pod, - bytemuck::{Pod, Zeroable}, -}; #[cfg(not(target_os = "solana"))] use { crate::{ @@ -12,7 +8,7 @@ use { pedersen::{Pedersen, PedersenCommitment, PedersenOpening}, }, errors::ProofError, - instruction::{combine_lo_hi_ciphertexts, split_u64, Role, Verifiable}, + instruction::{combine_lo_hi_ciphertexts, split_u64, Role}, range_proof::RangeProof, sigma_proofs::{ equality_proof::CtxtCommEqualityProof, validity_proof::AggregatedValidityProof, @@ -23,6 +19,13 @@ use { merlin::Transcript, std::convert::TryInto, }; +use { + crate::{ + instruction::{ProofType, ZkProofData}, + zk_token_elgamal::pod, + }, + bytemuck::{Pod, Zeroable}, +}; #[cfg(not(target_os = "solana"))] const TRANSFER_SOURCE_AMOUNT_BITS: usize = 64; @@ -42,20 +45,27 @@ lazy_static::lazy_static! { #[derive(Clone, Copy, Pod, Zeroable)] #[repr(C)] pub struct TransferData { + /// The context data for the transfer proof + pub context: TransferProofContext, + + /// Zero-knowledge proofs for Transfer + pub proof: TransferProof, +} + +#[derive(Clone, Copy, Pod, Zeroable)] +#[repr(C)] +pub struct TransferProofContext { /// Group encryption of the low 16 bits of the transfer amount - pub ciphertext_lo: pod::TransferAmountEncryption, + pub ciphertext_lo: pod::TransferAmountEncryption, // 128 bytes /// Group encryption of the high 48 bits of the transfer amount - pub ciphertext_hi: pod::TransferAmountEncryption, + pub ciphertext_hi: pod::TransferAmountEncryption, // 128 bytes /// The public encryption keys associated with the transfer: source, dest, and auditor - pub transfer_pubkeys: pod::TransferPubkeys, + pub transfer_pubkeys: pod::TransferPubkeys, // 96 bytes /// The final spendable ciphertext after the transfer - pub new_source_ciphertext: pod::ElGamalCiphertext, - - /// Zero-knowledge proofs for Transfer - pub proof: TransferProof, + pub new_source_ciphertext: pod::ElGamalCiphertext, // 64 bytes } #[cfg(not(target_os = "solana"))] @@ -116,6 +126,13 @@ impl TransferData { let pod_ciphertext_hi: pod::TransferAmountEncryption = ciphertext_hi.into(); let pod_new_source_ciphertext: pod::ElGamalCiphertext = new_source_ciphertext.into(); + let context = TransferProofContext { + ciphertext_lo: pod_ciphertext_lo, + ciphertext_hi: pod_ciphertext_hi, + transfer_pubkeys: pod_transfer_pubkeys, + new_source_ciphertext: pod_new_source_ciphertext, + }; + let mut transcript = TransferProof::transcript_new( &pod_transfer_pubkeys, &pod_ciphertext_lo, @@ -133,18 +150,12 @@ impl TransferData { &mut transcript, ); - Ok(Self { - ciphertext_lo: pod_ciphertext_lo, - ciphertext_hi: pod_ciphertext_hi, - transfer_pubkeys: pod_transfer_pubkeys, - new_source_ciphertext: pod_new_source_ciphertext, - proof, - }) + Ok(Self { context, proof }) } /// Extracts the lo ciphertexts associated with a transfer data fn ciphertext_lo(&self, role: Role) -> Result { - let ciphertext_lo: TransferAmountEncryption = self.ciphertext_lo.try_into()?; + let ciphertext_lo: TransferAmountEncryption = self.context.ciphertext_lo.try_into()?; let handle_lo = match role { Role::Source => Some(ciphertext_lo.source_handle), @@ -165,7 +176,7 @@ impl TransferData { /// Extracts the lo ciphertexts associated with a transfer data fn ciphertext_hi(&self, role: Role) -> Result { - let ciphertext_hi: TransferAmountEncryption = self.ciphertext_hi.try_into()?; + let ciphertext_hi: TransferAmountEncryption = self.context.ciphertext_hi.try_into()?; let handle_hi = match role { Role::Source => Some(ciphertext_hi.source_handle), @@ -201,21 +212,27 @@ impl TransferData { } } -#[cfg(not(target_os = "solana"))] -impl Verifiable for TransferData { - fn verify(&self) -> Result<(), ProofError> { +impl ZkProofData for TransferData { + const PROOF_TYPE: ProofType = ProofType::Transfer; + + fn context_data(&self) -> &TransferProofContext { + &self.context + } + + #[cfg(not(target_os = "solana"))] + fn verify_proof(&self) -> Result<(), ProofError> { // generate transcript and append all public inputs let mut transcript = TransferProof::transcript_new( - &self.transfer_pubkeys, - &self.ciphertext_lo, - &self.ciphertext_hi, - &self.new_source_ciphertext, + &self.context.transfer_pubkeys, + &self.context.ciphertext_lo, + &self.context.ciphertext_hi, + &self.context.new_source_ciphertext, ); - let ciphertext_lo = self.ciphertext_lo.try_into()?; - let ciphertext_hi = self.ciphertext_hi.try_into()?; - let transfer_pubkeys = self.transfer_pubkeys.try_into()?; - let new_spendable_ciphertext = self.new_source_ciphertext.try_into()?; + let ciphertext_lo = self.context.ciphertext_lo.try_into()?; + let ciphertext_hi = self.context.ciphertext_hi.try_into()?; + let transfer_pubkeys = self.context.transfer_pubkeys.try_into()?; + let new_spendable_ciphertext = self.context.new_source_ciphertext.try_into()?; self.proof.verify( &ciphertext_lo, @@ -537,7 +554,7 @@ mod test { ) .unwrap(); - assert!(transfer_data.verify().is_ok()); + assert!(transfer_data.verify_proof().is_ok()); // Case 2: transfer max amount @@ -558,7 +575,7 @@ mod test { ) .unwrap(); - assert!(transfer_data.verify().is_ok()); + assert!(transfer_data.verify_proof().is_ok()); // Case 3: general success case @@ -578,7 +595,7 @@ mod test { ) .unwrap(); - assert!(transfer_data.verify().is_ok()); + assert!(transfer_data.verify_proof().is_ok()); // Case 4: invalid destination or auditor pubkey let spendable_balance: u64 = 0; @@ -598,7 +615,7 @@ mod test { ) .unwrap(); - assert!(transfer_data.verify().is_err()); + assert!(transfer_data.verify_proof().is_err()); // auditor pubkey invalid let dest_pk = ElGamalKeypair::new_rand().public; @@ -612,7 +629,7 @@ mod test { ) .unwrap(); - assert!(transfer_data.verify().is_err()); + assert!(transfer_data.verify_proof().is_err()); } #[test] diff --git a/zk-token-sdk/src/instruction/transfer_with_fee.rs b/zk-token-sdk/src/instruction/transfer_with_fee.rs index 7fbb603e2ec333..a6bdee776b54aa 100644 --- a/zk-token-sdk/src/instruction/transfer_with_fee.rs +++ b/zk-token-sdk/src/instruction/transfer_with_fee.rs @@ -1,7 +1,3 @@ -use { - crate::zk_token_elgamal::pod, - bytemuck::{Pod, Zeroable}, -}; #[cfg(not(target_os = "solana"))] use { crate::{ @@ -14,7 +10,7 @@ use { errors::ProofError, instruction::{ combine_lo_hi_ciphertexts, combine_lo_hi_commitments, combine_lo_hi_openings, - combine_lo_hi_u64, split_u64, transfer::TransferAmountEncryption, Role, Verifiable, + combine_lo_hi_u64, split_u64, transfer::TransferAmountEncryption, Role, }, range_proof::RangeProof, sigma_proofs::{ @@ -29,6 +25,13 @@ use { std::convert::TryInto, subtle::{ConditionallySelectable, ConstantTimeGreater}, }; +use { + crate::{ + instruction::{ProofType, ZkProofData}, + zk_token_elgamal::pod, + }, + bytemuck::{Pod, Zeroable}, +}; #[cfg(not(target_os = "solana"))] const MAX_FEE_BASIS_POINTS: u64 = 10_000; @@ -61,29 +64,36 @@ lazy_static::lazy_static! { #[derive(Clone, Copy, Pod, Zeroable)] #[repr(C)] pub struct TransferWithFeeData { + /// The context data for the transfer with fee proof + pub context: TransferWithFeeProofContext, + + // transfer fee proof + pub proof: TransferWithFeeProof, +} + +#[derive(Clone, Copy, Pod, Zeroable)] +#[repr(C)] +pub struct TransferWithFeeProofContext { /// Group encryption of the low 16 bites of the transfer amount - pub ciphertext_lo: pod::TransferAmountEncryption, + pub ciphertext_lo: pod::TransferAmountEncryption, // 128 bytes /// Group encryption of the high 48 bits of the transfer amount - pub ciphertext_hi: pod::TransferAmountEncryption, + pub ciphertext_hi: pod::TransferAmountEncryption, // 128 bytes /// The public encryption keys associated with the transfer: source, dest, and auditor - pub transfer_with_fee_pubkeys: pod::TransferWithFeePubkeys, + pub transfer_with_fee_pubkeys: pod::TransferWithFeePubkeys, // 128 bytes /// The final spendable ciphertext after the transfer, - pub new_source_ciphertext: pod::ElGamalCiphertext, + pub new_source_ciphertext: pod::ElGamalCiphertext, // 64 bytes // transfer fee encryption of the low 16 bits of the transfer fee amount - pub fee_ciphertext_lo: pod::FeeEncryption, + pub fee_ciphertext_lo: pod::FeeEncryption, // 96 bytes // transfer fee encryption of the hi 32 bits of the transfer fee amount - pub fee_ciphertext_hi: pod::FeeEncryption, + pub fee_ciphertext_hi: pod::FeeEncryption, // 96 bytes // fee parameters - pub fee_parameters: pod::FeeParameters, - - // transfer fee proof - pub proof: TransferWithFeeProof, + pub fee_parameters: pod::FeeParameters, // 10 bytes } #[cfg(not(target_os = "solana"))] @@ -173,6 +183,16 @@ impl TransferWithFeeData { let pod_fee_ciphertext_lo: pod::FeeEncryption = fee_ciphertext_lo.to_pod(); let pod_fee_ciphertext_hi: pod::FeeEncryption = fee_ciphertext_hi.to_pod(); + let context = TransferWithFeeProofContext { + ciphertext_lo: pod_ciphertext_lo, + ciphertext_hi: pod_ciphertext_hi, + transfer_with_fee_pubkeys: pod_transfer_with_fee_pubkeys, + new_source_ciphertext: pod_new_source_ciphertext, + fee_ciphertext_lo: pod_fee_ciphertext_lo, + fee_ciphertext_hi: pod_fee_ciphertext_hi, + fee_parameters: fee_parameters.into(), + }; + let mut transcript = TransferWithFeeProof::transcript_new( &pod_transfer_with_fee_pubkeys, &pod_ciphertext_lo, @@ -196,21 +216,12 @@ impl TransferWithFeeData { &mut transcript, ); - Ok(Self { - ciphertext_lo: pod_ciphertext_lo, - ciphertext_hi: pod_ciphertext_hi, - transfer_with_fee_pubkeys: pod_transfer_with_fee_pubkeys, - new_source_ciphertext: pod_new_source_ciphertext, - fee_ciphertext_lo: pod_fee_ciphertext_lo, - fee_ciphertext_hi: pod_fee_ciphertext_hi, - fee_parameters: fee_parameters.into(), - proof, - }) + Ok(Self { context, proof }) } /// Extracts the lo ciphertexts associated with a transfer-with-fee data fn ciphertext_lo(&self, role: Role) -> Result { - let ciphertext_lo: TransferAmountEncryption = self.ciphertext_lo.try_into()?; + let ciphertext_lo: TransferAmountEncryption = self.context.ciphertext_lo.try_into()?; let handle_lo = match role { Role::Source => Some(ciphertext_lo.source_handle), @@ -231,7 +242,7 @@ impl TransferWithFeeData { /// Extracts the lo ciphertexts associated with a transfer-with-fee data fn ciphertext_hi(&self, role: Role) -> Result { - let ciphertext_hi: TransferAmountEncryption = self.ciphertext_hi.try_into()?; + let ciphertext_hi: TransferAmountEncryption = self.context.ciphertext_hi.try_into()?; let handle_hi = match role { Role::Source => Some(ciphertext_hi.source_handle), @@ -252,7 +263,7 @@ impl TransferWithFeeData { /// Extracts the lo fee ciphertexts associated with a transfer_with_fee data fn fee_ciphertext_lo(&self, role: Role) -> Result { - let fee_ciphertext_lo: FeeEncryption = self.fee_ciphertext_lo.try_into()?; + let fee_ciphertext_lo: FeeEncryption = self.context.fee_ciphertext_lo.try_into()?; let fee_handle_lo = match role { Role::Source => None, @@ -275,7 +286,7 @@ impl TransferWithFeeData { /// Extracts the hi fee ciphertexts associated with a transfer_with_fee data fn fee_ciphertext_hi(&self, role: Role) -> Result { - let fee_ciphertext_hi: FeeEncryption = self.fee_ciphertext_hi.try_into()?; + let fee_ciphertext_hi: FeeEncryption = self.context.fee_ciphertext_hi.try_into()?; let fee_handle_hi = match role { Role::Source => None, @@ -329,26 +340,32 @@ impl TransferWithFeeData { } } -#[cfg(not(target_os = "solana"))] -impl Verifiable for TransferWithFeeData { - fn verify(&self) -> Result<(), ProofError> { +impl ZkProofData for TransferWithFeeData { + const PROOF_TYPE: ProofType = ProofType::TransferWithFee; + + fn context_data(&self) -> &TransferWithFeeProofContext { + &self.context + } + + #[cfg(not(target_os = "solana"))] + fn verify_proof(&self) -> Result<(), ProofError> { let mut transcript = TransferWithFeeProof::transcript_new( - &self.transfer_with_fee_pubkeys, - &self.ciphertext_lo, - &self.ciphertext_hi, - &self.new_source_ciphertext, - &self.fee_ciphertext_lo, - &self.fee_ciphertext_hi, + &self.context.transfer_with_fee_pubkeys, + &self.context.ciphertext_lo, + &self.context.ciphertext_hi, + &self.context.new_source_ciphertext, + &self.context.fee_ciphertext_lo, + &self.context.fee_ciphertext_hi, ); - let ciphertext_lo = self.ciphertext_lo.try_into()?; - let ciphertext_hi = self.ciphertext_hi.try_into()?; - let pubkeys_transfer_with_fee = self.transfer_with_fee_pubkeys.try_into()?; - let new_source_ciphertext = self.new_source_ciphertext.try_into()?; + let ciphertext_lo = self.context.ciphertext_lo.try_into()?; + let ciphertext_hi = self.context.ciphertext_hi.try_into()?; + let pubkeys_transfer_with_fee = self.context.transfer_with_fee_pubkeys.try_into()?; + let new_source_ciphertext = self.context.new_source_ciphertext.try_into()?; - let fee_ciphertext_lo = self.fee_ciphertext_lo.try_into()?; - let fee_ciphertext_hi = self.fee_ciphertext_hi.try_into()?; - let fee_parameters = self.fee_parameters.into(); + let fee_ciphertext_lo = self.context.fee_ciphertext_lo.try_into()?; + let fee_ciphertext_hi = self.context.fee_ciphertext_hi.try_into()?; + let fee_parameters = self.context.fee_parameters.into(); self.proof.verify( &ciphertext_lo, @@ -886,7 +903,7 @@ mod test { ) .unwrap(); - assert!(fee_data.verify().is_ok()); + assert!(fee_data.verify_proof().is_ok()); // Case 2: transfer max amount let spendable_balance: u64 = u64::max_value(); @@ -910,7 +927,7 @@ mod test { ) .unwrap(); - assert!(fee_data.verify().is_ok()); + assert!(fee_data.verify_proof().is_ok()); // Case 3: general success case let spendable_balance: u64 = 120; @@ -933,7 +950,7 @@ mod test { ) .unwrap(); - assert!(fee_data.verify().is_ok()); + assert!(fee_data.verify_proof().is_ok()); // Case 4: invalid destination, auditor, or withdraw authority pubkeys let spendable_balance: u64 = 120; @@ -961,7 +978,7 @@ mod test { ) .unwrap(); - assert!(fee_data.verify().is_err()); + assert!(fee_data.verify_proof().is_err()); // auditor pubkey invalid let destination_pubkey: ElGamalPubkey = ElGamalKeypair::new_rand().public; @@ -978,7 +995,7 @@ mod test { ) .unwrap(); - assert!(fee_data.verify().is_err()); + assert!(fee_data.verify_proof().is_err()); // withdraw authority invalid let destination_pubkey: ElGamalPubkey = ElGamalKeypair::new_rand().public; @@ -995,6 +1012,6 @@ mod test { ) .unwrap(); - assert!(fee_data.verify().is_err()); + assert!(fee_data.verify_proof().is_err()); } } diff --git a/zk-token-sdk/src/instruction/withdraw.rs b/zk-token-sdk/src/instruction/withdraw.rs index 6d996bb6b59b09..1bd29d3510365d 100644 --- a/zk-token-sdk/src/instruction/withdraw.rs +++ b/zk-token-sdk/src/instruction/withdraw.rs @@ -1,7 +1,3 @@ -use { - crate::zk_token_elgamal::pod, - bytemuck::{Pod, Zeroable}, -}; #[cfg(not(target_os = "solana"))] use { crate::{ @@ -10,7 +6,6 @@ use { pedersen::{Pedersen, PedersenCommitment}, }, errors::ProofError, - instruction::Verifiable, range_proof::RangeProof, sigma_proofs::equality_proof::CtxtCommEqualityProof, transcript::TranscriptProtocol, @@ -18,6 +13,13 @@ use { merlin::Transcript, std::convert::TryInto, }; +use { + crate::{ + instruction::{ProofType, ZkProofData}, + zk_token_elgamal::pod, + }, + bytemuck::{Pod, Zeroable}, +}; #[cfg(not(target_os = "solana"))] const WITHDRAW_AMOUNT_BIT_LENGTH: usize = 64; @@ -32,15 +34,22 @@ const WITHDRAW_AMOUNT_BIT_LENGTH: usize = 64; #[derive(Clone, Copy, Pod, Zeroable)] #[repr(C)] pub struct WithdrawData { + /// The context data for the withdraw proof + pub context: WithdrawProofContext, // 128 bytes + + /// Range proof + pub proof: WithdrawProof, // 736 bytes +} + +#[derive(Clone, Copy, Pod, Zeroable)] +#[repr(C)] +pub struct WithdrawProofContext { /// The source account ElGamal pubkey pub pubkey: pod::ElGamalPubkey, // 32 bytes /// The source account available balance *after* the withdraw (encrypted by /// `source_pk` pub final_ciphertext: pod::ElGamalCiphertext, // 64 bytes - - /// Range proof - pub proof: WithdrawProof, // 736 bytes } #[cfg(not(target_os = "solana"))] @@ -64,24 +73,33 @@ impl WithdrawData { let pod_pubkey = pod::ElGamalPubkey(keypair.public.to_bytes()); let pod_final_ciphertext: pod::ElGamalCiphertext = final_ciphertext.into(); - let mut transcript = WithdrawProof::transcript_new(&pod_pubkey, &pod_final_ciphertext); - let proof = WithdrawProof::new(keypair, final_balance, &final_ciphertext, &mut transcript); - Ok(Self { + let context = WithdrawProofContext { pubkey: pod_pubkey, final_ciphertext: pod_final_ciphertext, - proof, - }) + }; + + let mut transcript = WithdrawProof::transcript_new(&pod_pubkey, &pod_final_ciphertext); + let proof = WithdrawProof::new(keypair, final_balance, &final_ciphertext, &mut transcript); + + Ok(Self { context, proof }) } } -#[cfg(not(target_os = "solana"))] -impl Verifiable for WithdrawData { - fn verify(&self) -> Result<(), ProofError> { - let mut transcript = WithdrawProof::transcript_new(&self.pubkey, &self.final_ciphertext); +impl ZkProofData for WithdrawData { + const PROOF_TYPE: ProofType = ProofType::Withdraw; + + fn context_data(&self) -> &WithdrawProofContext { + &self.context + } + + #[cfg(not(target_os = "solana"))] + fn verify_proof(&self) -> Result<(), ProofError> { + let mut transcript = + WithdrawProof::transcript_new(&self.context.pubkey, &self.context.final_ciphertext); - let elgamal_pubkey = self.pubkey.try_into()?; - let final_balance_ciphertext = self.final_ciphertext.try_into()?; + let elgamal_pubkey = self.context.pubkey.try_into()?; + let final_balance_ciphertext = self.context.final_ciphertext.try_into()?; self.proof .verify(&elgamal_pubkey, &final_balance_ciphertext, &mut transcript) } @@ -200,7 +218,7 @@ mod test { ¤t_ciphertext, ) .unwrap(); - assert!(data.verify().is_ok()); + assert!(data.verify_proof().is_ok()); // generate and verify proof with wrong balance let wrong_balance: u64 = 99; @@ -211,6 +229,6 @@ mod test { ¤t_ciphertext, ) .unwrap(); - assert!(data.verify().is_err()); + assert!(data.verify_proof().is_err()); } } diff --git a/zk-token-sdk/src/instruction/withdraw_withheld.rs b/zk-token-sdk/src/instruction/withdraw_withheld.rs index 327f6295d571da..7a92787709726e 100644 --- a/zk-token-sdk/src/instruction/withdraw_withheld.rs +++ b/zk-token-sdk/src/instruction/withdraw_withheld.rs @@ -1,7 +1,3 @@ -use { - crate::zk_token_elgamal::pod, - bytemuck::{Pod, Zeroable}, -}; #[cfg(not(target_os = "solana"))] use { crate::{ @@ -10,13 +6,19 @@ use { pedersen::PedersenOpening, }, errors::ProofError, - instruction::Verifiable, sigma_proofs::equality_proof::CtxtCtxtEqualityProof, transcript::TranscriptProtocol, }, merlin::Transcript, std::convert::TryInto, }; +use { + crate::{ + instruction::{ProofType, ZkProofData}, + zk_token_elgamal::pod, + }, + bytemuck::{Pod, Zeroable}, +}; /// This struct includes the cryptographic proof *and* the account data information needed to verify /// the proof @@ -28,15 +30,21 @@ use { #[derive(Clone, Copy, Pod, Zeroable)] #[repr(C)] pub struct WithdrawWithheldTokensData { - pub withdraw_withheld_authority_pubkey: pod::ElGamalPubkey, + pub context: WithdrawWithheldTokensProofContext, - pub destination_pubkey: pod::ElGamalPubkey, + pub proof: WithdrawWithheldTokensProof, +} - pub withdraw_withheld_authority_ciphertext: pod::ElGamalCiphertext, +#[derive(Clone, Copy, Pod, Zeroable)] +#[repr(C)] +pub struct WithdrawWithheldTokensProofContext { + pub withdraw_withheld_authority_pubkey: pod::ElGamalPubkey, // 32 bytes - pub destination_ciphertext: pod::ElGamalCiphertext, + pub destination_pubkey: pod::ElGamalPubkey, // 32 bytes - pub proof: WithdrawWithheldTokensProof, + pub withdraw_withheld_authority_ciphertext: pod::ElGamalCiphertext, // 64 bytes + + pub destination_ciphertext: pod::ElGamalCiphertext, // 64 bytes } #[cfg(not(target_os = "solana"))] @@ -58,6 +66,13 @@ impl WithdrawWithheldTokensData { pod::ElGamalCiphertext(withdraw_withheld_authority_ciphertext.to_bytes()); let pod_destination_ciphertext = pod::ElGamalCiphertext(destination_ciphertext.to_bytes()); + let context = WithdrawWithheldTokensProofContext { + withdraw_withheld_authority_pubkey: pod_withdraw_withheld_authority_pubkey, + destination_pubkey: pod_destination_pubkey, + withdraw_withheld_authority_ciphertext: pod_withdraw_withheld_authority_ciphertext, + destination_ciphertext: pod_destination_ciphertext, + }; + let mut transcript = WithdrawWithheldTokensProof::transcript_new( &pod_withdraw_withheld_authority_pubkey, &pod_destination_pubkey, @@ -74,32 +89,34 @@ impl WithdrawWithheldTokensData { &mut transcript, ); - Ok(Self { - withdraw_withheld_authority_pubkey: pod_withdraw_withheld_authority_pubkey, - destination_pubkey: pod_destination_pubkey, - withdraw_withheld_authority_ciphertext: pod_withdraw_withheld_authority_ciphertext, - destination_ciphertext: pod_destination_ciphertext, - proof, - }) + Ok(Self { context, proof }) } } -#[cfg(not(target_os = "solana"))] -impl Verifiable for WithdrawWithheldTokensData { - fn verify(&self) -> Result<(), ProofError> { +impl ZkProofData for WithdrawWithheldTokensData { + const PROOF_TYPE: ProofType = ProofType::WithdrawWithheldTokens; + + fn context_data(&self) -> &WithdrawWithheldTokensProofContext { + &self.context + } + + #[cfg(not(target_os = "solana"))] + fn verify_proof(&self) -> Result<(), ProofError> { let mut transcript = WithdrawWithheldTokensProof::transcript_new( - &self.withdraw_withheld_authority_pubkey, - &self.destination_pubkey, - &self.withdraw_withheld_authority_ciphertext, - &self.destination_ciphertext, + &self.context.withdraw_withheld_authority_pubkey, + &self.context.destination_pubkey, + &self.context.withdraw_withheld_authority_ciphertext, + &self.context.destination_ciphertext, ); let withdraw_withheld_authority_pubkey = - self.withdraw_withheld_authority_pubkey.try_into()?; - let destination_pubkey = self.destination_pubkey.try_into()?; - let withdraw_withheld_authority_ciphertext = - self.withdraw_withheld_authority_ciphertext.try_into()?; - let destination_ciphertext = self.destination_ciphertext.try_into()?; + self.context.withdraw_withheld_authority_pubkey.try_into()?; + let destination_pubkey = self.context.destination_pubkey.try_into()?; + let withdraw_withheld_authority_ciphertext = self + .context + .withdraw_withheld_authority_ciphertext + .try_into()?; + let destination_ciphertext = self.context.destination_ciphertext.try_into()?; self.proof.verify( &withdraw_withheld_authority_pubkey, @@ -210,7 +227,7 @@ mod test { ) .unwrap(); - assert!(withdraw_withheld_tokens_data.verify().is_ok()); + assert!(withdraw_withheld_tokens_data.verify_proof().is_ok()); let amount: u64 = 55; let withdraw_withheld_authority_ciphertext = @@ -224,7 +241,7 @@ mod test { ) .unwrap(); - assert!(withdraw_withheld_tokens_data.verify().is_ok()); + assert!(withdraw_withheld_tokens_data.verify_proof().is_ok()); let amount = u64::max_value(); let withdraw_withheld_authority_ciphertext = @@ -238,6 +255,6 @@ mod test { ) .unwrap(); - assert!(withdraw_withheld_tokens_data.verify().is_ok()); + assert!(withdraw_withheld_tokens_data.verify_proof().is_ok()); } } diff --git a/zk-token-sdk/src/lib.rs b/zk-token-sdk/src/lib.rs index c7a0391eaab8c8..14ce5b4380af47 100644 --- a/zk-token-sdk/src/lib.rs +++ b/zk-token-sdk/src/lib.rs @@ -37,3 +37,4 @@ pub mod instruction; pub mod zk_token_elgamal; pub mod zk_token_proof_instruction; pub mod zk_token_proof_program; +pub mod zk_token_proof_state; diff --git a/zk-token-sdk/src/zk_token_elgamal/pod.rs b/zk-token-sdk/src/zk_token_elgamal/pod.rs index 2658e3a1446357..fd728bc6a7caab 100644 --- a/zk-token-sdk/src/zk_token_elgamal/pod.rs +++ b/zk-token-sdk/src/zk_token_elgamal/pod.rs @@ -1,5 +1,10 @@ pub use bytemuck::{Pod, Zeroable}; -use std::fmt; +use { + crate::zk_token_proof_instruction::ProofType, + num_traits::{FromPrimitive, ToPrimitive}, + solana_program::instruction::InstructionError, + std::fmt, +}; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Pod, Zeroable)] #[repr(transparent)] @@ -29,6 +34,22 @@ impl From for u64 { } } +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Pod, Zeroable)] +#[repr(transparent)] +pub struct PodProofType(u8); +impl From for PodProofType { + fn from(proof_type: ProofType) -> Self { + Self(ToPrimitive::to_u8(&proof_type).unwrap()) + } +} +impl TryFrom for ProofType { + type Error = InstructionError; + + fn try_from(pod: PodProofType) -> Result { + FromPrimitive::from_u8(pod.0).ok_or(Self::Error::InvalidAccountData) + } +} + #[derive(Clone, Copy, Pod, Zeroable, PartialEq, Eq)] #[repr(transparent)] pub struct CompressedRistretto(pub [u8; 32]); diff --git a/zk-token-sdk/src/zk_token_proof_instruction.rs b/zk-token-sdk/src/zk_token_proof_instruction.rs index 867e0f783c901a..21d1dae975562d 100644 --- a/zk-token-sdk/src/zk_token_proof_instruction.rs +++ b/zk-token-sdk/src/zk_token_proof_instruction.rs @@ -1,18 +1,41 @@ ///! Instructions provided by the ZkToken Proof program pub use crate::instruction::*; use { - bytemuck::{bytes_of, Pod}, + bytemuck::bytes_of, num_derive::{FromPrimitive, ToPrimitive}, num_traits::{FromPrimitive, ToPrimitive}, - solana_program::instruction::Instruction, + solana_program::{ + instruction::{AccountMeta, Instruction}, + pubkey::Pubkey, + }, }; #[derive(Clone, Copy, Debug, FromPrimitive, ToPrimitive, PartialEq, Eq)] #[repr(u8)] pub enum ProofInstruction { - /// Verify a `CloseAccountData` struct + /// Close a zero-knowledge proof context state. /// /// Accounts expected by this instruction: + /// 0. `[writable]` The proof context account to close + /// 1. `[writable]` The destination account for lamports + /// 2. `[signer]` The context account's owner + /// + /// Data expected by this instruction: + /// None + /// + CloseContextState, + + /// Verify a close account zero-knowledge proof. + /// + /// This instruction can be configured to optionally create a proof context state account. + /// + /// Accounts expected by this instruction: + /// + /// * Creating a proof context account + /// 0. `[writable]` The proof context account + /// 1. `[]` The proof context account owner + /// + /// * Otherwise /// None /// /// Data expected by this instruction: @@ -20,9 +43,17 @@ pub enum ProofInstruction { /// VerifyCloseAccount, - /// Verify a `WithdrawData` struct + /// Verify a withdraw zero-knowledge proof. + /// + /// This instruction can be configured to optionally create a proof context state account. /// /// Accounts expected by this instruction: + /// + /// * Creating a proof context account + /// 0. `[writable]` The proof context account + /// 1. `[]` The proof context account owner + /// + /// * Otherwise /// None /// /// Data expected by this instruction: @@ -30,9 +61,17 @@ pub enum ProofInstruction { /// VerifyWithdraw, - /// Verify a `WithdrawWithheldTokensData` struct + /// Verify a withdraw withheld tokens zero-knowledge proof. + /// + /// This instruction can be configured to optionally create a proof context state account. /// /// Accounts expected by this instruction: + /// + /// * Creating a proof context account + /// 0. `[writable]` The proof context account + /// 1. `[]` The proof context account owner + /// + /// * Otherwise /// None /// /// Data expected by this instruction: @@ -40,9 +79,17 @@ pub enum ProofInstruction { /// VerifyWithdrawWithheldTokens, - /// Verify a `TransferData` struct + /// Verify a transfer zero-knowledge proof. + /// + /// This instruction can be configured to optionally create a proof context state account. /// /// Accounts expected by this instruction: + /// + /// * Creating a proof context account + /// 0. `[writable]` The proof context account + /// 1. `[]` The proof context account owner + /// + /// * Otherwise /// None /// /// Data expected by this instruction: @@ -50,9 +97,17 @@ pub enum ProofInstruction { /// VerifyTransfer, - /// Verify a `TransferWithFeeData` struct + /// Verify a transfer with fee zero-knowledge proof. + /// + /// This instruction can be configured to optionally create a proof context state account. /// /// Accounts expected by this instruction: + /// + /// * Creating a proof context account + /// 0. `[writable]` The proof context account + /// 1. `[]` The proof context account owner + /// + /// * Otherwise /// None /// /// Data expected by this instruction: @@ -60,9 +115,17 @@ pub enum ProofInstruction { /// VerifyTransferWithFee, - /// Verify a `PubkeyValidityData` struct + /// Verify a pubkey validity zero-knowledge proof. + /// + /// This instruction can be configured to optionally create a proof context state account. /// /// Accounts expected by this instruction: + /// + /// * Creating a proof context account + /// 0. `[writable]` The proof context account + /// 1. `[]` The proof context account owner + /// + /// * Otherwise /// None /// /// Data expected by this instruction: @@ -71,50 +134,124 @@ pub enum ProofInstruction { VerifyPubkeyValidity, } -impl ProofInstruction { - pub fn encode(&self, proof: &T) -> Instruction { - let mut data = vec![ToPrimitive::to_u8(self).unwrap()]; - data.extend_from_slice(bytes_of(proof)); - Instruction { - program_id: crate::zk_token_proof_program::id(), - accounts: vec![], - data, - } - } +/// Pubkeys associated with a context state account to be used as parameters to functions. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct ContextStateInfo<'a> { + pub context_state_account: &'a Pubkey, + pub context_state_authority: &'a Pubkey, +} - pub fn decode_type(input: &[u8]) -> Option { - input.first().and_then(|x| FromPrimitive::from_u8(*x)) - } +/// Create a `CloseContextState` instruction. +pub fn close_context_state( + context_state_info: ContextStateInfo, + destination_account: &Pubkey, +) -> Instruction { + let accounts = vec![ + AccountMeta::new(*context_state_info.context_state_account, false), + AccountMeta::new(*destination_account, false), + AccountMeta::new_readonly(*context_state_info.context_state_authority, true), + ]; - pub fn decode_data(input: &[u8]) -> Option<&T> { - if input.is_empty() { - None - } else { - bytemuck::try_from_bytes(&input[1..]).ok() - } + let data = vec![ToPrimitive::to_u8(&ProofInstruction::CloseContextState).unwrap()]; + + Instruction { + program_id: crate::zk_token_proof_program::id(), + accounts, + data, } } -pub fn verify_close_account(proof_data: &CloseAccountData) -> Instruction { - ProofInstruction::VerifyCloseAccount.encode(proof_data) +/// Create a `VerifyCloseAccount` instruction. +pub fn verify_close_account( + context_state_info: Option, + proof_data: &CloseAccountData, +) -> Instruction { + ProofInstruction::VerifyCloseAccount.encode_verify_proof(context_state_info, proof_data) } -pub fn verify_withdraw(proof_data: &WithdrawData) -> Instruction { - ProofInstruction::VerifyWithdraw.encode(proof_data) +/// Create a `VerifyWithdraw` instruction. +pub fn verify_withdraw( + context_state_info: Option, + proof_data: &WithdrawData, +) -> Instruction { + ProofInstruction::VerifyWithdraw.encode_verify_proof(context_state_info, proof_data) } -pub fn verify_withdraw_withheld_tokens(proof_data: &WithdrawWithheldTokensData) -> Instruction { - ProofInstruction::VerifyWithdrawWithheldTokens.encode(proof_data) +/// Create a `VerifyWithdrawWithheldTokens` instruction. +pub fn verify_withdraw_withheld_tokens( + context_state_info: Option, + proof_data: &WithdrawWithheldTokensData, +) -> Instruction { + ProofInstruction::VerifyWithdrawWithheldTokens + .encode_verify_proof(context_state_info, proof_data) } -pub fn verify_transfer(proof_data: &TransferData) -> Instruction { - ProofInstruction::VerifyTransfer.encode(proof_data) +/// Create a `VerifyTransfer` instruction. +pub fn verify_transfer( + context_state_info: Option, + proof_data: &TransferData, +) -> Instruction { + ProofInstruction::VerifyTransfer.encode_verify_proof(context_state_info, proof_data) } -pub fn verify_transfer_with_fee(proof_data: &TransferWithFeeData) -> Instruction { - ProofInstruction::VerifyTransferWithFee.encode(proof_data) +/// Create a `VerifyTransferWithFee` instruction. +pub fn verify_transfer_with_fee( + context_state_info: Option, + proof_data: &TransferWithFeeData, +) -> Instruction { + ProofInstruction::VerifyTransferWithFee.encode_verify_proof(context_state_info, proof_data) } -pub fn verify_pubkey_validity(proof_data: &PubkeyValidityData) -> Instruction { - ProofInstruction::VerifyPubkeyValidity.encode(proof_data) +/// Create a `VerifyPubkeyValidity` instruction. +pub fn verify_pubkey_validity( + context_state_info: Option, + proof_data: &PubkeyValidityData, +) -> Instruction { + ProofInstruction::VerifyPubkeyValidity.encode_verify_proof(context_state_info, proof_data) +} + +impl ProofInstruction { + pub fn encode_verify_proof( + &self, + context_state_info: Option, + proof_data: &T, + ) -> Instruction + where + T: Pod + ZkProofData, + U: Pod, + { + let accounts = if let Some(context_state_info) = context_state_info { + vec![ + AccountMeta::new(*context_state_info.context_state_account, false), + AccountMeta::new_readonly(*context_state_info.context_state_authority, false), + ] + } else { + vec![] + }; + + let mut data = vec![ToPrimitive::to_u8(self).unwrap()]; + data.extend_from_slice(bytes_of(proof_data)); + + Instruction { + program_id: crate::zk_token_proof_program::id(), + accounts, + data, + } + } + + pub fn instruction_type(input: &[u8]) -> Option { + input + .first() + .and_then(|instruction| FromPrimitive::from_u8(*instruction)) + } + + pub fn proof_data(input: &[u8]) -> Option<&T> + where + T: Pod + ZkProofData, + U: Pod, + { + input + .get(1..) + .and_then(|data| bytemuck::try_from_bytes(data).ok()) + } } diff --git a/zk-token-sdk/src/zk_token_proof_state.rs b/zk-token-sdk/src/zk_token_proof_state.rs new file mode 100644 index 00000000000000..d95aa4f11ec1c3 --- /dev/null +++ b/zk-token-sdk/src/zk_token_proof_state.rs @@ -0,0 +1,72 @@ +use { + crate::{zk_token_elgamal::pod::PodProofType, zk_token_proof_instruction::ProofType}, + bytemuck::{bytes_of, Pod, Zeroable}, + num_traits::ToPrimitive, + solana_program::{ + instruction::{InstructionError, InstructionError::InvalidAccountData}, + pubkey::Pubkey, + }, + std::mem::size_of, +}; + +/// The proof context account state +#[derive(Clone, Copy, Debug, PartialEq)] +#[repr(C)] +pub struct ProofContextState { + /// The proof context authority that can close the account + pub context_state_authority: Pubkey, + /// The proof type for the context data + pub proof_type: PodProofType, + /// The proof context data + pub proof_context: T, +} + +// `bytemuck::Pod` cannot be derived for generic structs unless the struct is marked +// `repr(packed)`, which may cause unnecessary complications when referencing its fields. Directly +// mark `ProofContextState` as `Zeroable` and `Pod` since since none of its fields has an alignment +// requirement greater than 1 and therefore, guaranteed to be `packed`. +unsafe impl Zeroable for ProofContextState {} +unsafe impl Pod for ProofContextState {} + +impl ProofContextState { + pub fn encode( + context_state_authority: &Pubkey, + proof_type: ProofType, + proof_context: &T, + ) -> Vec { + let mut buf = Vec::with_capacity(size_of::()); + buf.extend_from_slice(context_state_authority.as_ref()); + buf.push(ToPrimitive::to_u8(&proof_type).unwrap()); + buf.extend_from_slice(bytes_of(proof_context)); + buf + } + + /// Interpret a slice as a `ProofContextState`. + /// + /// This function requires a generic parameter. To access only the generic-independent fields + /// in `ProofContextState` without a generic parameter, use + /// `ProofContextStateMeta::try_from_bytes` instead. + pub fn try_from_bytes(input: &[u8]) -> Result<&Self, InstructionError> { + bytemuck::try_from_bytes(input).map_err(|_| InvalidAccountData) + } +} + +/// The `ProofContextState` without the proof context itself. This struct exists to facilitate the +/// decoding of generic-independent fields in `ProofContextState`. +#[derive(Clone, Copy, Debug, PartialEq, Pod, Zeroable)] +#[repr(C)] +pub struct ProofContextStateMeta { + /// The proof context authority that can close the account + pub context_state_authority: Pubkey, + /// The proof type for the context data + pub proof_type: PodProofType, +} + +impl ProofContextStateMeta { + pub fn try_from_bytes(input: &[u8]) -> Result<&Self, InstructionError> { + input + .get(..size_of::()) + .and_then(|data| bytemuck::try_from_bytes(data).ok()) + .ok_or(InvalidAccountData) + } +}