updating close account zk proof

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

@@ -5,7 +5,10 @@ use {
 #[cfg(not(target_arch = "bpf"))]
 use {
     crate::{
-        encryption::elgamal::{ElGamalCiphertext, ElGamalSecretKey},
+        encryption::{
+            elgamal::{ElGamalCiphertext, ElGamalKeypair, ElGamalPubkey},
+            pedersen::PedersenBase,
+        },
         errors::ProofError,
         instruction::Verifiable,
         transcript::TranscriptProtocol,
@@ -30,6 +33,9 @@ use {
 #[derive(Clone, Copy, Pod, Zeroable)]
 #[repr(C)]
 pub struct CloseAccountData {
+    /// The source account ElGamal pubkey
+    pub elgamal_pubkey: pod::ElGamalPubkey, // 32 bytes
+
     /// The source account available balance in encrypted form
     pub balance: pod::ElGamalCiphertext, // 64 bytes
 
@@ -39,10 +45,11 @@ pub struct CloseAccountData {
 
 #[cfg(not(target_arch = "bpf"))]
 impl CloseAccountData {
-    pub fn new(source_sk: &ElGamalSecretKey, balance: ElGamalCiphertext) -> Self {
+    pub fn new(source_keypair: &ElGamalKeypair, balance: ElGamalCiphertext) -> Self {
-        let proof = CloseAccountProof::new(source_sk, &balance);
+        let proof = CloseAccountProof::new(source_keypair, &balance);
 
         CloseAccountData {
+            elgamal_pubkey: source_keypair.public.into(),
             balance: balance.into(),
             proof,
         }
@@ -52,8 +59,9 @@ impl CloseAccountData {
 #[cfg(not(target_arch = "bpf"))]
 impl Verifiable for CloseAccountData {
     fn verify(&self) -> Result<(), ProofError> {
+        let elgamal_pubkey = self.elgamal_pubkey.try_into()?;
         let balance = self.balance.try_into()?;
-        self.proof.verify(&balance)
+        self.proof.verify(&elgamal_pubkey, &balance)
     }
 }
 
@@ -63,8 +71,9 @@ impl Verifiable for CloseAccountData {
 #[repr(C)]
 #[allow(non_snake_case)]
 pub struct CloseAccountProof {
-    pub R: pod::CompressedRistretto, // 32 bytes
+    pub Y_P: pod::CompressedRistretto, // 32 bytes
-    pub z: pod::Scalar,              // 32 bytes
+    pub Y_D: pod::CompressedRistretto, // 32 bytes
+    pub z: pod::Scalar,                // 32 bytes
 }
 
 #[allow(non_snake_case)]
@@ -74,67 +83,86 @@ impl CloseAccountProof {
         Transcript::new(b"CloseAccountProof")
     }
 
-    pub fn new(source_sk: &ElGamalSecretKey, balance: &ElGamalCiphertext) -> Self {
+    pub fn new(source_keypair: &ElGamalKeypair, balance: &ElGamalCiphertext) -> Self {
         let mut transcript = Self::transcript_new();
 
         // add a domain separator to record the start of the protocol
         transcript.close_account_proof_domain_sep();
 
         // extract the relevant scalar and Ristretto points from the input
-        let s = source_sk.get_scalar();
+        let P = source_keypair.public.get_point();
-        let C = balance.decrypt_handle.get_point();
+        let s = source_keypair.secret.get_scalar();
+
+        let C = balance.message_comm.get_point();
+        let D = balance.decrypt_handle.get_point();
+
+        // record ElGamal pubkey and ciphertext in the transcript
+        transcript.append_point(b"P", &P.compress());
+        transcript.append_point(b"C", &C.compress());
+        transcript.append_point(b"D", &D.compress());
 
         // generate a random masking factor that also serves as a nonce
-        let r = Scalar::random(&mut OsRng); // using OsRng for now
+        let y = Scalar::random(&mut OsRng);
-        let R = (r * C).compress();
+        let Y_P = (y * P).compress();
+        let Y_D = (y * D).compress();
 
-        // record R on transcript and receive a challenge scalar
+        // record Y in transcript and receive a challenge scalar
-        transcript.append_point(b"R", &R);
+        transcript.append_point(b"Y_P", &Y_P);
+        transcript.append_point(b"Y_D", &Y_D);
         let c = transcript.challenge_scalar(b"c");
 
         // compute the masked secret key
-        let z = c * s + r;
+        let z = c * s + y;
 
         CloseAccountProof {
-            R: R.into(),
+            Y_P: Y_P.into(),
+            Y_D: Y_D.into(),
             z: z.into(),
         }
     }
 
-    pub fn verify(&self, balance: &ElGamalCiphertext) -> Result<(), ProofError> {
+    pub fn verify(
+        &self,
+        elgamal_pubkey: &ElGamalPubkey,
+        balance: &ElGamalCiphertext,
+    ) -> Result<(), ProofError> {
         let mut transcript = Self::transcript_new();
 
         // add a domain separator to record the start of the protocol
         transcript.close_account_proof_domain_sep();
 
         // extract the relevant scalar and Ristretto points from the input
+        let P = elgamal_pubkey.get_point();
         let C = balance.message_comm.get_point();
         let D = balance.decrypt_handle.get_point();
 
-        let R = self.R.into();
+        let H = PedersenBase::default().H;
+
+        let Y_P = self.Y_P.into();
+        let Y_D = self.Y_D.into();
         let z = self.z.into();
 
-        // Edge case #1: if both C and D are identities, then this is a valid encryption of zero
+        // record ElGamal pubkey and ciphertext in the transcript
-        if C.is_identity() && D.is_identity() {
+        transcript.validate_and_append_point(b"P", &P.compress())?;
-            transcript.append_point(b"R", &R);
+        transcript.append_point(b"C", &C.compress());
-            return Ok(());
+        transcript.append_point(b"D", &D.compress());
-        }
 
-        // Edge case #2: if D is zeroed, but C is not, then this is an invalid ciphertext
+        // record Y in transcript and receive challenge scalars
-        if D.is_identity() {
+        transcript.validate_and_append_point(b"Y_P", &Y_P)?;
-            transcript.append_point(b"R", &R);
+        transcript.append_point(b"Y_D", &Y_D);
-            return Err(ProofError::VerificationError);
-        }
 
-        // generate a challenge scalar
-        transcript.validate_and_append_point(b"R", &R)?;
         let c = transcript.challenge_scalar(b"c");
+        let w = transcript.challenge_scalar(b"w"); // w used for multiscalar multiplication verification
 
         // decompress R or return verification error
-        let R = R.decompress().ok_or(ProofError::VerificationError)?;
+        let Y_P = Y_P.decompress().ok_or(ProofError::VerificationError)?;
+        let Y_D = Y_D.decompress().ok_or(ProofError::VerificationError)?;
 
         // check the required algebraic relation
-        let check = RistrettoPoint::multiscalar_mul(vec![z, -c, -Scalar::one()], vec![D, C, R]);
+        let check = RistrettoPoint::multiscalar_mul(
+            vec![z, -c, -Scalar::one(), w * z, -w * c, -w],
+            vec![P, H, Y_P, D, C, Y_D],
+        );
 
         if check.is_identity() {
             Ok(())
@@ -156,11 +184,24 @@ mod test {
 
     #[test]
     fn test_close_account_correctness() {
-        let source = ElGamalKeypair::default();
+        let source_keypair = ElGamalKeypair::default();
 
-        // invalid ciphertexts
+        // general case: encryption of 0
-        let balance = source.public.encrypt(0_u64);
+        let balance = source_keypair.public.encrypt(0_u64);
+        let proof = CloseAccountProof::new(&source_keypair, &balance);
+        assert!(proof.verify(&source_keypair.public, &balance).is_ok());
 
+        // general case: encryption of > 0
+        let balance = source_keypair.public.encrypt(1_u64);
+        let proof = CloseAccountProof::new(&source_keypair, &balance);
+        assert!(proof.verify(&source_keypair.public, &balance).is_err());
+
+        // // edge case: all zero ciphertext - such ciphertext should always be a valid encryption of 0
+        let zeroed_ct: ElGamalCiphertext = pod::ElGamalCiphertext::zeroed().try_into().unwrap();
+        let proof = CloseAccountProof::new(&source_keypair, &zeroed_ct);
+        assert!(proof.verify(&source_keypair.public, &zeroed_ct).is_ok());
+
+        // edge cases: only C or D is zero - such ciphertext is always invalid
         let zeroed_comm = Pedersen::with(0_u64, &PedersenOpening::default());
         let handle = balance.decrypt_handle;
 
@@ -169,32 +210,19 @@ mod test {
             decrypt_handle: handle,
         };
 
-        let proof = CloseAccountProof::new(&source.secret, &zeroed_comm_ciphertext);
+        let proof = CloseAccountProof::new(&source_keypair, &zeroed_comm_ciphertext);
-        assert!(proof.verify(&zeroed_comm_ciphertext).is_err());
+        assert!(proof
+            .verify(&source_keypair.public, &zeroed_comm_ciphertext)
+            .is_err());
 
         let zeroed_handle_ciphertext = ElGamalCiphertext {
             message_comm: balance.message_comm,
             decrypt_handle: PedersenDecryptHandle::default(),
         };
 
-        let proof = CloseAccountProof::new(&source.secret, &zeroed_handle_ciphertext);
+        let proof = CloseAccountProof::new(&source_keypair, &zeroed_handle_ciphertext);
-        assert!(proof.verify(&zeroed_handle_ciphertext).is_err());
+        assert!(proof
-
+            .verify(&source_keypair.public, &zeroed_handle_ciphertext)
-        // valid ciphertext, but encryption of non-zero amount
+            .is_err());
-        let balance = source.public.encrypt(55_u64);
-
-        let proof = CloseAccountProof::new(&source.secret, &balance);
-        assert!(proof.verify(&balance).is_err());
-
-        // all-zeroed ciphertext interpretted as a valid encryption of zero
-        let zeroed_ct: ElGamalCiphertext = pod::ElGamalCiphertext::zeroed().try_into().unwrap();
-        let proof = CloseAccountProof::new(&source.secret, &zeroed_ct);
-        assert!(proof.verify(&zeroed_ct).is_ok());
-
-        // general case: valid encryption of zero
-        let balance = source.public.encrypt(0_u64);
-
-        let proof = CloseAccountProof::new(&source.secret, &balance);
-        assert!(proof.verify(&balance).is_ok());
     }
 }