[zk-token-sdk] Add sigma proof program

zk-token-sdk/src/lib.rs

@@ -31,9 +31,9 @@ mod sigma_proofs;
 #[cfg(not(target_os = "solana"))]
 mod transcript;
 
-// TODO: re-organize visibility
 pub mod curve25519;
 pub mod instruction;
+pub mod programs;
 pub mod zk_token_elgamal;
 pub mod zk_token_proof_instruction;
 pub mod zk_token_proof_program;

zk-token-sdk/src/programs/mod.rs

@@ -0,0 +1,2 @@
+pub mod sigma_proof;
+pub mod state;

zk-token-sdk/src/programs/sigma_proof/instruction.rs

@@ -0,0 +1,138 @@
+use {
+    crate::{
+        instruction::{Pod, ZkProofData},
+        programs::{sigma_proof::id, state::ContextStateInfo},
+    },
+    bytemuck::bytes_of,
+    num_derive::{FromPrimitive, ToPrimitive},
+    num_traits::{FromPrimitive, ToPrimitive},
+    solana_program::{
+        instruction::{AccountMeta, Instruction},
+        pubkey::Pubkey,
+    },
+};
+
+#[derive(Clone, Copy, Debug, FromPrimitive, ToPrimitive, PartialEq, Eq)]
+#[repr(u8)]
+pub enum SigmaProofInstruction {
+    /// 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 public key validity zero-knowledge proof.
+    ///
+    /// A public key validity proof certifies that an ElGamal public key is well-formed and the
+    /// prover knows the corresponding secret key.
+    ///
+    /// Accounts expected by this instruction:
+    ///
+    ///   0. `[]` (Optional) Account to read the proof from
+    ///   1. `[writable]` (Optional) The proof context account
+    ///   2. `[]` (Optional) The proof context account owner
+    ///
+    /// The instruction expects either:
+    ///   i. `PubkeyValidityData` if proof is provided as instruction data
+    ///   ii. `u32` byte offset if proof is provided as an account
+    ///
+    VerifyPubkeyValidity,
+}
+
+/// 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(&SigmaProofInstruction::CloseContextState).unwrap()];
+
+    Instruction {
+        program_id: id(),
+        accounts,
+        data,
+    }
+}
+
+impl SigmaProofInstruction {
+    pub fn encode_verify_proof<T, U>(
+        &self,
+        context_state_info: Option<ContextStateInfo>,
+        proof_data: &T,
+    ) -> Instruction
+    where
+        T: Pod + ZkProofData<U>,
+        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: id(),
+            accounts,
+            data,
+        }
+    }
+
+    pub fn encode_verify_proof_from_account(
+        &self,
+        context_state_info: Option<ContextStateInfo>,
+        proof_account: &Pubkey,
+        offset: u32,
+    ) -> Instruction {
+        let accounts = if let Some(context_state_info) = context_state_info {
+            vec![
+                AccountMeta::new(*proof_account, false),
+                AccountMeta::new(*context_state_info.context_state_account, false),
+                AccountMeta::new_readonly(*context_state_info.context_state_authority, false),
+            ]
+        } else {
+            vec![AccountMeta::new(*proof_account, false)]
+        };
+
+        let mut data = vec![ToPrimitive::to_u8(self).unwrap()];
+        data.extend_from_slice(&offset.to_le_bytes());
+
+        Instruction {
+            program_id: id(),
+            accounts,
+            data,
+        }
+    }
+
+    pub fn instruction_type(input: &[u8]) -> Option<Self> {
+        input
+            .first()
+            .and_then(|instruction| FromPrimitive::from_u8(*instruction))
+    }
+
+    pub fn proof_data<T, U>(input: &[u8]) -> Option<&T>
+    where
+        T: Pod + ZkProofData<U>,
+        U: Pod,
+    {
+        input
+            .get(1..)
+            .and_then(|data| bytemuck::try_from_bytes(data).ok())
+    }
+}

zk-token-sdk/src/programs/sigma_proof/mod.rs

@@ -0,0 +1,6 @@
+//! The Sigma Proof program.
+
+pub mod instruction;
+
+// Program Id of the Sigma Proof program
+solana_program::declare_id!("SigmaProof111111111111111111111111111111111");

zk-token-sdk/src/programs/state.rs

@@ -0,0 +1,79 @@
+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 context state account addresses intended 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,
+}
+
+/// 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)
+    }
+}