mithril_stm/protocol/key_registration/
registration_entry.rs

1use std::cmp::Ordering;
2use std::hash::Hash;
3
4use serde::{Deserialize, Serialize};
5
6use crate::{
7    Initializer, RegisterError, Stake, StmResult, VerificationKeyForConcatenation,
8    VerificationKeyProofOfPossessionForConcatenation,
9};
10
11/// Represents a signer registration entry
12#[derive(PartialEq, Eq, Clone, Debug, Copy, Serialize, Deserialize)]
13pub struct RegistrationEntry(VerificationKeyForConcatenation, Stake);
14
15impl RegistrationEntry {
16    /// Creates a new registration entry. Verifies the proof of possession before creating the
17    /// entry. Fails if the proof of possession is invalid.
18    pub fn new(
19        bls_verification_key_proof_of_possession: VerificationKeyProofOfPossessionForConcatenation,
20        stake: Stake,
21    ) -> StmResult<Self> {
22        bls_verification_key_proof_of_possession
23            .verify_proof_of_possession()
24            .map_err(|_| {
25                RegisterError::KeyInvalid(Box::new(bls_verification_key_proof_of_possession.vk))
26            })?;
27        Ok(RegistrationEntry(
28            bls_verification_key_proof_of_possession.vk,
29            stake,
30        ))
31    }
32
33    /// Gets the BLS verification key.
34    pub fn get_bls_verification_key(&self) -> VerificationKeyForConcatenation {
35        self.0
36    }
37
38    /// Gets the stake associated with the registration entry.
39    pub fn get_stake(&self) -> Stake {
40        self.1
41    }
42
43    /// Converts the registration entry to bytes.
44    /// Uses 96 bytes for the BLS verification key and 8 bytes for the stake (u64 big-endian).
45    /// The order is backward compatible with previous implementations.
46    pub(crate) fn to_bytes(self) -> Vec<u8> {
47        let mut result = [0u8; 104];
48        result[..96].copy_from_slice(&self.0.to_bytes());
49        result[96..].copy_from_slice(&self.1.to_be_bytes());
50        result.to_vec()
51    }
52
53    /// Creates a registration entry from bytes.
54    /// Expects 96 bytes for the BLS verification key and 8 bytes for the stake (u64 big-endian).
55    /// The order is backward compatible with previous implementations.
56    pub(crate) fn from_bytes(bytes: &[u8]) -> StmResult<Self> {
57        let bls_verification_key = VerificationKeyForConcatenation::from_bytes(bytes)?;
58        let mut u64_bytes = [0u8; 8];
59        u64_bytes.copy_from_slice(&bytes[96..]);
60        let stake = Stake::from_be_bytes(u64_bytes);
61        Ok(RegistrationEntry(bls_verification_key, stake))
62    }
63}
64
65impl From<Initializer> for RegistrationEntry {
66    fn from(initializer: Initializer) -> Self {
67        Self(
68            initializer.bls_verification_key_proof_of_possession.vk,
69            initializer.stake,
70        )
71    }
72}
73
74impl Hash for RegistrationEntry {
75    /// Hashes the registration entry by hashing the stake first, then the verification key.
76    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
77        self.1.hash(state);
78        self.0.hash(state);
79    }
80
81    fn hash_slice<H: std::hash::Hasher>(data: &[Self], state: &mut H)
82    where
83        Self: Sized,
84    {
85        for piece in data {
86            piece.hash(state)
87        }
88    }
89}
90
91impl PartialOrd for RegistrationEntry {
92    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
93        Some(std::cmp::Ord::cmp(self, other))
94    }
95}
96
97impl Ord for RegistrationEntry {
98    /// Compares the registration entries by comparing the stake first, then the verification key.
99    fn cmp(&self, other: &Self) -> Ordering {
100        self.1.cmp(&other.1).then(self.0.cmp(&other.0))
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use rand_chacha::ChaCha20Rng;
107    use rand_core::SeedableRng;
108    use std::cmp::Ordering;
109
110    use crate::{
111        VerificationKeyProofOfPossessionForConcatenation, signature_scheme::BlsSigningKey,
112    };
113
114    use super::*;
115
116    fn create_registration_entry(rng: &mut ChaCha20Rng, stake: Stake) -> RegistrationEntry {
117        let sk = BlsSigningKey::generate(rng);
118        let pk = VerificationKeyProofOfPossessionForConcatenation::from(&sk);
119        RegistrationEntry::new(pk, stake).unwrap()
120    }
121
122    #[test]
123    fn test_ord_different_stakes() {
124        let mut rng = ChaCha20Rng::from_seed([0u8; 32]);
125
126        let entry_low_stake = create_registration_entry(&mut rng, 100);
127        let entry_high_stake = create_registration_entry(&mut rng, 200);
128
129        assert_eq!(entry_low_stake.cmp(&entry_high_stake), Ordering::Less);
130        assert_eq!(entry_high_stake.cmp(&entry_low_stake), Ordering::Greater);
131    }
132
133    #[test]
134    fn test_ord_same_stake_different_keys() {
135        let mut rng = ChaCha20Rng::from_seed([0u8; 32]);
136
137        let entry1 = create_registration_entry(&mut rng, 100);
138        let entry2 = create_registration_entry(&mut rng, 100);
139
140        let cmp_result = entry1.cmp(&entry2);
141        assert!(cmp_result == Ordering::Less || cmp_result == Ordering::Greater);
142
143        assert_eq!(entry2.cmp(&entry1), cmp_result.reverse());
144    }
145}