Repository: input-output-hk/mithril (rust)

Issue:
Remove VerificationKey and Stake from individual signature
## Why
After the change of removing the merkle path from the individual signature, the protocol relies on the aggregator to have the full Merkle tree to aggregate correctly. Therefore, we don't need the individual signature to contain the public key and stake.

## What
Currently, the signature is defined as:

```
pub struct StmSig {
    pub sigma: Signature,
    pub pk: StmVerificationKey,
    pub stake: Stake,
    pub indexes: Vec<Index>,
    pub signer_index: Index,
}
```

and should be changed to:

```
pub struct StmSig {
    pub sigma: Signature,
    pub indexes: Vec<Index>,
    pub signer_index: Index,
}
```

Pull request:
Simplify signature
## Content
This PR removes the verification key and the stake from the signature. These values are no longer required, as individual signatures no longer contain the merkle paths (it is now responsibility of the aggregator to include these values in the certificate). This required some changes in the multi_signer and single_signer, as these folders were verifying single signatures. 

Now, in order to verify aggregate signatures, they include not only the individual signatures but also the corresponding `RegParty`, to be able to verify the batched membership proof (batched merkle path).

## Pre-submit checklist

- Branch
  - [x] Commit sequence broadly makes sense
  - [x] Key commits have useful messages
- PR
  - [x] No clippy warnings in the CI
  - [x] Self-reviewed the diff
  - [x] Useful pull request description
  - [x] Reviewer requested

## Issue(s)
This PR closes #619 .


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index a1faa3e18..bb5dc1489 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2120,7 +2120,7 @@ dependencies = [
 
 [[package]]
 name = "mithril-stm"
-version = "0.1.0"
+version = "0.2.0"
 dependencies = [
  "bincode",
  "blake2 0.10.5",
diff --git a/mithril-aggregator/src/multi_signer.rs b/mithril-aggregator/src/multi_signer.rs
index fa22a9618..7c0c05c70 100644
--- a/mithril-aggregator/src/multi_signer.rs
+++ b/mithril-aggregator/src/multi_signer.rs
@@ -36,6 +36,10 @@ pub enum ProtocolError {
     #[error("signer already registered")]
     ExistingSigner(),
 
+    /// Signer was not registered.
+    #[error("signer did not register")]
+    UnregisteredParty(),
+
     /// Signer registration failed.
     #[error("signer registration failed")]
     FailedSignerRegistration(#[from] ProtocolRegistrationError),
@@ -696,16 +700,31 @@ 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(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.
+        let (vk, stake) = clerk
+            .get_reg_party(&signature.signer_index)
+            .ok_or_else(ProtocolError::UnregisteredParty)?;
+        signature
+            .verify(
+                &protocol_parameters,
+                &vk,
+                &stake,
+                &avk,
+                message.compute_hash().as_bytes(),
+            )
             .map_err(|e| ProtocolError::Core(e.to_string()))?;
 
         // Register single signature
@@ -714,7 +733,7 @@ impl MultiSigner for MultiSignerImpl {
             .as_ref()
             .ok_or_else(ProtocolError::UnavailableBeacon)?;
 
-        match self
+        return match self
             .single_signature_store
             .save_single_signatures(beacon, signatures)
             .await?
@@ -723,7 +742,7 @@ impl MultiSigner for MultiSignerImpl {
                 signatures.party_id.clone(),
             )),
             None => Ok(()),
-        }
+        };
     }
 
     /// Retrieves a multi signature from a message
diff --git a/mithril-signer/src/single_signer.rs b/mithril-signer/src/single_signer.rs
index 1215c28f0..ceef56f25 100644
--- 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/Cargo.toml b/mithril-stm/Cargo.toml
index 9fcdf44f8..7e4e2b3bd 100644
--- a/mithril-stm/Cargo.toml
+++ b/mithril-stm/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name    = "mithril-stm"
-version = "0.1.0"
+version = "0.2.0"
 edition = { workspace = true }
 authors = { workspace = true }
 documentation = { workspace = true }
diff --git a/mithril-stm/README.md b/mithril-stm/README.md
index 3f76308e7..af4721fb6 100644
--- a/mithril-stm/README.md
+++ b/mithril-stm/README.md
@@ -116,8 +116,8 @@ fn main() {
     let avk = clerk.compute_avk();
 
     // Check all parties can verify every sig
-    for s in sigs.iter() {
-        assert!(s.verify(&params, &avk, &msg).is_ok(), "Verification failed");
+    for (s, p) in sigs.iter().zip(ps.iter()) {
+        assert!(s.verify(&params, &p.verification_key(), &p.get_stake(), &avk, &msg).is_ok(), "Verification failed");
     }
 
     // Aggregate with random parties
@@ -143,7 +143,7 @@ fn main() {
 
 Here we give the benchmark results of STM for size and time. We run the benchmarks on macOS 12.6 on an Apple M1 Pro machine with 16 GB of RAM.
 
-Note that the size of an individual signature with one valid index is **176 bytes** and increases linearly in the length of valid indices (where an index is 8 bytes).
+Note that the size of an individual signature with one valid index is **72 bytes** (48 bytes from `sigma`, 8 bytes from `party_index`, 8 bytes for the `length` of winning indices and at least 8 bytes for a single winning `index`) and increases linearly in the length of valid indices (where an index is 8 bytes).
 
 ```shell
 +----------------------+
diff --git a/mithril-stm/src/merkle_tree.rs b/mithril-stm/src/merkle_tree.rs
index b67395b17..14e05502e 100644
--- a/mithril-stm/src/merkle_tree.rs
+++ b/mithril-stm/src/merkle_tree.rs
@@ -1,8 +1,9 @@
 //! Creation and verification of Merkle Trees
 use crate::error::MerkleTreeError;
 use crate::multi_sig::VerificationKey;
-use crate::stm::Stake;
-use blake2::digest::{Digest, FixedOutput};
+use crate::stm::{Stake, StmVerificationKey};
+use blake2::digest::{consts::U32, Digest, FixedOutput};
+use blake2::Blake2b;
 use serde::{Deserialize, Serialize};
 use std::cmp::Ordering;
 use std::convert::TryFrom;
@@ -74,6 +75,14 @@ pub struct MerkleTree<D: Digest> {
 }
 
 impl MTLeaf {
+    pub(crate) fn from_bytes(bytes: &[u8]) -> Result<Self, MerkleTreeError<Blake2b<U32>>> {
+        let pk = StmVerificationKey::from_bytes(bytes)
+            .map_err(|_| MerkleTreeError::SerializationError)?;
+        let mut u64_bytes = [0u8; 8];
+        u64_bytes.copy_from_slice(&bytes[96..]);
+        let stake = Stake::from_be_bytes(u64_bytes);
+        Ok(MTLeaf(pk, stake))
+    }
     pub(crate) fn to_bytes(self) -> [u8; 104] {
         let mut result = [0u8; 104];
         result[..96].copy_from_slice(&self.0.to_bytes());
@@ -82,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
index 047da6425..93ec2c60c 100644
--- a/mithril-stm/src/stm.rs
+++ b/mithril-stm/src/stm.rs
@@ -110,7 +110,7 @@ use crate::dense_mapping::ev_lt_phi;
 use crate::error::{
     AggregationError, RegisterError, StmAggregateSignatureError, StmSignatureError,
 };
-use crate::key_reg::ClosedKeyReg;
+use crate::key_reg::{ClosedKeyReg, RegParty};
 use crate::merkle_tree::{BatchPath, MTLeaf, MerkleTreeCommitmentBatchCompat};
 use crate::multi_sig::{Signature, SigningKey, VerificationKey, VerificationKeyPoP};
 use blake2::digest::{Digest, FixedOutput};
@@ -193,10 +193,6 @@ pub struct StmClerk<D: Clone + Digest> {
 pub struct StmSig {
     /// The signature from the underlying MSP scheme.
     pub sigma: Signature,
-    /// The Stm verification Key.
-    pub pk: StmVerificationKey,
-    /// The stake of the party that made this signature.
-    pub stake: Stake,
     /// The index(es) for which the signature is valid
     pub indexes: Vec<Index>,
     /// Merkle tree index of the signer.
@@ -224,7 +220,7 @@ pub struct StmAggrVerificationKey<D: Clone + Digest + FixedOutput> {
     deserialize = "BatchPath<D>: Deserialize<'de>"
 ))]
 pub struct StmAggrSig<D: Clone + Digest + FixedOutput> {
-    pub(crate) signatures: Vec<StmSig>,
+    pub(crate) signatures: Vec<(StmSig, RegParty)>,
     /// The list of unique merkle tree nodes that covers path for all signatures.
     pub batch_proof: BatchPath<D>,
 }
@@ -384,8 +380,6 @@ impl<D: Clone + Digest + FixedOutput> StmSigner<D> {
         if !indexes.is_empty() {
             Some(StmSig {
                 sigma,
-                pk: self.vk,
-                stake: self.stake,
                 indexes,
                 signer_index: self.mt_index,
             })
@@ -404,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> {
@@ -436,6 +440,15 @@ impl<D: Digest + Clone + FixedOutput> StmClerk<D> {
     ) -> Result<StmAggrSig<D>, AggregationError> {
         let mut unique_sigs = self.dedup_sigs_for_indices(msg, sigs)?;
         unique_sigs.sort_unstable();
+        let signatures = unique_sigs
+            .iter()
+            .map(|sig| {
+                (
+                    sig.clone(),
+                    self.closed_reg.reg_parties[sig.signer_index as usize],
+                )
+            })
+            .collect(); // todo: look into this conversion
 
         let mt_index_list = unique_sigs
             .iter()
@@ -445,7 +458,7 @@ impl<D: Digest + Clone + FixedOutput> StmClerk<D> {
         let batch_proof = self.closed_reg.merkle_tree.get_batched_path(mt_index_list);
 
         Ok(StmAggrSig {
-            signatures: unique_sigs,
+            signatures,
             batch_proof,
         })
     }
@@ -466,9 +479,16 @@ impl<D: Digest + Clone + FixedOutput> StmClerk<D> {
         let avk = StmAggrVerificationKey::from(&self.closed_reg);
         let mut sig_by_index: BTreeMap<Index, &StmSig> = BTreeMap::new();
         let mut removal_idx_by_vk: HashMap<&StmSig, Vec<Index>> = HashMap::new();
-
-        for sig in sigs.iter() {
-            if sig.verify(&self.params, &avk, msg).is_err() {
+        let reg_parties = sigs
+            .iter()
+            .map(|sig| self.closed_reg.reg_parties[sig.signer_index as usize])
+            .collect::<Vec<RegParty>>(); // todo: look into this
+
+        for (sig, reg_party) in sigs.iter().zip(reg_parties.iter()) {
+            if sig
+                .verify(&self.params, &reg_party.0, &reg_party.1, &avk, msg)
+                .is_err()
+            {
                 continue;
             }
 
@@ -535,6 +555,14 @@ 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 its index.
+    pub fn get_reg_party(&self, party_index: &Index) -> Option<(StmVerificationKey, Stake)> {
+        self.closed_reg
+            .reg_parties
+            .get(*party_index as usize)
+            .map(|&r| r.into())
+    }
 }
 
 impl StmSig {
@@ -543,12 +571,14 @@ impl StmSig {
     pub fn verify<D: Clone + Digest + FixedOutput>(
         &self,
         params: &StmParameters,
+        pk: &StmVerificationKey,
+        stake: &Stake,
         avk: &StmAggrVerificationKey<D>,
         msg: &[u8],
     ) -> Result<(), StmSignatureError> {
         let msgp = avk.mt_commitment.concat_with_msg(msg);
-        self.sigma.verify(&msgp, &self.pk)?;
-        self.check_indices(params, &msgp, avk)?;
+        self.sigma.verify(&msgp, pk)?;
+        self.check_indices(params, stake, &msgp, avk)?;
 
         Ok(())
     }
@@ -557,6 +587,7 @@ impl StmSig {
     pub(crate) fn check_indices<D: Clone + Digest + FixedOutput>(
         &self,
         params: &StmParameters,
+        stake: &Stake,
         msgp: &[u8],
         avk: &StmAggrVerificationKey<D>,
     ) -> Result<(), StmSignatureError> {
@@ -567,7 +598,7 @@ impl StmSig {
 
             let ev = self.sigma.eval(msgp, index);
 
-            if !ev_lt_phi(params.phi_f, ev, self.stake, avk.total_stake) {
+            if !ev_lt_phi(params.phi_f, ev, *stake, avk.total_stake) {
                 return Err(StmSignatureError::LotteryLost);
             }
         }
@@ -586,14 +617,12 @@ impl StmSig {
     /// * Merkle index of the signer.
     pub fn to_bytes(&self) -> Vec<u8> {
         let mut output = Vec::new();
-        output.extend_from_slice(&self.stake.to_be_bytes());
         output.extend_from_slice(&(self.indexes.len() as u64).to_be_bytes());
 
         for index in &self.indexes {
             output.extend_from_slice(&index.to_be_bytes());
         }
 
-        output.extend_from_slice(&self.pk.to_bytes());
         output.extend_from_slice(&self.sigma.to_bytes());
 
         output.extend_from_slice(&self.signer_index.to_be_bytes());
@@ -606,29 +635,23 @@ impl StmSig {
     ) -> Result<StmSig, StmSignatureError> {
         let mut u64_bytes = [0u8; 8];
 
-        u64_bytes.copy_from_slice(&bytes[..8]);
-        let stake = u64::from_be_bytes(u64_bytes);
-
-        u64_bytes.copy_from_slice(&bytes[8..16]);
+        u64_bytes.copy_from_slice(&bytes[0..8]);
         let nr_indexes = u64::from_be_bytes(u64_bytes) as usize;
 
         let mut indexes = Vec::new();
         for i in 0..nr_indexes {
-            u64_bytes.copy_from_slice(&bytes[16 + i * 8..24 + i * 8]);
+            u64_bytes.copy_from_slice(&bytes[8 + i * 8..16 + i * 8]);
             indexes.push(u64::from_be_bytes(u64_bytes));
         }
 
-        let offset = 16 + nr_indexes * 8;
-        let pk = StmVerificationKey::from_bytes(&bytes[offset..offset + 96])?;
-        let sigma = Signature::from_bytes(&bytes[offset + 96..offset + 144])?;
+        let offset = 8 + nr_indexes * 8;
+        let sigma = Signature::from_bytes(&bytes[offset..offset + 48])?;
 
-        u64_bytes.copy_from_slice(&bytes[offset + 144..offset + 152]);
+        u64_bytes.copy_from_slice(&bytes[offset + 48..offset + 56]);
         let signer_index = u64::from_be_bytes(u64_bytes);
 
         Ok(StmSig {
             sigma,
-            pk,
-            stake,
             indexes,
             signer_index,
         })
@@ -692,17 +715,16 @@ impl<D: Clone + Digest + FixedOutput> StmAggrSig<D> {
         let mut unique_indices = HashSet::new();
         let mut leaves = Vec::new();
 
-        for sig in &self.signatures {
+        for (sig, reg_party) in &self.signatures {
             let msgp = avk.mt_commitment.concat_with_msg(msg);
-            sig.check_indices(parameters, &msgp, avk)?;
+            sig.check_indices(parameters, &reg_party.1, &msgp, avk)?;
 
             for &index in &sig.indexes {
                 unique_indices.insert(index);
                 nr_indices += 1;
             }
 
-            let mt_leaf = MTLeaf(sig.pk, sig.stake);
-            leaves.push(mt_leaf);
+            leaves.push(*reg_party);
         }
 
         if nr_indices != unique_indices.len() {
@@ -720,12 +742,12 @@ impl<D: Clone + Digest + FixedOutput> StmAggrSig<D> {
         let signatures = self
             .signatures
             .iter()
-            .map(|sig| sig.sigma)
+            .map(|(sig, _)| sig.sigma)
             .collect::<Vec<Signature>>();
         let vks = self
             .signatures
             .iter()
-            .map(|sig| sig.pk)
+            .map(|(_, reg_party)| reg_party.0)
             .collect::<Vec<VerificationKey>>();
 
         Signature::verify_aggregate(msg.as_slice(), &vks, &signatures)?;
@@ -736,18 +758,19 @@ impl<D: Clone + Digest + FixedOutput> StmAggrSig<D> {
     /// # Layout
     /// * Number of signatures (as u64)
     /// * Size of a signature
-    /// * Signatures
+    /// * Pairs of Signatures and Registered Parties
     /// * Batch proof
     pub fn to_bytes(&self) -> Vec<u8> {
         let mut out = Vec::new();
         out.extend_from_slice(&u64::try_from(self.signatures.len()).unwrap().to_be_bytes());
         out.extend_from_slice(
-            &u64::try_from(self.signatures[0].to_bytes().len())
+            &u64::try_from(self.signatures[0].0.to_bytes().len())
                 .unwrap()
                 .to_be_bytes(),
         );
-        for sig in &self.signatures {
-            out.extend_from_slice(&sig.to_bytes())
+        for (sig, reg_party) in &self.signatures {
+            out.extend_from_slice(&sig.to_bytes());
+            out.extend_from_slice(&reg_party.to_bytes());
         }
         let proof = &self.batch_proof;
         out.extend_from_slice(&proof.to_bytes());
@@ -769,12 +792,18 @@ impl<D: Clone + Digest + FixedOutput> StmAggrSig<D> {
 
         let mut signatures = Vec::with_capacity(size);
         for i in 0..size {
-            signatures.push(StmSig::from_bytes::<D>(
-                &bytes[16 + i * sig_size..16 + (i + 1) * sig_size],
-            )?);
+            signatures.push((
+                StmSig::from_bytes::<D>(
+                    &bytes[16 + i * (sig_size + 104)..16 + sig_size + i * (sig_size + 104)],
+                )?,
+                MTLeaf::from_bytes(
+                    &bytes[16 + sig_size + i * (sig_size + 104)..16 + (i + 1) * (sig_size + 104)],
+                )
+                .map_err(|_| StmAggregateSignatureError::SerializationError)?,
+            ));
         }
 
-        let offset = 16 + sig_size * size;
+        let offset = 16 + (sig_size + 104) * size;
         let batch_proof = BatchPath::from_bytes(&bytes[offset..])?;
 
         Ok(StmAggrSig {
@@ -903,28 +932,26 @@ mod tests {
         #[test]
         /// Test that `dedup_sigs_for_indices` only takes valid signatures.
         fn test_dedup(msg in any::<[u8; 16]>()) {
-            let nparties = 10usize;
             let false_msg = [1u8; 20];
-            let params = StmParameters { m: (nparties as u64), k: 1, phi_f: 1.0 };
-            let ps = setup_equal_parties(params, nparties);
-            let p = &ps[0];
-            let clerk = StmClerk::from_signer(p);
+            let params = StmParameters { m: 1, k: 1, phi_f: 1.0 };
+            let ps = setup_equal_parties(params, 1);
+            let clerk = StmClerk::from_signer(&ps[0]);
             let avk = clerk.compute_avk();
-            let mut sigs = Vec::with_capacity(nparties);
+            let mut sigs = Vec::with_capacity(2);
 
 
-            if let Some(sig) = p.sign(&false_msg) {
+            if let Some(sig) = ps[0].sign(&false_msg) {
                 sigs.push(sig);
             }
 
-            if let Some(sig) = p.sign(&msg) {
+            if let Some(sig) = ps[0].sign(&msg) {
                 sigs.push(sig);
             }
 
             let dedup_result = clerk.dedup_sigs_for_indices(&msg, &sigs);
             assert!(dedup_result.is_ok(), "dedup failure {:?}", dedup_result);
             for passed_sigs in dedup_result.unwrap() {
-                let verify_result = passed_sigs.verify(&params, &avk, &msg);
+                let verify_result = passed_sigs.verify(&params, &ps[0].vk, &ps[0].stake, &avk, &msg);
                 assert!(verify_result.is_ok(), "verify {:?}", verify_result);
             }
         }
@@ -934,15 +961,13 @@ mod tests {
         #[test]
         /// Test that when a party creates a signature it can be verified
         fn test_sig(msg in any::<[u8;16]>()) {
-            let nparties = 2;
-            let params = StmParameters { m: (nparties as u64), k: 1, phi_f: 0.2 };
-            let ps = setup_equal_parties(params, nparties);
-            let p = &ps[0];
-            let clerk = StmClerk::from_signer(p);
+            let params = StmParameters { m: 1, k: 1, phi_f: 0.2 };
+            let ps = setup_equal_parties(params, 1);
+            let clerk = StmClerk::from_signer(&ps[0]);
             let avk = clerk.compute_avk();
 
-            if let Some(sig) = p.sign(&msg) {
-                assert!(sig.verify(&params, &avk, &msg).is_ok());
+            if let Some(sig) = ps[0].sign(&msg) {
+                assert!(sig.verify(&params, &ps[0].vk, &ps[0].stake, &avk, &msg).is_ok());
             }
         }
     }
@@ -1005,21 +1030,19 @@ mod tests {
 
         #[test]
         fn test_sig_serialize_deserialize(msg in any::<[u8;16]>()) {
-            let nparties = 2;
-            let params = StmParameters { m: (nparties as u64), k: 1, phi_f: 0.2 };
-            let ps = setup_equal_parties(params, nparties);
-            let p = &ps[0];
-            let clerk = StmClerk::from_signer(p);
+            let params = StmParameters { m: 1, k: 1, phi_f: 0.2 };
+            let ps = setup_equal_parties(params, 1);
+            let clerk = StmClerk::from_signer(&ps[0]);
             let avk = clerk.compute_avk();
 
-            if let Some(sig) = p.sign(&msg) {
+            if let Some(sig) = ps[0].sign(&msg) {
                 let bytes = sig.to_bytes();
                 let sig_deser = StmSig::from_bytes::<D>(&bytes).unwrap();
-                assert!(sig_deser.verify(&params, &avk, &msg).is_ok());
+                assert!(sig_deser.verify(&params, &ps[0].vk, &ps[0].stake, &avk, &msg).is_ok());
 
                 let encoded = bincode::serialize(&sig).unwrap();
                 let decoded: StmSig = bincode::deserialize(&encoded).unwrap();
-                assert!(decoded.verify(&params, &avk, &msg).is_ok());
+                assert!(decoded.verify(&params, &ps[0].vk, &ps[0].stake, &avk, &msg).is_ok());
             }
         }
 
@@ -1164,7 +1187,7 @@ mod tests {
         fn test_invalid_proof_index_unique(tc in arb_proof_setup(10)) {
             with_proof_mod(tc, |aggr, clerk, _msg| {
                 for sig in aggr.signatures.iter_mut() {
-                    for index in sig.indexes.iter_mut() {
+                    for index in sig.0.indexes.iter_mut() {
                        *index %= clerk.params.k - 1
                     }
                 }


Test patch:
diff --git a/mithril-stm/tests/integration.rs b/mithril-stm/tests/integration.rs
index 6a8758f16..c29d64f45 100644
--- a/mithril-stm/tests/integration.rs
+++ b/mithril-stm/tests/integration.rs
@@ -1,5 +1,7 @@
 use mithril_stm::key_reg::KeyReg;
-use mithril_stm::stm::{StmClerk, StmInitializer, StmParameters, StmSig, StmSigner};
+use mithril_stm::stm::{
+    Stake, StmClerk, StmInitializer, StmParameters, StmSig, StmSigner, StmVerificationKey,
+};
 use mithril_stm::AggregationError;
 
 use blake2::{digest::consts::U32, Blake2b};
@@ -34,9 +36,11 @@ fn test_full_protocol() {
     let mut key_reg = KeyReg::init();
 
     let mut ps: Vec<StmInitializer> = Vec::with_capacity(nparties as usize);
+    let mut reg_parties: Vec<(StmVerificationKey, Stake)> = Vec::with_capacity(nparties as usize);
     for stake in parties {
         let p = StmInitializer::setup(params, stake, &mut rng);
         key_reg.register(stake, p.verification_key()).unwrap();
+        reg_parties.push((p.verification_key().vk, stake));
         ps.push(p);
     }
 
@@ -60,8 +64,11 @@ fn test_full_protocol() {
     let avk = clerk.compute_avk();
 
     // Check all parties can verify every sig
-    for s in sigs.iter() {
-        assert!(s.verify(&params, &avk, &msg).is_ok(), "Verification failed");
+    for (s, (vk, stake)) in sigs.iter().zip(reg_parties.iter()) {
+        assert!(
+            s.verify(&params, vk, stake, &avk, &msg).is_ok(),
+            "Verification failed"
+        );
     }
 
     // Aggregate with random parties