From f283187a72478c500217cd558535e88c84e3cf5a Mon Sep 17 00:00:00 2001
From: Piotr Roslaniec <p.roslaniec@gmail.com>
Date: Fri, 19 Jan 2024 15:19:44 +0100
Subject: [PATCH] fix: prevent precomputed shares from being created with
 inapprioriate variant

---
 ferveo/src/api.rs             | 10 ++++++++--
 ferveo/src/bindings_python.rs |  5 +++++
 ferveo/src/lib.rs             |  4 ++++
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/ferveo/src/api.rs b/ferveo/src/api.rs
index 0a8bf2aa..3c2295d1 100644
--- 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
index 411b42c5..f897c8f6 100644
--- 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
index 394afb1a..c316c815 100644
--- 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>;