aggregator: Convert unwraps to ?.

chainwayxyz · Sep 9, 2024 · caff0e6 · caff0e6
1 parent 14f002e
commit caff0e6
Showing 1 changed file with 16 additions and 12 deletions.
diff --git a/core/src/aggregator.rs b/core/src/aggregator.rs
@@ -9,7 +9,7 @@ use crate::{
     traits::rpc::AggregatorServer,
     transaction_builder::TransactionBuilder,
     utils::{self, handle_taproot_witness_new},
-    ByteArray32, EVMAddress, UTXO,
+    ByteArray32, ByteArray66, EVMAddress, UTXO,
 };
 use async_trait::async_trait;
 use bitcoin::{address::NetworkUnchecked, Address, OutPoint};
@@ -210,15 +210,15 @@ impl Aggregator {
     ) -> Result<Vec<MuSigAggNonce>, BridgeError> {
         let mut agg_nonces = Vec::new();
         for i in 0..pub_nonces[0].len() {
-            let agg_nonce = aggregate_nonces(
-                pub_nonces
-                    .iter()
-                    .map(|v| v.get(i).cloned().unwrap()) // TODO: remove unwrap
-                    .collect::<Vec<_>>(),
-            );
+            let pub_nonces = pub_nonces
+                .iter()
+                .map(|v| v.get(i).cloned())
+                .collect::<Option<Vec<ByteArray66>>>()
+                .ok_or(BridgeError::NoncesNotFound)?;
 
-            agg_nonces.push(agg_nonce);
+            agg_nonces.push(aggregate_nonces(pub_nonces));
         }
+
         Ok(agg_nonces)
     }
 
@@ -236,18 +236,22 @@ impl Aggregator {
             agg_nonces,
             partial_sigs
         );
+
         let mut slash_or_take_sigs = Vec::new();
         for i in 0..partial_sigs[0].len() {
+            let partial_sigs = partial_sigs
+                .iter()
+                .map(|v| v.get(i).cloned())
+                .collect::<Option<Vec<ByteArray32>>>()
+                .ok_or(BridgeError::NoncesNotFound)?;
+
             let agg_sig = self.aggregate_slash_or_take_partial_sigs(
                 deposit_outpoint,
                 kickoff_utxos[i].clone(),
                 self.config.operators_xonly_pks[i],
                 i,
                 &agg_nonces[i].clone(),
-                partial_sigs
-                    .iter()
-                    .map(|v| v.get(i).cloned().unwrap()) // TODO: remove unwrap
-                    .collect::<Vec<_>>(),
+                partial_sigs,
             )?;
 
             slash_or_take_sigs.push(secp256k1::schnorr::Signature::from_slice(&agg_sig)?);