[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

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 }
 

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,
+}

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,
+    }
+}

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");

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,
+}

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<T: Pod> {
+    const PROOF_TYPE: ProofType;
+
+    fn context_data(&self) -> &T;
+
+    #[cfg(not(target_os = "solana"))]
+    fn verify_proof(&self) -> Result<(), ProofVerificationError>;
+}

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<ProofType> for PodProofType {
+    fn from(proof_type: ProofType) -> Self {
+        Self(ToPrimitive::to_u8(&proof_type).unwrap())
+    }
+}
+impl TryFrom<PodProofType> for ProofType {
+    type Error = ProofDataError;
+
+    fn try_from(pod: PodProofType) -> Result<Self, Self::Error> {
+        FromPrimitive::from_u8(pod.0).ok_or(Self::Error::InvalidProofType)
+    }
+}

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<T: Pod> {
+    /// 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<T: Pod> Zeroable for ProofContextState<T> {}
+unsafe impl<T: Pod> Pod for ProofContextState<T> {}
+
+impl<T: Pod> ProofContextState<T> {
+    pub fn encode(
+        context_state_authority: &Pubkey,
+        proof_type: ProofType,
+        proof_context: &T,
+    ) -> Vec<u8> {
+        let mut buf = Vec::with_capacity(size_of::<Self>());
+        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::<ProofContextStateMeta>())
+            .and_then(|data| bytemuck::try_from_bytes(data).ok())
+            .ok_or(InvalidAccountData)
+    }
+}

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;