From 4c033ef0c9f1d707e51f49a6ecb652499f325247 Mon Sep 17 00:00:00 2001
From: Fiono11 <ruipedromorais11@gmail.com>
Date: Mon, 29 Apr 2024 17:52:59 +0100
Subject: [PATCH] Add encryption error

---
 src/olaf/errors.rs      |  2 ++
 src/olaf/simplpedpop.rs | 27 +++++++++++++++------------
 src/olaf/tests.rs       |  4 ++--
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/src/olaf/errors.rs b/src/olaf/errors.rs
index 008ce51..a62ba82 100644
--- a/src/olaf/errors.rs
+++ b/src/olaf/errors.rs
@@ -66,6 +66,8 @@ pub enum DKGError {
     },
     /// Decryption error when decrypting an encrypted secret share.
     DecryptionError(chacha20poly1305::Error),
+    /// Encryption error when encrypting the secret share.
+    EncryptionError(chacha20poly1305::Error),
     /// Incorrect number of coefficient commitments.
     InvalidSecretPolynomialCommitment {
         /// The expected value.
diff --git a/src/olaf/simplpedpop.rs b/src/olaf/simplpedpop.rs
index 6a2c6f9..2bf105a 100644
--- a/src/olaf/simplpedpop.rs
+++ b/src/olaf/simplpedpop.rs
@@ -139,7 +139,7 @@ impl SecretShare {
         decryption_key: &Scalar,
         encryption_key: &RistrettoPoint,
         context: &[u8],
-    ) -> EncryptedSecretShare {
+    ) -> DKGResult<EncryptedSecretShare> {
         let shared_secret = decryption_key * encryption_key;
 
         let mut transcript = Transcript::new(b"encryption");
@@ -152,9 +152,11 @@ impl SecretShare {
         let cipher: ChaCha20Poly1305 = make_aead::<Transcript, ChaCha20Poly1305>(transcript);
         let nonce = Nonce::from_slice(&bytes[..]);
 
-        let ciphertext: Vec<u8> = cipher.encrypt(nonce, &self.0.as_bytes()[..]).unwrap();
+        let ciphertext: Vec<u8> = cipher
+            .encrypt(nonce, &self.0.as_bytes()[..])
+            .map_err(DKGError::EncryptionError)?;
 
-        EncryptedSecretShare(ciphertext)
+        Ok(EncryptedSecretShare(ciphertext))
     }
 }
 
@@ -397,12 +399,12 @@ pub mod round2 {
             deckey: Scalar,
             enckey: RistrettoPoint,
             context: &[u8],
-        ) -> PrivateMessage {
-            let encrypted_secret_share = secret_share.encrypt(&deckey, &enckey, context);
+        ) -> DKGResult<PrivateMessage> {
+            let encrypted_secret_share = secret_share.encrypt(&deckey, &enckey, context)?;
 
-            PrivateMessage {
+            Ok(PrivateMessage {
                 encrypted_secret_share,
-            }
+            })
         }
     }
 
@@ -477,7 +479,7 @@ pub mod round2 {
             &round1_private_data.secret_polynomial,
             round1_private_data.secret_key,
             secret_commitment,
-        );
+        )?;
 
         Ok((public_data, messages))
     }
@@ -664,8 +666,9 @@ pub mod round2 {
         secret_polynomial: &SecretPolynomial,
         secret_key: SecretKey,
         secret_commitment: &SecretCommitment,
-    ) -> Messages {
+    ) -> DKGResult<Messages> {
         let mut private_messages = BTreeMap::new();
+
         let enc_keys: Vec<RistrettoPoint> = round2_public_data
             .round1_public_messages
             .values()
@@ -691,7 +694,7 @@ pub mod round2 {
                     secret_key.key,
                     enc_keys[i],
                     identifier.0.as_bytes(),
-                ),
+                )?,
             );
         }
 
@@ -704,10 +707,10 @@ pub mod round2 {
 
         let public_message = PublicMessage { certificate };
 
-        Messages {
+        Ok(Messages {
             private_messages,
             public_message,
-        }
+        })
     }
 }
 
diff --git a/src/olaf/tests.rs b/src/olaf/tests.rs
index 64e7ed3..494b59d 100644
--- a/src/olaf/tests.rs
+++ b/src/olaf/tests.rs
@@ -396,7 +396,7 @@ mod tests {
                 .get_mut(&participants_sets_of_participants[0].own_identifier)
                 .unwrap();
 
-            private_message.encrypted_secret_share = enc_share;
+            private_message.encrypted_secret_share = enc_share.unwrap();
 
             let result = round3(
                 &participants_sets_of_participants,
@@ -956,7 +956,7 @@ mod tests {
             let original_share = SecretShare(Scalar::random(&mut rng));
 
             let encrypted_share = original_share.encrypt(&deckey, &enckey, context);
-            let decrypted_share = encrypted_share.decrypt(&deckey, &enckey, context);
+            let decrypted_share = encrypted_share.unwrap().decrypt(&deckey, &enckey, context);
 
             assert_eq!(
                 original_share.0,