refactor(stm): move proof type encoding byte to aggregate signature

Author: jpraynaud

mithril-stm/src/aggregate_signature/proof/concatenation.rs

@@ -59,7 +59,7 @@ impl<D: Clone + Digest + FixedOutput + Send + Sync> ConcatenationProof<D> {
         ))
     }
 
-    /// Verify aggregate signature, by checking that
+    /// Verify concatenation proof, by checking that
     /// * each signature contains only valid indices,
     /// * the lottery is indeed won by each one of them,
     /// * the merkle tree path is valid,
@@ -129,19 +129,13 @@ impl<D: Clone + Digest + FixedOutput + Send + Sync> ConcatenationProof<D> {
         Ok(())
     }
 
-    /// Convert multi signature to bytes
+    /// Convert concatenation proof to bytes
     /// # Layout
-    /// * Aggregate signature type (u8)
     /// * Number of the pairs of Signatures and Registered Parties (SigRegParty) (as u64)
     /// * Pairs of Signatures and Registered Parties (prefixed with their size as u64)
     /// * Batch proof
     pub fn to_bytes(&self) -> Vec<u8> {
         let mut out = Vec::new();
-        // This proof type is not strictly necessary, but it will help to identify
-        // the type of the proof used to aggregate when implementing multiple aggregation proof systems.
-        // We use '0' for concatenation proof.
-        let proof_type: u8 = 0;
-        out.extend_from_slice(&proof_type.to_be_bytes());
         out.extend_from_slice(&u64::try_from(self.signatures.len()).unwrap().to_be_bytes());
         for sig_reg in &self.signatures {
             out.extend_from_slice(&u64::try_from(sig_reg.to_bytes().len()).unwrap().to_be_bytes());
@@ -153,21 +147,12 @@ impl<D: Clone + Digest + FixedOutput + Send + Sync> ConcatenationProof<D> {
         out
     }
 
-    ///Extract a `AggregateSignature` from a byte slice.
+    ///Extract a concatenation proof from a byte slice.
     pub fn from_bytes(
         bytes: &[u8],
     ) -> Result<ConcatenationProof<D>, StmAggregateSignatureError<D>> {
-        let mut u8_bytes = [0u8; 1];
         let mut bytes_index = 0;
 
-        u8_bytes
-            .copy_from_slice(bytes.get(..1).ok_or(StmAggregateSignatureError::SerializationError)?);
-        bytes_index += 1;
-        let proof_type = u8::from_be_bytes(u8_bytes);
-        if proof_type != 0 {
-            return Err(StmAggregateSignatureError::SerializationError);
-        }
-
         let mut u64_bytes = [0u8; 8];
         u64_bytes.copy_from_slice(
             bytes

mithril-stm/src/aggregate_signature/signature.rs

@@ -1,21 +1,64 @@
 use std::collections::HashMap;
+use std::fmt::Display;
 use std::hash::Hash;
 
 use blake2::digest::{Digest, FixedOutput};
 use serde::{Deserialize, Serialize};
-use strum::{Display, EnumDiscriminants};
 
 use crate::error::StmAggregateSignatureError;
 use crate::merkle_tree::MerkleBatchPath;
 use crate::{AggregateVerificationKey, Parameters, SingleSignatureWithRegisteredParty};
 
 use super::ConcatenationProof;
 
-/// A STM aggregate signature.
-#[derive(Debug, Clone, Serialize, Deserialize, EnumDiscriminants)]
-#[strum(serialize_all = "PascalCase")]
-#[strum_discriminants(derive(Serialize, Hash, Display))] // TODO: replace strum with custom implementation
-#[strum_discriminants(name(AggregateSignatureType))]
+/// The type of STM aggregate signature.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub enum AggregateSignatureType {
+    /// Concatenation proof system.
+    Concatenation,
+}
+
+impl AggregateSignatureType {
+    /// The prefix byte used in the byte representation of the aggregate signature type.
+    ///
+    /// IMPORTANT: This value is used in serialization/deserialization. Changing it will break compatibility.
+    pub fn to_bytes_encoding_prefix(&self) -> u8 {
+        match self {
+            AggregateSignatureType::Concatenation => 0,
+        }
+    }
+
+    /// Create an aggregate signature type from a prefix byte.
+    ///
+    /// IMPORTANT: This value is used in serialization/deserialization. Changing it will break compatibility.
+    pub fn from_bytes_encoding_prefix(byte: u8) -> Option<Self> {
+        match byte {
+            0 => Some(AggregateSignatureType::Concatenation),
+            _ => None,
+        }
+    }
+}
+
+impl<D: Clone + Digest + FixedOutput + Send + Sync> From<&AggregateSignature<D>>
+    for AggregateSignatureType
+{
+    fn from(aggr_sig: &AggregateSignature<D>) -> Self {
+        match aggr_sig {
+            AggregateSignature::Concatenation(_) => AggregateSignatureType::Concatenation,
+        }
+    }
+}
+
+impl Display for AggregateSignatureType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            AggregateSignatureType::Concatenation => write!(f, "Concatenation"),
+        }
+    }
+}
+
+/// An STM aggregate signature.
+#[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(bound(
     serialize = "MerkleBatchPath<D>: Serialize",
     deserialize = "MerkleBatchPath<D>: Deserialize<'de>"
@@ -75,18 +118,36 @@ impl<D: Clone + Digest + FixedOutput + Send + Sync> AggregateSignature<D> {
 
     /// Convert an aggregate signature to bytes
     pub fn to_bytes(&self) -> Vec<u8> {
-        match self {
-            AggregateSignature::Concatenation(stm_aggr_sig) => stm_aggr_sig.to_bytes(),
-        }
+        let mut aggregate_signature_bytes = Vec::new();
+        let aggregate_signature_type: AggregateSignatureType = self.into();
+        aggregate_signature_bytes
+            .extend_from_slice(&[aggregate_signature_type.to_bytes_encoding_prefix()]);
+
+        let mut proof_bytes = match self {
+            AggregateSignature::Concatenation(concatenation_proof) => {
+                concatenation_proof.to_bytes()
+            }
+        };
+        aggregate_signature_bytes.append(&mut proof_bytes);
+
+        aggregate_signature_bytes
     }
 
     /// Extract an aggregate signature from a byte slice.
     pub fn from_bytes(bytes: &[u8]) -> Result<Self, StmAggregateSignatureError<D>> {
-        // TODO: Move first byte of concatenation proof here to identify the proof system
-        Ok(Self::Concatenation(ConcatenationProof::from_bytes(bytes)?))
+        let proof_type_byte = bytes.get(0).ok_or(StmAggregateSignatureError::SerializationError)?;
+        let proof_bytes = &bytes[1..];
+        let proof_type = AggregateSignatureType::from_bytes_encoding_prefix(*proof_type_byte)
+            .ok_or(StmAggregateSignatureError::SerializationError)?;
+        match proof_type {
+            AggregateSignatureType::Concatenation => Ok(AggregateSignature::Concatenation(
+                ConcatenationProof::from_bytes(proof_bytes)?,
+            )),
+        }
     }
 
     /// Extract the list of signatures.
+    // TODO: transfer this function to the concatenation proof ? Some proofs might not fully carry this information
     pub fn signatures(&self) -> Vec<SingleSignatureWithRegisteredParty> {
         match self {
             AggregateSignature::Concatenation(stm_aggr_sig) => stm_aggr_sig.signatures.clone(),
@@ -112,3 +173,24 @@ impl<D: Clone + Digest + FixedOutput + Send + Sync> AggregateSignature<D> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    mod aggregate_signature_type {
+        use super::*;
+
+        #[test]
+        fn golden_bytes_encoding_prefix() {
+            assert_eq!(
+                0u8,
+                AggregateSignatureType::Concatenation.to_bytes_encoding_prefix()
+            );
+            assert_eq!(
+                AggregateSignatureType::from_bytes_encoding_prefix(0u8),
+                Some(AggregateSignatureType::Concatenation)
+            );
+        }
+    }
+}