solana_zk_token_sdk/instruction/
withdraw.rs

1#[cfg(not(target_os = "solana"))]
2use {
3    crate::{
4        encryption::{
5            elgamal::{ElGamal, ElGamalCiphertext, ElGamalKeypair, ElGamalPubkey},
6            pedersen::{Pedersen, PedersenCommitment},
7        },
8        errors::{ProofGenerationError, ProofVerificationError},
9        range_proof::RangeProof,
10        sigma_proofs::ciphertext_commitment_equality_proof::CiphertextCommitmentEqualityProof,
11        transcript::TranscriptProtocol,
12    },
13    merlin::Transcript,
14    std::convert::TryInto,
15};
16use {
17    crate::{
18        instruction::{ProofType, ZkProofData},
19        zk_token_elgamal::pod,
20    },
21    bytemuck_derive::{Pod, Zeroable},
22};
23
24#[cfg(not(target_os = "solana"))]
25const WITHDRAW_AMOUNT_BIT_LENGTH: usize = 64;
26
27/// The instruction data that is needed for the `ProofInstruction::VerifyWithdraw` instruction.
28///
29/// It includes the cryptographic proof as well as the context data information needed to verify
30/// the proof.
31#[derive(Clone, Copy, Pod, Zeroable)]
32#[repr(C)]
33pub struct WithdrawData {
34    /// The context data for the withdraw proof
35    pub context: WithdrawProofContext, // 128 bytes
36
37    /// Range proof
38    pub proof: WithdrawProof, // 736 bytes
39}
40
41/// The context data needed to verify a withdraw proof.
42#[derive(Clone, Copy, Pod, Zeroable)]
43#[repr(C)]
44pub struct WithdrawProofContext {
45    /// The source account ElGamal pubkey
46    pub pubkey: pod::ElGamalPubkey, // 32 bytes
47
48    /// The source account available balance *after* the withdraw (encrypted by
49    /// `source_pk`
50    pub final_ciphertext: pod::ElGamalCiphertext, // 64 bytes
51}
52
53#[cfg(not(target_os = "solana"))]
54impl WithdrawData {
55    pub fn new(
56        amount: u64,
57        keypair: &ElGamalKeypair,
58        current_balance: u64,
59        current_ciphertext: &ElGamalCiphertext,
60    ) -> Result<Self, ProofGenerationError> {
61        // subtract withdraw amount from current balance
62        //
63        // errors if current_balance < amount
64        let final_balance = current_balance
65            .checked_sub(amount)
66            .ok_or(ProofGenerationError::NotEnoughFunds)?;
67
68        // encode withdraw amount as an ElGamal ciphertext and subtract it from
69        // current source balance
70        let final_ciphertext = current_ciphertext - &ElGamal::encode(amount);
71
72        let pod_pubkey = pod::ElGamalPubkey(keypair.pubkey().into());
73        let pod_final_ciphertext: pod::ElGamalCiphertext = final_ciphertext.into();
74
75        let context = WithdrawProofContext {
76            pubkey: pod_pubkey,
77            final_ciphertext: pod_final_ciphertext,
78        };
79
80        let mut transcript = context.new_transcript();
81        let proof = WithdrawProof::new(keypair, final_balance, &final_ciphertext, &mut transcript)?;
82
83        Ok(Self { context, proof })
84    }
85}
86
87impl ZkProofData<WithdrawProofContext> for WithdrawData {
88    const PROOF_TYPE: ProofType = ProofType::Withdraw;
89
90    fn context_data(&self) -> &WithdrawProofContext {
91        &self.context
92    }
93
94    #[cfg(not(target_os = "solana"))]
95    fn verify_proof(&self) -> Result<(), ProofVerificationError> {
96        let mut transcript = self.context.new_transcript();
97
98        let elgamal_pubkey = self.context.pubkey.try_into()?;
99        let final_balance_ciphertext = self.context.final_ciphertext.try_into()?;
100        self.proof
101            .verify(&elgamal_pubkey, &final_balance_ciphertext, &mut transcript)
102    }
103}
104
105#[allow(non_snake_case)]
106#[cfg(not(target_os = "solana"))]
107impl WithdrawProofContext {
108    fn new_transcript(&self) -> Transcript {
109        let mut transcript = Transcript::new(b"WithdrawProof");
110
111        transcript.append_pubkey(b"pubkey", &self.pubkey);
112        transcript.append_ciphertext(b"ciphertext", &self.final_ciphertext);
113
114        transcript
115    }
116}
117
118/// The withdraw proof.
119///
120/// It contains a ciphertext-commitment equality proof and a 64-bit range proof.
121#[derive(Clone, Copy, Pod, Zeroable)]
122#[repr(C)]
123#[allow(non_snake_case)]
124pub struct WithdrawProof {
125    /// New Pedersen commitment
126    pub commitment: pod::PedersenCommitment,
127
128    /// Associated equality proof
129    pub equality_proof: pod::CiphertextCommitmentEqualityProof,
130
131    /// Associated range proof
132    pub range_proof: pod::RangeProofU64, // 672 bytes
133}
134
135#[allow(non_snake_case)]
136#[cfg(not(target_os = "solana"))]
137impl WithdrawProof {
138    pub fn new(
139        keypair: &ElGamalKeypair,
140        final_balance: u64,
141        final_ciphertext: &ElGamalCiphertext,
142        transcript: &mut Transcript,
143    ) -> Result<Self, ProofGenerationError> {
144        // generate a Pedersen commitment for `final_balance`
145        let (commitment, opening) = Pedersen::new(final_balance);
146        let pod_commitment: pod::PedersenCommitment = commitment.into();
147
148        transcript.append_commitment(b"commitment", &pod_commitment);
149
150        // generate equality_proof
151        let equality_proof = CiphertextCommitmentEqualityProof::new(
152            keypair,
153            final_ciphertext,
154            &opening,
155            final_balance,
156            transcript,
157        );
158
159        let range_proof =
160            RangeProof::new(vec![final_balance], vec![64], vec![&opening], transcript)?;
161
162        Ok(Self {
163            commitment: pod_commitment,
164            equality_proof: equality_proof.into(),
165            range_proof: range_proof
166                .try_into()
167                .map_err(|_| ProofGenerationError::ProofLength)?,
168        })
169    }
170
171    pub fn verify(
172        &self,
173        pubkey: &ElGamalPubkey,
174        final_ciphertext: &ElGamalCiphertext,
175        transcript: &mut Transcript,
176    ) -> Result<(), ProofVerificationError> {
177        transcript.append_commitment(b"commitment", &self.commitment);
178
179        let commitment: PedersenCommitment = self.commitment.try_into()?;
180        let equality_proof: CiphertextCommitmentEqualityProof = self.equality_proof.try_into()?;
181        let range_proof: RangeProof = self.range_proof.try_into()?;
182
183        // verify equality proof
184        equality_proof.verify(pubkey, final_ciphertext, &commitment, transcript)?;
185
186        // verify range proof
187        range_proof.verify(
188            vec![&commitment],
189            vec![WITHDRAW_AMOUNT_BIT_LENGTH],
190            transcript,
191        )?;
192
193        Ok(())
194    }
195}
196
197#[cfg(test)]
198mod test {
199    use {super::*, crate::encryption::elgamal::ElGamalKeypair};
200
201    #[test]
202    fn test_withdraw_correctness() {
203        // generate and verify proof for the proper setting
204        let keypair = ElGamalKeypair::new_rand();
205
206        let current_balance: u64 = 77;
207        let current_ciphertext = keypair.pubkey().encrypt(current_balance);
208
209        let withdraw_amount: u64 = 55;
210
211        let data = WithdrawData::new(
212            withdraw_amount,
213            &keypair,
214            current_balance,
215            &current_ciphertext,
216        )
217        .unwrap();
218        assert!(data.verify_proof().is_ok());
219
220        // generate and verify proof with wrong balance
221        let wrong_balance: u64 = 99;
222        let data = WithdrawData::new(
223            withdraw_amount,
224            &keypair,
225            wrong_balance,
226            &current_ciphertext,
227        )
228        .unwrap();
229        assert!(data.verify_proof().is_err());
230    }
231}
solana_zk_token_sdk/instruction/withdraw.rs

solana_zk_token_sdk/instruction/
withdraw.rs
