Adapt rest of the repo to changes

input-output-hk · Dec 2, 2022 · 1d39536 · 1d39536
1 parent 7f13680
commit 1d39536
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 24 deletions.
diff --git a/mithril-aggregator/src/multi_signer.rs b/mithril-aggregator/src/multi_signer.rs
@@ -703,34 +703,54 @@ impl MultiSigner for MultiSignerImpl {
             .get_protocol_parameters()
             .await?
             .ok_or_else(ProtocolError::UnavailableProtocolParameters)?;
-        let avk = &self
+
+        let clerk = self
             .clerk
             .as_ref()
-            .ok_or_else(ProtocolError::UnavailableClerk)?
-            .compute_avk();
+            .ok_or_else(ProtocolError::UnavailableClerk)?;
 
-        signatures
+        let signature = signatures
             .to_protocol_signature()
-            .map_err(ProtocolError::Codec)?
-            .verify(&protocol_parameters, avk, message.compute_hash().as_bytes())
-            .map_err(|e| ProtocolError::Core(e.to_string()))?;
-
-        // Register single signature
-        let beacon = self
-            .current_beacon
-            .as_ref()
-            .ok_or_else(ProtocolError::UnavailableBeacon)?;
+            .map_err(ProtocolError::Codec)?;
+
+        let avk = clerk.compute_avk();
+
+        // If there is no reg_party, then we simply received a signature from a non-registered
+        // party, and we can ignore the request.
+        if let Some((vk, stake)) = clerk.get_reg_party(&signature.signer_index) {
+            signature
+                .verify(
+                    &protocol_parameters,
+                    &vk,
+                    &stake,
+                    &avk,
+                    message.compute_hash().as_bytes(),
+                )
+                .map_err(|e| ProtocolError::Core(e.to_string()))?;
 
-        match self
-            .single_signature_store
-            .save_single_signatures(beacon, signatures)
-            .await?
-        {
-            Some(_) => Err(ProtocolError::ExistingSingleSignature(
-                signatures.party_id.clone(),
-            )),
-            None => Ok(()),
+            // Register single signature
+            let beacon = self
+                .current_beacon
+                .as_ref()
+                .ok_or_else(ProtocolError::UnavailableBeacon)?;
+
+            match self
+                .single_signature_store
+                .save_single_signatures(beacon, signatures)
+                .await?
+            {
+                Some(_) => {
+                    return Err(ProtocolError::ExistingSingleSignature(
+                        signatures.party_id.clone(),
+                    ));
+                }
+                None => {
+                    return Ok(());
+                }
+            }
         }
+
+        Ok(())
     }
 
     /// Retrieves a multi signature from a message

diff --git a/mithril-signer/src/single_signer.rs b/mithril-signer/src/single_signer.rs
@@ -271,7 +271,13 @@ mod tests {
         let decoded_sig: ProtocolSingleSignature = key_decode_hex(&sign_result.signature).unwrap();
         assert!(
             decoded_sig
-                .verify(&protocol_parameters, &avk, &expected_message)
+                .verify(
+                    &protocol_parameters,
+                    &protocol_signer.verification_key(),
+                    &protocol_signer.get_stake(),
+                    &avk,
+                    &expected_message
+                )
                 .is_ok(),
             "produced single signature should be valid"
         );

diff --git a/mithril-stm/src/merkle_tree.rs b/mithril-stm/src/merkle_tree.rs
@@ -91,6 +91,12 @@ impl MTLeaf {
     }
 }
 
+impl From<MTLeaf> for (StmVerificationKey, Stake) {
+    fn from(leaf: MTLeaf) -> (StmVerificationKey, Stake) {
+        (leaf.0, leaf.1)
+    }
+}
+
 impl PartialOrd for MTLeaf {
     /// Ordering of MT Values.
     ///

diff --git a/mithril-stm/src/stm.rs b/mithril-stm/src/stm.rs
@@ -169,7 +169,6 @@ pub struct StmInitializer {
 /// Participant in the protocol can sign messages.
 /// This instance can only be generated out of an `StmInitializer` and a `ClosedKeyReg`.
 /// This ensures that a `MerkleTree` root is not computed before all participants have registered.
-#[allow(dead_code)]
 #[derive(Debug, Clone)]
 pub struct StmSigner<D: Digest> {
     mt_index: u64,
@@ -399,6 +398,16 @@ impl<D: Clone + Digest + FixedOutput> StmSigner<D> {
     pub fn get_closed_reg(self) -> ClosedKeyReg<D> {
         self.closed_reg
     }
+
+    /// Extract the verification key.
+    pub fn verification_key(&self) -> StmVerificationKey {
+        self.vk
+    }
+
+    /// Extract stake from the signer.
+    pub fn get_stake(&self) -> Stake {
+        self.stake
+    }
 }
 
 impl<D: Digest + Clone + FixedOutput> StmClerk<D> {
@@ -546,6 +555,15 @@ impl<D: Digest + Clone + FixedOutput> StmClerk<D> {
     pub fn compute_avk(&self) -> StmAggrVerificationKey<D> {
         StmAggrVerificationKey::from(&self.closed_reg)
     }
+
+    /// Get the (VK, stake) of a party given it's index.
+    pub fn get_reg_party(&self, party_index: &Index) -> Option<(StmVerificationKey, Stake)> {
+        if *party_index as usize >= self.closed_reg.reg_parties.len() {
+            return None;
+        }
+
+        Some(self.closed_reg.reg_parties[*party_index as usize].into())
+    }
 }
 
 impl StmSig {