refactor

Author: Mr-Leshiy

rust/catalyst-voting/src/crypto/zk_unit_vector/mod.rs

@@ -43,7 +43,7 @@ pub struct UnitVectorProof(
 /// the proof will be invalid.
 pub fn generate_unit_vector_proof<R: CryptoRngCore>(
     unit_vector: &[Scalar], mut encryption_randomness: Vec<Scalar>,
-    mut ciphertexts: Vec<Ciphertext>, public_key: &PublicKey, commitment_key: &PublicKey,
+    mut ciphertexts: Vec<Ciphertext>, public_key: &PublicKey, commitment_key: &GroupElement,
     rng: &mut R,
 ) -> UnitVectorProof {
     let i = unit_vector
@@ -249,7 +249,7 @@ mod tests {
     #[proptest]
     fn zk_unit_vector_test(
         secret_key: SecretKey, secret_commitment_key: SecretKey,
-        #[strategy(2..10_usize)] unit_vector_size: usize,
+        #[strategy(1..10_usize)] unit_vector_size: usize,
         #[strategy(0..#unit_vector_size)] unit_vector_index: usize,
     ) {
         let mut rng = OsRng;
@@ -300,7 +300,7 @@ mod tests {
     #[proptest]
     fn not_a_unit_vector_test(
         secret_key: SecretKey, secret_commitment_key: SecretKey,
-        #[any(size_range(2..10_usize).lift())] random_vector: Vec<Scalar>,
+        #[any(size_range(1..10_usize).lift())] random_vector: Vec<Scalar>,
     ) {
         let mut rng = OsRng;
 

rust/catalyst-voting/src/voter/mod.rs

@@ -26,7 +26,7 @@ pub struct EncryptionRandomness(Vec<Scalar>);
 
 impl EncryptionRandomness {
     /// Randomly generate the `EncryptionRandomness`.
-    pub fn generate<R: CryptoRngCore>(rng: &mut R, voting_options: usize) -> Self {
+    fn generate<R: CryptoRngCore>(rng: &mut R, voting_options: usize) -> Self {
         Self((0..voting_options).map(|_| Scalar::random(rng)).collect())
     }
 }
@@ -79,30 +79,15 @@ impl Vote {
     }
 }
 
-/// Encrypted vote error
-#[derive(thiserror::Error, Debug)]
-pub enum EncryptedVoteError {
-    /// Incorrect randomness length
-    #[error(
-        "Invalid randomness length, the length of randomness: {0}, should be equal to the number of voting options: {1}."
-    )]
-    IncorrectRandomnessLength(usize, usize),
-}
-
 /// Create a new encrypted vote from the given vote and public key.
 /// More detailed described [here](https://input-output-hk.github.io/catalyst-voices/architecture/08_concepts/voting_transaction/crypto/#vote-encryption)
 ///
 /// # Errors
 ///   - `EncryptedVoteError`
-pub fn encrypt_vote(
-    vote: &Vote, public_key: &PublicKey, randomness: &EncryptionRandomness,
-) -> Result<EncryptedVote, EncryptedVoteError> {
-    if vote.voting_options != randomness.0.len() {
-        return Err(EncryptedVoteError::IncorrectRandomnessLength(
-            randomness.0.len(),
-            vote.voting_options,
-        ));
-    }
+pub fn encrypt_vote<R: CryptoRngCore>(
+    vote: &Vote, public_key: &PublicKey, rng: &mut R,
+) -> (EncryptedVote, EncryptionRandomness) {
+    let randomness = EncryptionRandomness::generate(rng, vote.voting_options);
 
     let unit_vector = vote.to_unit_vector();
     let ciphers = unit_vector
@@ -111,16 +96,12 @@ pub fn encrypt_vote(
         .map(|(m, r)| encrypt(m, public_key, r))
         .collect();
 
-    Ok(EncryptedVote(ciphers))
+    (EncryptedVote(ciphers), randomness)
 }
 
 #[cfg(test)]
 mod tests {
-    use proptest::sample::size_range;
-    use test_strategy::proptest;
-
     use super::*;
-    use crate::crypto::elgamal::SecretKey;
 
     #[test]
     fn vote_test() {
@@ -150,17 +131,4 @@ mod tests {
         assert!(Vote::new(3, voting_options).is_err());
         assert!(Vote::new(4, voting_options).is_err());
     }
-
-    #[proptest]
-    fn encrypt_test(
-        secret_key: SecretKey, #[strategy(1..10usize)] voting_options: usize,
-        #[any(size_range(#voting_options).lift())] randomness: Vec<Scalar>,
-    ) {
-        let public_key = secret_key.public_key();
-        let vote = Vote::new(0, voting_options).unwrap();
-
-        let encrypted =
-            encrypt_vote(&vote, &public_key, &EncryptionRandomness(randomness)).unwrap();
-        assert_eq!(encrypted.0.len(), vote.voting_options);
-    }
 }

rust/catalyst-voting/tests/voting_test.rs

@@ -2,7 +2,7 @@
 
 use catalyst_voting::{
     decrypt_tally, encrypt_vote, generate_tally_proof, tally, verify_tally_proof,
-    DecryptionTallySetup, EncryptionRandomness, SecretKey, Vote,
+    DecryptionTallySetup, SecretKey, Vote,
 };
 use proptest::prelude::ProptestConfig;
 use test_strategy::{proptest, Arbitrary};
@@ -29,15 +29,10 @@ fn voting_test(voters: [Voter; 100]) {
         .map(|voter| Vote::new(voter.choice, VOTING_OPTIONS).unwrap())
         .collect();
 
-    let voters_randomness: Vec<_> = (0..voters.len())
-        .map(|_| EncryptionRandomness::generate(&mut rng, VOTING_OPTIONS))
-        .collect();
-
-    let encrypted_votes: Vec<_> = votes
+    let (encrypted_votes, _randomness): (Vec<_>, Vec<_>) = votes
         .iter()
-        .zip(voters_randomness.iter())
-        .map(|(vote, r)| encrypt_vote(vote, &election_public_key, r).unwrap())
-        .collect();
+        .map(|vote| encrypt_vote(vote, &election_public_key, &mut rng))
+        .unzip();
 
     let voting_powers: Vec<_> = voters
         .iter()