Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[zk-sdk] Add error types and zero ciphertext instruction #1453

Merged
merged 7 commits into from
May 23, 2024
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 48 additions & 1 deletion zk-sdk/src/elgamal_program/errors.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,54 @@
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("not enough funds in account")]
NotEnoughFunds,
#[error("transfer fee calculation error")]
FeeCalculation,
#[error("illegal number of commitments")]

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should these errors be changed or removed?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh great point! I removed it. Thanks!

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok cool! Can you also get rid of NotEnoughFunds?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh oops, removed. Thanks!

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());
}
}
Loading