From a2af4caab1d83b17eb9cf7c57dc2132e475e9e15 Mon Sep 17 00:00:00 2001 From: samkim-crypto Date: Wed, 8 May 2024 20:38:52 +0900 Subject: [PATCH] [zk-sdk] Add the `program` module (#1188) * add `program` module * add `instruction` module * add `proof_data` module * add `state` module * update program id to `ZkE1Gama1Proof11111111111111111111111111111` * rename `program` module to `elgamal-program` * move `ProofVerificationError` to `elgamal-program` module --- Cargo.lock | 3 + zk-sdk/Cargo.toml | 3 + zk-sdk/src/elgamal_program/errors.rs | 7 ++ zk-sdk/src/elgamal_program/instruction.rs | 85 +++++++++++++++++++ zk-sdk/src/elgamal_program/mod.rs | 16 ++++ .../src/elgamal_program/proof_data/errors.rs | 15 ++++ zk-sdk/src/elgamal_program/proof_data/mod.rs | 24 ++++++ zk-sdk/src/elgamal_program/proof_data/pod.rs | 21 +++++ zk-sdk/src/elgamal_program/state.rs | 72 ++++++++++++++++ zk-sdk/src/lib.rs | 1 + 10 files changed, 247 insertions(+) create mode 100644 zk-sdk/src/elgamal_program/errors.rs create mode 100644 zk-sdk/src/elgamal_program/instruction.rs create mode 100644 zk-sdk/src/elgamal_program/mod.rs create mode 100644 zk-sdk/src/elgamal_program/proof_data/errors.rs create mode 100644 zk-sdk/src/elgamal_program/proof_data/mod.rs create mode 100644 zk-sdk/src/elgamal_program/proof_data/pod.rs create mode 100644 zk-sdk/src/elgamal_program/state.rs diff --git a/Cargo.lock b/Cargo.lock index 1ff05bdd52dd1d..4b5f146a683f5c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7709,10 +7709,13 @@ dependencies = [ "aes-gcm-siv", "base64 0.22.0", "bincode", + "bytemuck", "curve25519-dalek", "itertools", "lazy_static", "merlin", + "num-derive", + "num-traits", "rand 0.7.3", "serde", "serde_json", diff --git a/zk-sdk/Cargo.toml b/zk-sdk/Cargo.toml index 34d817425b2536..974c47b7e3ecc7 100644 --- a/zk-sdk/Cargo.toml +++ b/zk-sdk/Cargo.toml @@ -11,6 +11,9 @@ edition = { workspace = true } [dependencies] base64 = { workspace = true } +bytemuck = { workspace = true } +num-derive = { workspace = true } +num-traits = { workspace = true } solana-program = { workspace = true } thiserror = { workspace = true } diff --git a/zk-sdk/src/elgamal_program/errors.rs b/zk-sdk/src/elgamal_program/errors.rs new file mode 100644 index 00000000000000..514140d56d2c87 --- /dev/null +++ b/zk-sdk/src/elgamal_program/errors.rs @@ -0,0 +1,7 @@ +use thiserror::Error; + +#[derive(Error, Clone, Debug, Eq, PartialEq)] +pub enum ProofVerificationError { + #[error("Invalid proof context")] + ProofContext, +} diff --git a/zk-sdk/src/elgamal_program/instruction.rs b/zk-sdk/src/elgamal_program/instruction.rs new file mode 100644 index 00000000000000..4372dcf024b538 --- /dev/null +++ b/zk-sdk/src/elgamal_program/instruction.rs @@ -0,0 +1,85 @@ +//! Instructions provided by the [`ZK ElGamal proof`] program. +//! +//! There are two types of instructions in the proof program: proof verification instructions and +//! the `CloseContextState` instruction. +//! +//! Each proof verification instruction verifies a certain type of zero-knowledge proof. These +//! instructions are processed by the program in two steps: +//! 1. The program verifies the zero-knowledge proof. +//! 2. The program optionally stores the context component of the zero-knowledge proof to a +//! dedicated [`context-state`] account. +//! +//! In step 1, the zero-knowledge proof can either be included directly as the instruction data or +//! pre-written to an account. The progrma determines whether the proof is provided as instruction +//! data or pre-written to an account by inspecting the length of the data. If the instruction data +//! is exactly 5 bytes (instruction discriminator + unsigned 32-bit integer), then the program +//! assumes that the first account provided with the instruction contains the zero-knowledge proof +//! and verifies the account data at the offset specified in the instruction data. Otherwise, the +//! program assumes that the zero-knowledge proof is provided as part of the instruction data. +//! +//! In step 2, the program determines whether to create a context-state account by inspecting the +//! number of accounts provided with the instruction. If two additional accounts are provided with +//! the instruction after verifying the zero-knowledge proof, then the program writes the context +//! data to the specified context-state account. +//! +//! NOTE: A context-state account must be pre-allocated to the exact size of the context data that +//! is expected for a proof type before it is included as part of a proof verification instruction. +//! +//! The `CloseContextState` instruction closes a context state account. A transaction containing +//! this instruction must be signed by the context account's owner. This instruction can be used by +//! the account owner to relcaim lamports for storage. +//! +//! [`ZK ElGamal proof`]: https://docs.solanalabs.com/runtime/zk-token-proof +//! [`context-state`]: https://docs.solanalabs.com/runtime/zk-token-proof#context-data + +use { + num_derive::{FromPrimitive, ToPrimitive}, + num_traits::ToPrimitive, + solana_program::{ + instruction::{AccountMeta, Instruction}, + pubkey::Pubkey, + }, +}; + +#[derive(Clone, Copy, Debug, FromPrimitive, ToPrimitive, PartialEq, Eq)] +#[repr(u8)] +pub enum ProofInstruction { + /// 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, +} + +/// 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, +} + +/// 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), + ]; + + let data = vec![ToPrimitive::to_u8(&ProofInstruction::CloseContextState).unwrap()]; + + Instruction { + program_id: crate::elgamal_program::id(), + accounts, + data, + } +} diff --git a/zk-sdk/src/elgamal_program/mod.rs b/zk-sdk/src/elgamal_program/mod.rs new file mode 100644 index 00000000000000..b8adc8744893dc --- /dev/null +++ b/zk-sdk/src/elgamal_program/mod.rs @@ -0,0 +1,16 @@ +//! The native ZK ElGamal proof program. +//! +//! The program verifies a number of zero-knowledge proofs that are tailored to work with Pedersen +//! commitments and ElGamal encryption over the elliptic curve curve25519. A general overview of +//! the program as well as the technical details of some of the proof instructions can be found in +//! the [`ZK ElGamal proof`] documentation. +//! +//! [`ZK ElGamal proof`]: https://docs.solanalabs.com/runtime/zk-token-proof + +pub mod errors; +pub mod instruction; +pub mod proof_data; +pub mod state; + +// Program Id of the ZK ElGamal Proof program +solana_program::declare_id!("ZkE1Gama1Proof11111111111111111111111111111"); diff --git a/zk-sdk/src/elgamal_program/proof_data/errors.rs b/zk-sdk/src/elgamal_program/proof_data/errors.rs new file mode 100644 index 00000000000000..423981a1a81499 --- /dev/null +++ b/zk-sdk/src/elgamal_program/proof_data/errors.rs @@ -0,0 +1,15 @@ +use thiserror::Error; + +#[derive(Error, Clone, Debug, Eq, PartialEq)] +pub enum ProofDataError { + #[error("decryption error")] + Decryption, + #[error("missing ciphertext")] + MissingCiphertext, + #[error("illegal amount bit length")] + IllegalAmountBitLength, + #[error("arithmetic overflow")] + Overflow, + #[error("invalid proof type")] + InvalidProofType, +} diff --git a/zk-sdk/src/elgamal_program/proof_data/mod.rs b/zk-sdk/src/elgamal_program/proof_data/mod.rs new file mode 100644 index 00000000000000..da3989d5ef74ff --- /dev/null +++ b/zk-sdk/src/elgamal_program/proof_data/mod.rs @@ -0,0 +1,24 @@ +use { + crate::elgamal_program::errors::ProofVerificationError, + bytemuck::Pod, + num_derive::{FromPrimitive, ToPrimitive}, +}; + +pub mod errors; +pub mod pod; + +#[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, +} + +pub trait ZkProofData { + const PROOF_TYPE: ProofType; + + fn context_data(&self) -> &T; + + #[cfg(not(target_os = "solana"))] + fn verify_proof(&self) -> Result<(), ProofVerificationError>; +} diff --git a/zk-sdk/src/elgamal_program/proof_data/pod.rs b/zk-sdk/src/elgamal_program/proof_data/pod.rs new file mode 100644 index 00000000000000..7d3f346a684d59 --- /dev/null +++ b/zk-sdk/src/elgamal_program/proof_data/pod.rs @@ -0,0 +1,21 @@ +use { + crate::elgamal_program::proof_data::{errors::ProofDataError, ProofType}, + bytemuck::{Pod, Zeroable}, + num_traits::{FromPrimitive, ToPrimitive}, +}; + +#[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 = ProofDataError; + + fn try_from(pod: PodProofType) -> Result { + FromPrimitive::from_u8(pod.0).ok_or(Self::Error::InvalidProofType) + } +} diff --git a/zk-sdk/src/elgamal_program/state.rs b/zk-sdk/src/elgamal_program/state.rs new file mode 100644 index 00000000000000..fe9532f588dd52 --- /dev/null +++ b/zk-sdk/src/elgamal_program/state.rs @@ -0,0 +1,72 @@ +use { + crate::elgamal_program::proof_data::{pod::PodProofType, 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) + } +} diff --git a/zk-sdk/src/lib.rs b/zk-sdk/src/lib.rs index 2c7e591e02a2a3..8e398406dc0638 100644 --- a/zk-sdk/src/lib.rs +++ b/zk-sdk/src/lib.rs @@ -19,6 +19,7 @@ // `clippy::op_ref` is turned off to prevent clippy from warning that this is not idiomatic code. #![allow(clippy::arithmetic_side_effects, clippy::op_ref)] +pub mod elgamal_program; #[cfg(not(target_os = "solana"))] pub mod encryption; pub mod errors;