Skip to content

Commit

Permalink
[zk-sdk] Add error types and zero ciphertext instruction (#1453)
Browse files Browse the repository at this point in the history
* add `ProofGenerationError` and `ProofVerificationError`

* add `zero_ciphertext` proof data module from zk-token-sdk verbatim

* clean up `zero_ciphertext` proof data module

* add `VerifyZeroCiphertext` instruction

* cargo fmt

* remove `ProofGenerationError::FeeCalculation`

* remove `ProofGenerationError::NotEnoughFunds`
  • Loading branch information
samkim-crypto authored May 23, 2024
1 parent 32cdbd6 commit 9035690
Show file tree
Hide file tree
Showing 4 changed files with 263 additions and 2 deletions.
45 changes: 44 additions & 1 deletion zk-sdk/src/elgamal_program/errors.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,50 @@
use thiserror::Error;
use {
crate::{
errors::ElGamalError,
range_proof::errors::{RangeProofGenerationError, RangeProofVerificationError},
sigma_proofs::errors::*,
},
thiserror::Error,
};

#[cfg(not(target_os = "solana"))]
#[derive(Error, Clone, Debug, Eq, PartialEq)]
pub enum ProofGenerationError {
#[error("illegal number of commitments")]
IllegalCommitmentLength,
#[error("illegal amount bit length")]
IllegalAmountBitLength,
#[error("invalid commitment")]
InvalidCommitment,
#[error("range proof generation failed")]
RangeProof(#[from] RangeProofGenerationError),
#[error("unexpected proof length")]
ProofLength,
}

#[derive(Error, Clone, Debug, Eq, PartialEq)]
pub enum ProofVerificationError {
#[error("range proof verification failed")]
RangeProof(#[from] RangeProofVerificationError),
#[error("sigma proof verification failed")]
SigmaProof(SigmaProofType, SigmaProofVerificationError),
#[error("ElGamal ciphertext or public key error")]
ElGamal(#[from] ElGamalError),
#[error("Invalid proof context")]
ProofContext,
#[error("illegal commitment length")]
IllegalCommitmentLength,
#[error("illegal amount bit length")]
IllegalAmountBitLength,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SigmaProofType {
ZeroCiphertext,
}

impl From<ZeroCiphertextProofVerificationError> for ProofVerificationError {
fn from(err: ZeroCiphertextProofVerificationError) -> Self {
Self::SigmaProof(SigmaProofType::ZeroCiphertext, err.0)
}
}
92 changes: 91 additions & 1 deletion zk-sdk/src/elgamal_program/instruction.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,10 @@
//! [`context-state`]: https://docs.solanalabs.com/runtime/zk-token-proof#context-data
use {
crate::elgamal_program::proof_data::ZkProofData,
bytemuck::{bytes_of, Pod},
num_derive::{FromPrimitive, ToPrimitive},
num_traits::ToPrimitive,
num_traits::{FromPrimitive, ToPrimitive},
solana_program::{
instruction::{AccountMeta, Instruction},
pubkey::Pubkey,
Expand All @@ -55,6 +57,22 @@ pub enum ProofInstruction {
/// None
///
CloseContextState,

/// Verify a zero-ciphertext proof.
///
/// A zero-ciphertext proof certifies that an ElGamal ciphertext encrypts the value zero.
///
/// 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. `ZeroCiphertextProofData` if proof is provided as instruction data
/// ii. `u32` byte offset if proof is provided as an account
///
VerifyZeroCiphertext,
}

/// Pubkeys associated with a context state account to be used as parameters to functions.
Expand Down Expand Up @@ -83,3 +101,75 @@ pub fn close_context_state(
data,
}
}

impl ProofInstruction {
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: crate::elgamal_program::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: crate::elgamal_program::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())
}
}
2 changes: 2 additions & 0 deletions zk-sdk/src/elgamal_program/proof_data/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,14 @@ use {

pub mod errors;
pub mod pod;
pub mod zero_ciphertext;

#[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,
ZeroCiphertext,
}

pub trait ZkProofData<T: Pod> {
Expand Down
126 changes: 126 additions & 0 deletions zk-sdk/src/elgamal_program/proof_data/zero_ciphertext.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,126 @@
//! The zero-ciphertext proof instruction.
//!
//! A zero-ciphertext proof is defined with respect to a twisted ElGamal ciphertext. The proof
//! certifies that a given ciphertext encrypts the message 0 in the field (`Scalar::zero()`). To
//! generate the proof, a prover must provide the decryption key for the ciphertext.
use {
crate::{
elgamal_program::{
errors::{ProofGenerationError, ProofVerificationError},
proof_data::{ProofType, ZkProofData},
},
encryption::pod::elgamal::{PodElGamalCiphertext, PodElGamalPubkey},
sigma_proofs::pod::PodZeroCiphertextProof,
},
bytemuck::{bytes_of, Pod, Zeroable},
};
#[cfg(not(target_os = "solana"))]
use {
crate::{
encryption::elgamal::{ElGamalCiphertext, ElGamalKeypair},
sigma_proofs::zero_ciphertext::ZeroCiphertextProof,
},
merlin::Transcript,
std::convert::TryInto,
};

/// The instruction data that is needed for the `ProofInstruction::ZeroCiphertext` instruction.
///
/// It includes the cryptographic proof as well as the context data information needed to verify
/// the proof.
#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct ZeroCiphertextProofData {
/// The context data for the zero-ciphertext proof
pub context: ZeroCiphertextProofContext, // 96 bytes

/// Proof that the ciphertext is zero
pub proof: PodZeroCiphertextProof, // 96 bytes
}

/// The context data needed to verify a zero-ciphertext proof.
#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct ZeroCiphertextProofContext {
/// The ElGamal pubkey associated with the ElGamal ciphertext
pub pubkey: PodElGamalPubkey, // 32 bytes

/// The ElGamal ciphertext that encrypts zero
pub ciphertext: PodElGamalCiphertext, // 64 bytes
}

#[cfg(not(target_os = "solana"))]
impl ZeroCiphertextProofData {
pub fn new(
keypair: &ElGamalKeypair,
ciphertext: &ElGamalCiphertext,
) -> Result<Self, ProofGenerationError> {
let pod_pubkey = PodElGamalPubkey(keypair.pubkey().into());
let pod_ciphertext = PodElGamalCiphertext(ciphertext.to_bytes());

let context = ZeroCiphertextProofContext {
pubkey: pod_pubkey,
ciphertext: pod_ciphertext,
};

let mut transcript = context.new_transcript();
let proof = ZeroCiphertextProof::new(keypair, ciphertext, &mut transcript).into();

Ok(ZeroCiphertextProofData { context, proof })
}
}

impl ZkProofData<ZeroCiphertextProofContext> for ZeroCiphertextProofData {
const PROOF_TYPE: ProofType = ProofType::ZeroCiphertext;

fn context_data(&self) -> &ZeroCiphertextProofContext {
&self.context
}

#[cfg(not(target_os = "solana"))]
fn verify_proof(&self) -> Result<(), ProofVerificationError> {
let mut transcript = self.context.new_transcript();
let pubkey = self.context.pubkey.try_into()?;
let ciphertext = self.context.ciphertext.try_into()?;
let proof: ZeroCiphertextProof = self.proof.try_into()?;
proof
.verify(&pubkey, &ciphertext, &mut transcript)
.map_err(|e| e.into())
}
}

#[allow(non_snake_case)]
#[cfg(not(target_os = "solana"))]
impl ZeroCiphertextProofContext {
fn new_transcript(&self) -> Transcript {
let mut transcript = Transcript::new(b"zero-ciphertext-instruction");

transcript.append_message(b"pubkey", bytes_of(&self.pubkey));
transcript.append_message(b"ciphertext", bytes_of(&self.ciphertext));

transcript
}
}

#[cfg(test)]
mod test {
use super::*;

#[test]
fn test_zero_ciphertext_proof_instruction_correctness() {
let keypair = ElGamalKeypair::new_rand();

// general case: encryption of 0
let ciphertext = keypair.pubkey().encrypt(0_u64);
let zero_ciphertext_proof_data =
ZeroCiphertextProofData::new(&keypair, &ciphertext).unwrap();
assert!(zero_ciphertext_proof_data.verify_proof().is_ok());

// general case: encryption of > 0
let ciphertext = keypair.pubkey().encrypt(1_u64);
let zero_ciphertext_proof_data =
ZeroCiphertextProofData::new(&keypair, &ciphertext).unwrap();
assert!(zero_ciphertext_proof_data.verify_proof().is_err());
}
}

0 comments on commit 9035690

Please sign in to comment.