Add additional zkp for fee (solana-labs#23112)

* zk-token-sdk: add equality proof for fee

* zk-token-sdk: tweak some naming conventions for readability

* zk-token-sdk: add verify withdraw withheld instruction

* zk-token-sdk: add test for withdraw withheld verification

* zk-token-sdk: more renaming of variables for readability

* zk-token-sdk: cargo fmt

* zk-token-sdk: minor

* zk-token-sdk: resolve bpf compilation warnings

* zk-token-sdk: minor update to doc

programs/zk-token-proof/src/lib.rs

@@ -49,6 +49,10 @@ pub fn process_instruction(
             ic_msg!(invoke_context, "VerifyWithdraw");
             verify::<WithdrawData>(input, invoke_context)
         }
+        ProofInstruction::VerifyWithdrawWithheldTokens => {
+            ic_msg!(invoke_context, "VerifyWithdraw");
+            verify::<WithdrawData>(input, invoke_context)
+        }
         ProofInstruction::VerifyTransfer => {
             ic_msg!(invoke_context, "VerifyTransfer");
             verify::<TransferData>(input, invoke_context)

zk-token-sdk/src/instruction/close_account.rs

@@ -98,7 +98,7 @@ impl CloseAccountProof {
     ) -> Self {
         let proof = ZeroBalanceProof::new(keypair, ciphertext, transcript);
 
-        CloseAccountProof {
+        Self {
             proof: proof.into(),
         }
     }

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

@@ -2,6 +2,7 @@ pub mod close_account;
 pub mod transfer;
 pub mod transfer_with_fee;
 pub mod withdraw;
+pub mod withdraw_withheld;
 
 #[cfg(not(target_arch = "bpf"))]
 use {
@@ -17,6 +18,7 @@ use {
 pub use {
     close_account::CloseAccountData, transfer::TransferData,
     transfer_with_fee::TransferWithFeeData, withdraw::WithdrawData,
+    withdraw_withheld::WithdrawWithheldTokensData,
 };
 
 /// Constant for 2^32

zk-token-sdk/src/instruction/transfer.rs

@@ -15,22 +15,31 @@ use {
         errors::ProofError,
         instruction::{combine_u32_ciphertexts, split_u64_into_u32, Role, Verifiable, TWO_32},
         range_proof::RangeProof,
-        sigma_proofs::{equality_proof::EqualityProof, validity_proof::AggregatedValidityProof},
+        sigma_proofs::{
+            equality_proof::CtxtCommEqualityProof, validity_proof::AggregatedValidityProof,
+        },
         transcript::TranscriptProtocol,
     },
     arrayref::{array_ref, array_refs},
     merlin::Transcript,
     std::convert::TryInto,
 };
 
+#[cfg(not(target_arch = "bpf"))]
+const TRANSFER_SOURCE_AMOUNT_BIT_LENGTH: usize = 64;
+#[cfg(not(target_arch = "bpf"))]
+const TRANSFER_AMOUNT_LO_BIT_LENGTH: usize = 32;
+#[cfg(not(target_arch = "bpf"))]
+const TRANSFER_AMOUNT_HI_BIT_LENGTH: usize = 32;
+
 #[derive(Clone)]
 #[repr(C)]
 #[cfg(not(target_arch = "bpf"))]
 pub struct TransferAmountEncryption {
     pub commitment: PedersenCommitment,
-    pub source: DecryptHandle,
-    pub dest: DecryptHandle,
-    pub auditor: DecryptHandle,
+    pub handle_source: DecryptHandle,
+    pub handle_dest: DecryptHandle,
+    pub handle_auditor: DecryptHandle,
 }
 
 #[cfg(not(target_arch = "bpf"))]
@@ -44,9 +53,9 @@ impl TransferAmountEncryption {
         let (commitment, opening) = Pedersen::new(amount);
         let transfer_amount_encryption = Self {
             commitment,
-            source: pubkey_source.decrypt_handle(&opening),
-            dest: pubkey_dest.decrypt_handle(&opening),
-            auditor: pubkey_auditor.decrypt_handle(&opening),
+            handle_source: pubkey_source.decrypt_handle(&opening),
+            handle_dest: pubkey_dest.decrypt_handle(&opening),
+            handle_auditor: pubkey_auditor.decrypt_handle(&opening),
         };
 
         (transfer_amount_encryption, opening)
@@ -55,9 +64,9 @@ impl TransferAmountEncryption {
     pub fn to_pod(&self) -> pod::TransferAmountEncryption {
         pod::TransferAmountEncryption {
             commitment: self.commitment.into(),
-            source: self.source.into(),
-            dest: self.dest.into(),
-            auditor: self.auditor.into(),
+            handle_source: self.handle_source.into(),
+            handle_dest: self.handle_dest.into(),
+            handle_auditor: self.handle_auditor.into(),
         }
     }
 }
@@ -113,22 +122,22 @@ impl TransferData {
 
         let transfer_amount_lo_source = ElGamalCiphertext {
             commitment: ciphertext_lo.commitment,
-            handle: ciphertext_lo.source,
+            handle: ciphertext_lo.handle_source,
         };
 
         let transfer_amount_hi_source = ElGamalCiphertext {
             commitment: ciphertext_hi.commitment,
-            handle: ciphertext_hi.source,
+            handle: ciphertext_hi.handle_source,
         };
 
         let ciphertext_new_source = ciphertext_old_source
             - combine_u32_ciphertexts(&transfer_amount_lo_source, &transfer_amount_hi_source);
 
         // generate transcript and append all public inputs
         let pod_transfer_pubkeys = pod::TransferPubkeys {
-            source: keypair_source.public.into(),
-            dest: (*pubkey_dest).into(),
-            auditor: (*pubkey_auditor).into(),
+            pubkey_source: keypair_source.public.into(),
+            pubkey_dest: (*pubkey_dest).into(),
+            pubkey_auditor: (*pubkey_auditor).into(),
         };
         let pod_ciphertext_lo: pod::TransferAmountEncryption = ciphertext_lo.into();
         let pod_ciphertext_hi: pod::TransferAmountEncryption = ciphertext_hi.into();
@@ -165,9 +174,9 @@ impl TransferData {
         let ciphertext_lo: TransferAmountEncryption = self.ciphertext_lo.try_into()?;
 
         let handle_lo = match role {
-            Role::Source => ciphertext_lo.source,
-            Role::Dest => ciphertext_lo.dest,
-            Role::Auditor => ciphertext_lo.auditor,
+            Role::Source => ciphertext_lo.handle_source,
+            Role::Dest => ciphertext_lo.handle_dest,
+            Role::Auditor => ciphertext_lo.handle_auditor,
         };
 
         Ok(ElGamalCiphertext {
@@ -181,9 +190,9 @@ impl TransferData {
         let ciphertext_hi: TransferAmountEncryption = self.ciphertext_hi.try_into()?;
 
         let handle_hi = match role {
-            Role::Source => ciphertext_hi.source,
-            Role::Dest => ciphertext_hi.dest,
-            Role::Auditor => ciphertext_hi.auditor,
+            Role::Source => ciphertext_hi.handle_source,
+            Role::Dest => ciphertext_hi.handle_dest,
+            Role::Auditor => ciphertext_hi.handle_auditor,
         };
 
         Ok(ElGamalCiphertext {
@@ -247,7 +256,7 @@ pub struct TransferProof {
     pub commitment_new_source: pod::PedersenCommitment,
 
     /// Associated equality proof
-    pub equality_proof: pod::EqualityProof,
+    pub equality_proof: pod::CtxtCommEqualityProof,
 
     /// Associated ciphertext validity proof
     pub validity_proof: pod::AggregatedValidityProof,
@@ -267,19 +276,19 @@ impl TransferProof {
     ) -> Transcript {
         let mut transcript = Transcript::new(b"transfer-proof");
 
-        transcript.append_pubkey(b"pubkey_source", &transfer_pubkeys.source);
-        transcript.append_pubkey(b"pubkey_dest", &transfer_pubkeys.dest);
-        transcript.append_pubkey(b"pubkey_auditor", &transfer_pubkeys.auditor);
+        transcript.append_pubkey(b"pubkey-source", &transfer_pubkeys.pubkey_source);
+        transcript.append_pubkey(b"pubkey-dest", &transfer_pubkeys.pubkey_dest);
+        transcript.append_pubkey(b"pubkey-auditor", &transfer_pubkeys.pubkey_auditor);
 
         transcript.append_commitment(b"comm-lo-amount", &ciphertext_lo.commitment);
-        transcript.append_handle(b"handle-lo-source", &ciphertext_lo.source);
-        transcript.append_handle(b"handle-lo-dest", &ciphertext_lo.dest);
-        transcript.append_handle(b"handle-lo-auditor", &ciphertext_lo.auditor);
+        transcript.append_handle(b"handle-lo-source", &ciphertext_lo.handle_source);
+        transcript.append_handle(b"handle-lo-dest", &ciphertext_lo.handle_dest);
+        transcript.append_handle(b"handle-lo-auditor", &ciphertext_lo.handle_auditor);
 
         transcript.append_commitment(b"comm-hi-amount", &ciphertext_hi.commitment);
-        transcript.append_handle(b"handle-hi-source", &ciphertext_hi.source);
-        transcript.append_handle(b"handle-hi-dest", &ciphertext_hi.dest);
-        transcript.append_handle(b"handle-hi-auditor", &ciphertext_hi.auditor);
+        transcript.append_handle(b"handle-hi-source", &ciphertext_hi.handle_source);
+        transcript.append_handle(b"handle-hi-dest", &ciphertext_hi.handle_dest);
+        transcript.append_handle(b"handle-hi-auditor", &ciphertext_hi.handle_auditor);
 
         transcript.append_ciphertext(b"ciphertext-new-source", ciphertext_new_source);
 
@@ -302,7 +311,7 @@ impl TransferProof {
         transcript.append_commitment(b"commitment-new-source", &pod_commitment_new_source);
 
         // generate equality_proof
-        let equality_proof = EqualityProof::new(
+        let equality_proof = CtxtCommEqualityProof::new(
             keypair_source,
             ciphertext_new_source,
             source_new_balance,
@@ -325,7 +334,11 @@ impl TransferProof {
                 transfer_amount_lo as u64,
                 transfer_amount_hi as u64,
             ],
-            vec![64, 32, 32],
+            vec![
+                TRANSFER_SOURCE_AMOUNT_BIT_LENGTH,
+                TRANSFER_AMOUNT_LO_BIT_LENGTH,
+                TRANSFER_AMOUNT_HI_BIT_LENGTH,
+            ],
             vec![&opening_source, opening_lo, opening_hi],
             transcript,
         );
@@ -343,39 +356,38 @@ impl TransferProof {
         ciphertext_lo: &TransferAmountEncryption,
         ciphertext_hi: &TransferAmountEncryption,
         transfer_pubkeys: &TransferPubkeys,
-        new_spendable_ciphertext: &ElGamalCiphertext,
+        ciphertext_new_spendable: &ElGamalCiphertext,
         transcript: &mut Transcript,
     ) -> Result<(), ProofError> {
         transcript.append_commitment(b"commitment-new-source", &self.commitment_new_source);
 
         let commitment: PedersenCommitment = self.commitment_new_source.try_into()?;
-        let equality_proof: EqualityProof = self.equality_proof.try_into()?;
+        let equality_proof: CtxtCommEqualityProof = self.equality_proof.try_into()?;
         let aggregated_validity_proof: AggregatedValidityProof = self.validity_proof.try_into()?;
         let range_proof: RangeProof = self.range_proof.try_into()?;
 
         // verify equality proof
         //
         // TODO: we can also consider verifying equality and range proof in a batch
         equality_proof.verify(
-            &transfer_pubkeys.source,
-            new_spendable_ciphertext,
+            &transfer_pubkeys.pubkey_source,
+            ciphertext_new_spendable,
             &commitment,
             transcript,
         )?;
 
-        println!("equality pass");
-
         // verify validity proof
         aggregated_validity_proof.verify(
-            (&transfer_pubkeys.dest, &transfer_pubkeys.auditor),
+            (
+                &transfer_pubkeys.pubkey_dest,
+                &transfer_pubkeys.pubkey_auditor,
+            ),
             (&ciphertext_lo.commitment, &ciphertext_hi.commitment),
-            (&ciphertext_lo.dest, &ciphertext_hi.dest),
-            (&ciphertext_lo.auditor, &ciphertext_hi.auditor),
+            (&ciphertext_lo.handle_dest, &ciphertext_hi.handle_dest),
+            (&ciphertext_lo.handle_auditor, &ciphertext_hi.handle_auditor),
             transcript,
         )?;
 
-        println!("validity pass");
-
         // verify range proof
         let commitment_new_source = self.commitment_new_source.try_into()?;
         range_proof.verify(
@@ -397,34 +409,36 @@ impl TransferProof {
 #[repr(C)]
 #[cfg(not(target_arch = "bpf"))]
 pub struct TransferPubkeys {
-    pub source: ElGamalPubkey,
-    pub dest: ElGamalPubkey,
-    pub auditor: ElGamalPubkey,
+    pub pubkey_source: ElGamalPubkey,
+    pub pubkey_dest: ElGamalPubkey,
+    pub pubkey_auditor: ElGamalPubkey,
 }
 
 #[cfg(not(target_arch = "bpf"))]
 impl TransferPubkeys {
     // TODO: use constructor instead
     pub fn to_bytes(&self) -> [u8; 96] {
         let mut bytes = [0u8; 96];
-        bytes[..32].copy_from_slice(&self.source.to_bytes());
-        bytes[32..64].copy_from_slice(&self.dest.to_bytes());
-        bytes[64..96].copy_from_slice(&self.auditor.to_bytes());
+        bytes[..32].copy_from_slice(&self.pubkey_source.to_bytes());
+        bytes[32..64].copy_from_slice(&self.pubkey_dest.to_bytes());
+        bytes[64..96].copy_from_slice(&self.pubkey_auditor.to_bytes());
         bytes
     }
 
     pub fn from_bytes(bytes: &[u8]) -> Result<Self, ProofError> {
         let bytes = array_ref![bytes, 0, 96];
-        let (source, dest, auditor) = array_refs![bytes, 32, 32, 32];
+        let (pubkey_source, pubkey_dest, pubkey_auditor) = array_refs![bytes, 32, 32, 32];
 
-        let source = ElGamalPubkey::from_bytes(source).ok_or(ProofError::Verification)?;
-        let dest = ElGamalPubkey::from_bytes(dest).ok_or(ProofError::Verification)?;
-        let auditor = ElGamalPubkey::from_bytes(auditor).ok_or(ProofError::Verification)?;
+        let pubkey_source =
+            ElGamalPubkey::from_bytes(pubkey_source).ok_or(ProofError::Verification)?;
+        let pubkey_dest = ElGamalPubkey::from_bytes(pubkey_dest).ok_or(ProofError::Verification)?;
+        let pubkey_auditor =
+            ElGamalPubkey::from_bytes(pubkey_auditor).ok_or(ProofError::Verification)?;
 
         Ok(Self {
-            source,
-            dest,
-            auditor,
+            pubkey_source,
+            pubkey_dest,
+            pubkey_auditor,
         })
     }
 }