fix: prevent precomputed shares from being created with inapprioriate…

… variant
nucypher · Jan 19, 2024 · feee2f3 · feee2f3
1 parent e79b4e5
commit feee2f3
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 24 deletions.
diff --git a/ferveo-wasm/tests/node.rs b/ferveo-wasm/tests/node.rs
@@ -7,9 +7,6 @@ use itertools::zip_eq;
 use wasm_bindgen_test::*;
 
 type TestSetup = (
-    u32,
-    u32,
-    u32,
     Vec<Keypair>,
     Vec<Validator>,
     ValidatorArray,
@@ -19,11 +16,9 @@ type TestSetup = (
     Ciphertext,
 );
 
-fn setup_dkg() -> TestSetup {
-    let tau = 1;
-    let shares_num: u32 = 16;
-    let security_threshold = shares_num * 2 / 3;
+const TAU: u32 = 0;
 
+fn setup_dkg(shares_num: u32, security_threshold: u32) -> TestSetup {
     let validator_keypairs = (0..shares_num as usize)
         .map(gen_keypair)
         .collect::<Vec<Keypair>>();
@@ -38,7 +33,7 @@ fn setup_dkg() -> TestSetup {
     // validator, including themselves
     let messages = validators.iter().map(|sender| {
         let dkg = Dkg::new(
-            tau,
+            TAU,
             shares_num,
             security_threshold,
             &validators_js,
@@ -54,7 +49,7 @@ fn setup_dkg() -> TestSetup {
     // every validator can aggregate the transcripts
 
     let mut dkg = Dkg::new(
-        tau,
+        TAU,
         shares_num,
         security_threshold,
         &validators_js,
@@ -80,9 +75,6 @@ fn setup_dkg() -> TestSetup {
     let ciphertext = ferveo_encrypt(&msg, &aad, &dkg.public_key()).unwrap();
 
     (
-        tau,
-        shares_num,
-        security_threshold,
         validator_keypairs,
         validators,
         validators_js,
@@ -95,24 +87,23 @@ fn setup_dkg() -> TestSetup {
 
 #[wasm_bindgen_test]
 fn tdec_simple() {
+    let shares_num = 16;
+    let security_threshold = 10;
     let (
-        tau,
-        shares_num,
-        security_threshold,
         validator_keypairs,
         validators,
         validators_js,
         messages_js,
         msg,
         aad,
         ciphertext,
-    ) = setup_dkg();
+    ) = setup_dkg(shares_num, security_threshold);
 
     // Having aggregated the transcripts, the validators can now create decryption shares
     let decryption_shares = zip_eq(validators, validator_keypairs)
         .map(|(validator, keypair)| {
             let mut dkg = Dkg::new(
-                tau,
+                TAU,
                 shares_num,
                 security_threshold,
                 &validators_js,
@@ -149,24 +140,23 @@ fn tdec_simple() {
 
 #[wasm_bindgen_test]
 fn tdec_precomputed() {
+    let shares_num = 16;
+    let security_threshold = shares_num; // Must be equal to shares_num in precomputed variant
     let (
-        tau,
-        shares_num,
-        security_threshold,
         validator_keypairs,
         validators,
         validators_js,
         messages_js,
         msg,
         aad,
         ciphertext,
-    ) = setup_dkg();
+    ) = setup_dkg(shares_num, security_threshold);
 
     // Having aggregated the transcripts, the validators can now create decryption shares
     let decryption_shares = zip_eq(validators, validator_keypairs)
         .map(|(validator, keypair)| {
             let mut dkg = Dkg::new(
-                tau,
+                TAU,
                 shares_num,
                 security_threshold,
                 &validators_js,

diff --git a/ferveo/src/api.rs b/ferveo/src/api.rs
@@ -309,6 +309,14 @@ impl AggregatedTranscript {
         aad: &[u8],
         validator_keypair: &Keypair,
     ) -> Result<DecryptionSharePrecomputed> {
+        if dkg.0.dkg_params.shares_num()
+            != dkg.0.dkg_params.security_threshold()
+        {
+            return Err(Error::InvalidDkgParametersForPrecomputedVariant(
+                dkg.0.dkg_params.shares_num(),
+                dkg.0.dkg_params.security_threshold(),
+            ));
+        }
         let domain_points: Vec<_> = dkg
             .0
             .domain
@@ -455,8 +463,6 @@ mod test_ferveo_api {
         let rng = &mut StdRng::seed_from_u64(0);
 
         // In precomputed variant, the security threshold is equal to the number of shares
-        // TODO: Refactor DKG constructor to not require security threshold or this case.
-        //  Or figure out a different way to simplify the precomputed variant API.
         let security_threshold = shares_num;
 
         let (messages, validators, validator_keypairs) =

diff --git a/ferveo/src/bindings_python.rs b/ferveo/src/bindings_python.rs
@@ -104,6 +104,11 @@ impl From<FerveoPythonError> for PyErr {
                         "{index}"
                     ))
                 },
+                Error::InvalidDkgParametersForPrecomputedVariant(num_shares, security_threshold) => {
+                    InvalidDkgParameters::new_err(format!(
+                        "num_shares: {num_shares}, security_threshold: {security_threshold}"
+                    ))
+                },
             },
             _ => default(),
         }

diff --git a/ferveo/src/lib.rs b/ferveo/src/lib.rs
@@ -114,6 +114,10 @@ pub enum Error {
     /// Failed to access a share for a given share index
     #[error("Invalid share index: {0}")]
     InvalidShareIndex(u32),
+
+    /// Failed to produce a precomputed variant decryption share
+    #[error("Invalid DKG parameters for precomputed variant: number of shares {0}, threshold {1}")]
+    InvalidDkgParametersForPrecomputedVariant(u32, u32),
 }
 
 pub type Result<T> = std::result::Result<T, Error>;