Move extract_key method to Subset

Author: akonring

timeboost-sequencer/src/decrypt.rs

@@ -718,11 +718,11 @@ impl Worker {
         }
 
         if let Some((&subset, _)) = counts.iter().find(|(_, count)| **count >= threshold) {
-            let acc = DkgAccumulator::from_subset(current.to_owned(), subset.to_owned());
-            let dec_key = acc
-                .extract_key(&self.dkg_sk, prev)
-                .map_err(|e| DecrypterError::Dkg(e.to_string()))?;
+            let acc = DkgAccumulator::from_subset(current.clone(), subset.to_owned());
             self.tracker.insert(committee.id(), acc);
+            let dec_key = subset
+                .extract_key(current.to_owned(), &self.dkg_sk, prev)
+                .map_err(|e| DecrypterError::Dkg(e.to_string()))?;
 
             self.dec_key.set(dec_key);
             self.state = WorkerState::Running;
@@ -765,11 +765,11 @@ impl Worker {
         }
 
         if let Some((&subset, _)) = counts.iter().find(|(_, count)| **count >= threshold) {
-            let acc = DkgAccumulator::from_subset(current.to_owned(), subset.to_owned());
-            let next_dec_key = acc
-                .extract_key(&self.dkg_sk, Some(prev.to_owned()))
-                .map_err(|e| DecrypterError::Dkg(e.to_string()))?;
+            let acc = DkgAccumulator::from_subset(current.clone(), subset.to_owned());
             self.tracker.insert(current.committee().id(), acc);
+            let next_dec_key = subset
+                .extract_key(current.clone(), &self.dkg_sk, Some(prev.to_owned()))
+                .map_err(|e| DecrypterError::Dkg(e.to_string()))?;
 
             info!(committee_id = %current.committee().id(), node = %self.label, "handover finished");
             self.state = WorkerState::HandoverComplete;
@@ -828,14 +828,14 @@ impl Worker {
         let acc = self
             .tracker
             .entry(*committee_id)
-            .or_insert_with(|| DkgAccumulator::new_dkg(key_store.to_owned()));
+            .or_insert_with(|| DkgAccumulator::new_dkg(key_store.clone()));
 
         acc.try_add(bundle)
             .map_err(|e| DecrypterError::Dkg(format!("unable to add dkg bundle: {e}")))?;
 
-        if acc.try_finalize().is_some() {
-            let dec_key = acc
-                .extract_key(&self.dkg_sk, None)
+        if let Some(subset) = acc.try_finalize() {
+            let dec_key = subset
+                .extract_key(key_store.to_owned(), &self.dkg_sk, None)
                 .map_err(|e| DecrypterError::Dkg(e.to_string()))?;
             self.dec_key.set(dec_key);
             self.state = WorkerState::Running;
@@ -888,8 +888,8 @@ impl Worker {
 
             if committee.contains_key(&self.label) {
                 // node is a member of the next committee; decrypting reshares immediately
-                let next_dec_key = acc
-                    .extract_key(&self.dkg_sk, Some(current))
+                let next_dec_key = subset
+                    .extract_key(next, &self.dkg_sk, Some(current))
                     .map_err(|e| DecrypterError::Dkg(e.to_string()))?;
                 self.state = WorkerState::ResharingComplete(next_dec_key);
             } else {

timeboost-types/src/decryption.rs

@@ -484,70 +484,6 @@ impl DkgAccumulator {
             complete: true,
         }
     }
-
-    /// Extract the new threshold decryption key from the accumulator.
-    pub fn extract_key(
-        &self,
-        dkg_sk: &LabeledDkgDecKey,
-        prev: Option<KeyStore>,
-    ) -> anyhow::Result<DecryptionKey> {
-        if !self.completed() {
-            return Err(anyhow!(
-                "attempt to extract key from incomplete accumulator"
-            ));
-        };
-        let vess = Vess::new_fast();
-
-        match &self.mode {
-            AccumulatorMode::Dkg => {
-                let mut dealings_iter = ResultIter::new(self.bundles().iter().map(|b| {
-                    vess.decrypt_share(self.committee(), dkg_sk, b.vess_ct(), DKG_AAD)
-                        .map(|s| (s, b.comm().clone()))
-                }));
-
-                let dec_key = DecryptionKey::from_dkg(
-                    self.committee().size().into(),
-                    dkg_sk.node_idx(),
-                    &mut dealings_iter,
-                )?;
-
-                dealings_iter.result()?;
-
-                Ok(dec_key)
-            }
-            AccumulatorMode::Resharing(combkey) => {
-                let Some(prev) = prev else {
-                    return Err(anyhow!("previous key store missing"));
-                };
-                let dealings: Vec<_> = self
-                    .bundles()
-                    .iter()
-                    .enumerate()
-                    .map(|(i, b)| {
-                        let node_idx = b.origin().0.into();
-                        let pub_share = combkey
-                            .get_pub_share(node_idx)
-                            .ok_or(VessError::FailedVerification)?;
-                        let s = vess.decrypt_reshare(
-                            self.committee(),
-                            dkg_sk,
-                            b.vess_ct(),
-                            DKG_AAD,
-                            *pub_share,
-                        )?;
-                        Ok((i, s, b.comm().clone()))
-                    })
-                    .collect::<Result<Vec<_>, VessError>>()?;
-
-                DecryptionKey::from_resharing(
-                    prev.committee(),
-                    self.committee(),
-                    dkg_sk.node_idx(),
-                    dealings.into_iter(),
-                )
-            }
-        }
-    }
 }
 
 /// A unified subset that can represent both DKG and Resharing results.
@@ -605,6 +541,66 @@ impl DkgSubset {
     pub fn is_resharing(&self) -> bool {
         self.combkey.is_some()
     }
+
+    /// Extract the new threshold decryption key from the subset.
+    pub fn extract_key(
+        &self,
+        curr: KeyStore,
+        dkg_sk: &LabeledDkgDecKey,
+        prev: Option<KeyStore>,
+    ) -> anyhow::Result<DecryptionKey> {
+        let vess = Vess::new_fast();
+
+        match &self.combkey {
+            None => {
+                let mut dealings_iter = ResultIter::new(self.bundles().iter().map(|b| {
+                    vess.decrypt_share(curr.committee(), dkg_sk, b.vess_ct(), DKG_AAD)
+                        .map(|s| (s, b.comm().clone()))
+                }));
+
+                let dec_key = DecryptionKey::from_dkg(
+                    curr.committee().size().into(),
+                    dkg_sk.node_idx(),
+                    &mut dealings_iter,
+                )?;
+
+                dealings_iter.result()?;
+
+                Ok(dec_key)
+            }
+            Some(combkey) => {
+                let Some(prev) = prev else {
+                    return Err(anyhow!("previous key store missing"));
+                };
+                let dealings: Vec<_> = self
+                    .bundles()
+                    .iter()
+                    .enumerate()
+                    .map(|(i, b)| {
+                        let node_idx = b.origin().0.into();
+                        let pub_share = combkey
+                            .get_pub_share(node_idx)
+                            .ok_or(VessError::FailedVerification)?;
+                        let s = vess.decrypt_reshare(
+                            curr.committee(),
+                            dkg_sk,
+                            b.vess_ct(),
+                            DKG_AAD,
+                            *pub_share,
+                        )?;
+                        Ok((i, s, b.comm().clone()))
+                    })
+                    .collect::<Result<Vec<_>, VessError>>()?;
+
+                DecryptionKey::from_resharing(
+                    prev.committee(),
+                    curr.committee(),
+                    dkg_sk.node_idx(),
+                    dealings.into_iter(),
+                )
+            }
+        }
+    }
 }
 
 /// Wrapper iterator that bridges type conversion