diff --git a/zk-token-sdk/src/encryption/elgamal.rs b/zk-token-sdk/src/encryption/elgamal.rs index 7e83a69b99..7c5ca32312 100644 --- a/zk-token-sdk/src/encryption/elgamal.rs +++ b/zk-token-sdk/src/encryption/elgamal.rs @@ -21,20 +21,24 @@ use { }; /// Handle for the (twisted) ElGamal encryption scheme -pub struct ElGamal; +pub struct ElGamal { + pub pk: ElGamalPubkey, + pub sk: ElGamalSecretKey, +} + impl ElGamal { /// Generates the public and secret keys for ElGamal encryption. #[cfg(not(target_arch = "bpf"))] #[allow(clippy::new_ret_no_self)] - pub fn new() -> (ElGamalPubkey, ElGamalSecretKey) { - ElGamal::with(&mut OsRng) // using OsRng for now + pub fn new() -> Self { + Self::with(&mut OsRng) // using OsRng for now } /// On input a randomness generator, the function generates the public and /// secret keys for ElGamal encryption. #[cfg(not(target_arch = "bpf"))] #[allow(non_snake_case)] - pub fn with(rng: &mut T) -> (ElGamalPubkey, ElGamalSecretKey) { + pub fn with(rng: &mut T) -> Self { // sample a non-zero scalar let mut s: Scalar; loop { @@ -48,7 +52,10 @@ impl ElGamal { let H = PedersenBase::default().H; let P = s.invert() * H; - (ElGamalPubkey(P), ElGamalSecretKey(s)) + Self { + pk: ElGamalPubkey(P), + sk: ElGamalSecretKey(s), + } } /// On input a public key and a message to be encrypted, the function @@ -101,7 +108,7 @@ impl ElGamal { /// On input a secret key and a ciphertext, the function decrypts the /// ciphertext for a u32 value. pub fn decrypt_u32(sk: &ElGamalSecretKey, ct: &ElGamalCiphertext) -> Option { - let discrete_log_instance = ElGamal::decrypt(sk, ct); + let discrete_log_instance = Self::decrypt(sk, ct); discrete_log_instance.decode_u32() } @@ -112,7 +119,7 @@ impl ElGamal { ct: &ElGamalCiphertext, hashmap: &HashMap<[u8; 32], u32>, ) -> Option { - let discrete_log_instance = ElGamal::decrypt(sk, ct); + let discrete_log_instance = Self::decrypt(sk, ct); discrete_log_instance.decode_u32_online(hashmap) } } @@ -372,7 +379,7 @@ mod tests { #[test] fn test_encrypt_decrypt_correctness() { - let (pk, sk) = ElGamal::new(); + let ElGamal { pk, sk } = ElGamal::new(); let msg: u32 = 57; let ct = ElGamal::encrypt(&pk, msg); @@ -389,8 +396,8 @@ mod tests { #[test] fn test_decrypt_handle() { - let (pk_1, sk_1) = ElGamal::new(); - let (pk_2, sk_2) = ElGamal::new(); + let ElGamal { pk: pk_1, sk: sk_1 } = ElGamal::new(); + let ElGamal { pk: pk_2, sk: sk_2 } = ElGamal::new(); let msg: u32 = 77; let (comm, open) = Pedersen::new(msg); @@ -416,7 +423,7 @@ mod tests { #[test] fn test_homomorphic_addition() { - let (pk, _) = ElGamal::new(); + let ElGamal { pk, sk: _ } = ElGamal::new(); let msg_0: u64 = 57; let msg_1: u64 = 77; @@ -441,7 +448,7 @@ mod tests { #[test] fn test_homomorphic_subtraction() { - let (pk, _) = ElGamal::new(); + let ElGamal { pk, sk: _ } = ElGamal::new(); let msg_0: u64 = 77; let msg_1: u64 = 55; @@ -466,7 +473,7 @@ mod tests { #[test] fn test_homomorphic_multiplication() { - let (pk, _) = ElGamal::new(); + let ElGamal { pk, sk: _ } = ElGamal::new(); let msg_0: u64 = 57; let msg_1: u64 = 77; @@ -482,7 +489,7 @@ mod tests { #[test] fn test_homomorphic_division() { - let (pk, _) = ElGamal::new(); + let ElGamal { pk, sk: _ } = ElGamal::new(); let msg_0: u64 = 55; let msg_1: u64 = 5; @@ -498,7 +505,7 @@ mod tests { #[test] fn test_serde_ciphertext() { - let (pk, _) = ElGamal::new(); + let ElGamal { pk, sk: _ } = ElGamal::new(); let msg: u64 = 77; let ct = pk.encrypt(msg); @@ -510,7 +517,7 @@ mod tests { #[test] fn test_serde_pubkey() { - let (pk, _) = ElGamal::new(); + let ElGamal { pk, sk: _ } = ElGamal::new(); let encoded = bincode::serialize(&pk).unwrap(); let decoded: ElGamalPubkey = bincode::deserialize(&encoded).unwrap(); @@ -520,7 +527,7 @@ mod tests { #[test] fn test_serde_secretkey() { - let (_, sk) = ElGamal::new(); + let ElGamal { pk: _, sk } = ElGamal::new(); let encoded = bincode::serialize(&sk).unwrap(); let decoded: ElGamalSecretKey = bincode::deserialize(&encoded).unwrap(); diff --git a/zk-token-sdk/src/instruction/close_account.rs b/zk-token-sdk/src/instruction/close_account.rs index cbd705c358..e33d44bac8 100644 --- a/zk-token-sdk/src/instruction/close_account.rs +++ b/zk-token-sdk/src/instruction/close_account.rs @@ -152,10 +152,10 @@ mod test { #[test] fn test_close_account_correctness() { - let (source_pk, source_sk) = ElGamal::new(); + let source = ElGamal::new(); // invalid ciphertexts - let balance = source_pk.encrypt(0_u64); + let balance = source.pk.encrypt(0_u64); let zeroed_comm = Pedersen::with(0_u64, &PedersenOpening::default()); let handle = balance.decrypt_handle; @@ -165,7 +165,7 @@ mod test { decrypt_handle: handle, }; - let proof = CloseAccountProof::new(&source_sk, &zeroed_comm_ciphertext); + let proof = CloseAccountProof::new(&source.sk, &zeroed_comm_ciphertext); assert!(proof.verify(&zeroed_comm_ciphertext).is_err()); let zeroed_handle_ciphertext = ElGamalCiphertext { @@ -173,24 +173,24 @@ mod test { decrypt_handle: PedersenDecryptHandle::default(), }; - let proof = CloseAccountProof::new(&source_sk, &zeroed_handle_ciphertext); + let proof = CloseAccountProof::new(&source.sk, &zeroed_handle_ciphertext); assert!(proof.verify(&zeroed_handle_ciphertext).is_err()); // valid ciphertext, but encryption of non-zero amount - let balance = source_pk.encrypt(55_u64); + let balance = source.pk.encrypt(55_u64); - let proof = CloseAccountProof::new(&source_sk, &balance); + let proof = CloseAccountProof::new(&source.sk, &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_sk, &zeroed_ct); + let proof = CloseAccountProof::new(&source.sk, &zeroed_ct); assert!(proof.verify(&zeroed_ct).is_ok()); // general case: valid encryption of zero - let balance = source_pk.encrypt(0_u64); + let balance = source.pk.encrypt(0_u64); - let proof = CloseAccountProof::new(&source_sk, &balance); + let proof = CloseAccountProof::new(&source.sk, &balance); assert!(proof.verify(&balance).is_ok()); } } diff --git a/zk-token-sdk/src/instruction/transfer.rs b/zk-token-sdk/src/instruction/transfer.rs index 69b9fb0fb5..88ef2bea6c 100644 --- a/zk-token-sdk/src/instruction/transfer.rs +++ b/zk-token-sdk/src/instruction/transfer.rs @@ -468,9 +468,12 @@ mod test { #[test] fn test_transfer_correctness() { // ElGamal keys for source, destination, and auditor accounts - let (source_pk, source_sk) = ElGamal::new(); - let (dest_pk, _) = ElGamal::new(); - let (auditor_pk, _) = ElGamal::new(); + let ElGamal { + pk: source_pk, + sk: source_sk, + } = ElGamal::new(); + let dest_pk = ElGamal::new().pk; + let auditor_pk = ElGamal::new().pk; // create source account spendable ciphertext let spendable_balance: u64 = 77; @@ -496,9 +499,15 @@ mod test { #[test] fn test_source_dest_ciphertext() { // ElGamal keys for source, destination, and auditor accounts - let (source_pk, source_sk) = ElGamal::new(); - let (dest_pk, dest_sk) = ElGamal::new(); - let (auditor_pk, _) = ElGamal::new(); + let ElGamal { + pk: source_pk, + sk: source_sk, + } = ElGamal::new(); + let ElGamal { + pk: dest_pk, + sk: dest_sk, + } = ElGamal::new(); + let auditor_pk = ElGamal::new().pk; // create source account spendable ciphertext let spendable_balance: u64 = 77; diff --git a/zk-token-sdk/src/instruction/update_account_pk.rs b/zk-token-sdk/src/instruction/update_account_pk.rs index adb818bc38..95f3a707d2 100644 --- a/zk-token-sdk/src/instruction/update_account_pk.rs +++ b/zk-token-sdk/src/instruction/update_account_pk.rs @@ -219,50 +219,50 @@ mod test { #[test] fn test_update_account_public_key_general_cases() { - let (current_pk, current_sk) = ElGamal::new(); - let (new_pk, new_sk) = ElGamal::new(); + let current = ElGamal::new(); + let new = ElGamal::new(); // If current_ct and new_ct encrypt same values, then the proof verification should succeed let balance: u64 = 77; - let current_ct = current_pk.encrypt(balance); - let new_ct = new_pk.encrypt(balance); + let current_ct = current.pk.encrypt(balance); + let new_ct = new.pk.encrypt(balance); - let proof = UpdateAccountPkProof::new(balance, ¤t_sk, &new_sk, ¤t_ct, &new_ct); + let proof = UpdateAccountPkProof::new(balance, ¤t.sk, &new.sk, ¤t_ct, &new_ct); assert!(proof.verify(¤t_ct, &new_ct).is_ok()); // If current_ct and new_ct encrypt different values, then the proof verification should fail - let new_ct = new_pk.encrypt(55_u64); + let new_ct = new.pk.encrypt(55_u64); - let proof = UpdateAccountPkProof::new(balance, ¤t_sk, &new_sk, ¤t_ct, &new_ct); + let proof = UpdateAccountPkProof::new(balance, ¤t.sk, &new.sk, ¤t_ct, &new_ct); assert!(proof.verify(¤t_ct, &new_ct).is_err()); } #[test] fn test_update_account_public_key_zeroed_ciphertexts() { - let (current_pk, current_sk) = ElGamal::new(); - let (new_pk, new_sk) = ElGamal::new(); + let current = ElGamal::new(); + let new = ElGamal::new(); // A zeroed cipehrtext should be considered as an account balance of 0 let balance: u64 = 0; let zeroed_ct_as_current_ct: ElGamalCiphertext = pod::ElGamalCiphertext::zeroed().try_into().unwrap(); - let new_ct: ElGamalCiphertext = new_pk.encrypt(balance); + let new_ct: ElGamalCiphertext = new.pk.encrypt(balance); let proof = UpdateAccountPkProof::new( balance, - ¤t_sk, - &new_sk, + ¤t.sk, + &new.sk, &zeroed_ct_as_current_ct, &new_ct, ); assert!(proof.verify(&zeroed_ct_as_current_ct, &new_ct).is_ok()); - let current_ct = current_pk.encrypt(balance); + let current_ct = current.pk.encrypt(balance); let zeroed_ct_as_new_ct: ElGamalCiphertext = pod::ElGamalCiphertext::zeroed().try_into().unwrap(); let proof = UpdateAccountPkProof::new( balance, - ¤t_sk, - &new_sk, + ¤t.sk, + &new.sk, ¤t_ct, &zeroed_ct_as_new_ct, ); @@ -274,8 +274,8 @@ mod test { pod::ElGamalCiphertext::zeroed().try_into().unwrap(); let proof = UpdateAccountPkProof::new( balance, - ¤t_sk, - &new_sk, + ¤t.sk, + &new.sk, &zeroed_ct_as_current_ct, &zeroed_ct_as_new_ct, ); @@ -286,11 +286,11 @@ mod test { #[test] fn test_update_account_public_key_partially_zeroed_ciphertexts() { - let (current_pk, current_sk) = ElGamal::new(); - let (new_pk, new_sk) = ElGamal::new(); + let current = ElGamal::new(); + let new = ElGamal::new(); let balance = 0_u64; - let balance_ciphertext = new_pk.encrypt(balance); + let balance_ciphertext = new.pk.encrypt(balance); let zeroed_comm = Pedersen::with(0_u64, &PedersenOpening::default()); let handle = balance_ciphertext.decrypt_handle; @@ -300,12 +300,12 @@ mod test { message_comm: zeroed_comm, decrypt_handle: handle, }; - let new_ct: ElGamalCiphertext = new_pk.encrypt(balance); + let new_ct: ElGamalCiphertext = new.pk.encrypt(balance); let proof = UpdateAccountPkProof::new( balance, - ¤t_sk, - &new_sk, + ¤t.sk, + &new.sk, &zeroed_comm_ciphertext, &new_ct, ); @@ -318,8 +318,8 @@ mod test { let proof = UpdateAccountPkProof::new( balance, - ¤t_sk, - &new_sk, + ¤t.sk, + &new.sk, &zeroed_handle_ciphertext, &new_ct, ); @@ -330,12 +330,12 @@ mod test { message_comm: zeroed_comm, decrypt_handle: handle, }; - let current_ct: ElGamalCiphertext = current_pk.encrypt(balance); + let current_ct: ElGamalCiphertext = current.pk.encrypt(balance); let proof = UpdateAccountPkProof::new( balance, - ¤t_sk, - &new_sk, + ¤t.sk, + &new.sk, ¤t_ct, &zeroed_comm_ciphertext, ); @@ -348,8 +348,8 @@ mod test { let proof = UpdateAccountPkProof::new( balance, - ¤t_sk, - &new_sk, + ¤t.sk, + &new.sk, ¤t_ct, &zeroed_handle_ciphertext, ); diff --git a/zk-token-sdk/src/instruction/withdraw.rs b/zk-token-sdk/src/instruction/withdraw.rs index 5d68a0993b..6929f153a8 100644 --- a/zk-token-sdk/src/instruction/withdraw.rs +++ b/zk-token-sdk/src/instruction/withdraw.rs @@ -179,17 +179,17 @@ mod test { #[ignore] fn test_withdraw_correctness() { // generate and verify proof for the proper setting - let (source_pk, source_sk) = ElGamal::new(); + let ElGamal { pk, sk } = ElGamal::new(); let current_balance: u64 = 77; - let current_balance_ct = source_pk.encrypt(current_balance); + let current_balance_ct = pk.encrypt(current_balance); let withdraw_amount: u64 = 55; let data = WithdrawData::new( withdraw_amount, - source_pk, - &source_sk, + pk, + &sk, current_balance, current_balance_ct, ); @@ -197,13 +197,7 @@ mod test { // generate and verify proof with wrong balance let wrong_balance: u64 = 99; - let data = WithdrawData::new( - withdraw_amount, - source_pk, - &source_sk, - wrong_balance, - current_balance_ct, - ); + let data = WithdrawData::new(withdraw_amount, pk, &sk, wrong_balance, current_balance_ct); assert!(data.verify().is_err()); // TODO: test for ciphertexts that encrypt numbers outside the 0, 2^64 range diff --git a/zk-token-sdk/src/zk_token_elgamal/ops.rs b/zk-token-sdk/src/zk_token_elgamal/ops.rs index 10498491cc..8ed92d6ca8 100644 --- a/zk-token-sdk/src/zk_token_elgamal/ops.rs +++ b/zk-token-sdk/src/zk_token_elgamal/ops.rs @@ -252,7 +252,7 @@ mod tests { // spendable_ct should be an encryption of 0 for any public key when // `PedersenOpen::default()` is used - let (pk, _) = ElGamal::new(); + let pk = ElGamal::new().pk; let balance: u64 = 0; assert_eq!( spendable_ct, @@ -276,7 +276,7 @@ mod tests { let added_ct = ops::add_to(&spendable_balance, 55).unwrap(); - let (pk, _) = ElGamal::new(); + let pk = ElGamal::new().pk; let expected: pod::ElGamalCiphertext = pk.encrypt_with(55_u64, &PedersenOpening::default()).into(); @@ -286,7 +286,7 @@ mod tests { #[test] fn test_subtract_from() { let amount = 77_u64; - let (pk, _) = ElGamal::new(); + let pk = ElGamal::new().pk; let open = PedersenOpening::random(&mut OsRng); let encrypted_amount: pod::ElGamalCiphertext = pk.encrypt_with(amount, &open).into(); @@ -312,9 +312,9 @@ mod tests { let (amount_lo, amount_hi) = split_u64_into_u32(transfer_amount); // generate public keys - let (source_pk, _) = ElGamal::new(); - let (dest_pk, _) = ElGamal::new(); - let (auditor_pk, _) = ElGamal::new(); + let source_pk = ElGamal::new().pk; + let dest_pk = ElGamal::new().pk; + let auditor_pk = ElGamal::new().pk; // commitments associated with TransferRangeProof let (comm_lo, open_lo) = Pedersen::new(amount_lo);