mithril_aggregator/database/repository/
signer_registration_store.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use anyhow::Context;
5use async_trait::async_trait;
6
7use mithril_common::entities::{Epoch, PartyId, Signer, SignerWithStake};
8use mithril_common::StdResult;
9use mithril_persistence::sqlite::{ConnectionExtensions, SqliteConnection};
10
11use crate::database::query::{
12    DeleteSignerRegistrationRecordQuery, GetSignerRegistrationRecordQuery,
13    InsertOrReplaceSignerRegistrationRecordQuery,
14};
15use crate::database::record::SignerRegistrationRecord;
16use crate::services::EpochPruningTask;
17use crate::VerificationKeyStorer;
18
19/// Service to deal with signer_registration (read & write).
20pub struct SignerRegistrationStore {
21    connection: Arc<SqliteConnection>,
22
23    /// Number of epochs before previous records will be deleted at the next registration round
24    /// opening
25    verification_key_epoch_retention_limit: Option<u64>,
26}
27
28impl SignerRegistrationStore {
29    /// Create a new [SignerRegistrationStore] service
30    pub fn new(
31        connection: Arc<SqliteConnection>,
32        verification_key_epoch_retention_limit: Option<u64>,
33    ) -> Self {
34        Self {
35            connection,
36            verification_key_epoch_retention_limit,
37        }
38    }
39}
40
41#[async_trait]
42impl VerificationKeyStorer for SignerRegistrationStore {
43    async fn save_verification_key(
44        &self,
45        epoch: Epoch,
46        signer: SignerWithStake,
47    ) -> StdResult<Option<SignerWithStake>> {
48        let existing_record = self
49            .connection
50            .fetch_first(GetSignerRegistrationRecordQuery::by_signer_id_and_epoch(
51                signer.party_id.to_owned(),
52                epoch,
53            )?)
54            .with_context(|| {
55                format!(
56                    "Get signer registration record failure with signer_id: '{}', epoch: '{}'",
57                    signer.party_id, epoch
58                )
59            })?;
60
61        let _updated_record = self
62            .connection
63            .fetch_first(InsertOrReplaceSignerRegistrationRecordQuery::one(
64                SignerRegistrationRecord::from_signer_with_stake(signer, epoch),
65            ))
66            .with_context(|| format!("persist verification key failure, epoch: {epoch}"))?;
67
68        match existing_record {
69            None => Ok(None),
70            Some(previous_record) => Ok(Some(previous_record.into())),
71        }
72    }
73
74    async fn get_verification_keys(
75        &self,
76        epoch: Epoch,
77    ) -> StdResult<Option<HashMap<PartyId, Signer>>> {
78        let cursor = self
79            .connection
80            .fetch(GetSignerRegistrationRecordQuery::by_epoch(epoch)?)
81            .with_context(|| format!("get verification key failure, epoch: {epoch}"))?;
82
83        let signer_with_stakes: HashMap<PartyId, Signer> =
84            HashMap::from_iter(cursor.map(|record| (record.signer_id.to_owned(), record.into())));
85
86        match signer_with_stakes.is_empty() {
87            true => Ok(None),
88            false => Ok(Some(signer_with_stakes)),
89        }
90    }
91
92    async fn get_signers(&self, epoch: Epoch) -> StdResult<Option<Vec<SignerWithStake>>> {
93        let cursor = self
94            .connection
95            .fetch(GetSignerRegistrationRecordQuery::by_epoch(epoch)?)
96            .with_context(|| format!("get verification key failure, epoch: {epoch}"))?;
97
98        let signer_with_stakes: Vec<SignerWithStake> = cursor.map(|record| record.into()).collect();
99
100        match signer_with_stakes.is_empty() {
101            true => Ok(None),
102            false => Ok(Some(signer_with_stakes)),
103        }
104    }
105
106    async fn prune_verification_keys(&self, max_epoch_to_prune: Epoch) -> StdResult<()> {
107        self.connection
108            .apply(DeleteSignerRegistrationRecordQuery::below_epoch_threshold(
109                max_epoch_to_prune,
110            ))?;
111
112        Ok(())
113    }
114}
115
116#[async_trait]
117impl EpochPruningTask for SignerRegistrationStore {
118    fn pruned_data(&self) -> &'static str {
119        "Signer registration"
120    }
121
122    async fn prune(&self, epoch: Epoch) -> StdResult<()> {
123        let registration_epoch = epoch.offset_to_recording_epoch();
124
125        if let Some(retention_limit) = self.verification_key_epoch_retention_limit {
126            self.prune_verification_keys(registration_epoch - retention_limit)
127                .await
128                .with_context(|| {
129                    format!(
130                        "VerificationKeyStorer can not prune verification keys below epoch: '{}'",
131                        registration_epoch - retention_limit
132                    )
133                })?;
134        }
135
136        Ok(())
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use crate::database::test_helper::{insert_signer_registrations, main_db_connection};
143
144    use mithril_common::entities::{Epoch, PartyId, Signer, SignerWithStake};
145    use mithril_common::test_utils::fake_keys;
146    use std::collections::HashMap;
147    use std::sync::Arc;
148
149    use crate::VerificationKeyStorer;
150
151    use super::*;
152
153    /// Build simple fake signers with stakes.
154    /// It could be done by `fake_data::signers_with_stakes` which produce verification keys dynamically
155    /// but take longer.
156    fn build_fake_signers_with_stakes(nb: u64) -> Vec<SignerWithStake> {
157        let verification_keys = fake_keys::signer_verification_key();
158        let nb_keys = verification_keys.len() as u64;
159        (1..=nb)
160            .map(|party_idx| SignerWithStake {
161                party_id: format!("party_id:{party_idx}"),
162                verification_key: verification_keys[(party_idx % nb_keys) as usize]
163                    .try_into()
164                    .unwrap(),
165                verification_key_signature: None,
166                operational_certificate: None,
167                kes_period: None,
168                stake: 10,
169            })
170            .collect()
171    }
172
173    fn build_signers(
174        nb_epoch: u64,
175        signers_per_epoch: usize,
176    ) -> HashMap<Epoch, Vec<SignerWithStake>> {
177        (1..=nb_epoch)
178            .map(|epoch| {
179                (
180                    Epoch(epoch),
181                    build_fake_signers_with_stakes(signers_per_epoch as u64),
182                )
183            })
184            .collect()
185    }
186
187    fn insert_golden_signer_registration(connection: &SqliteConnection) {
188        connection
189            .execute(
190                r#"
191                insert into signer_registration
192                values(
193                    'pool1t9uuagsat8hlr0n0ga4wzge0jxlyjuhl6mugrm8atc285vzkf2e',
194                    292,
195                    '7b22766b223a5b3132382c3134322c31352c37322c35342c37332c32392c3135372c39302c3134392c33342c3235352c35312c31382c34342c33322c36302c34362c3130302c31342c3136342c39362c3138362c31382c32372c3231312c3130322c3130362c35352c332c3137302c3234302c3131342c3134372c3134362c3234382c31352c32312c3232392c3133322c3234362c3230322c3136322c34312c3135312c3138362c3136332c3232302c31342c3231372c3235352c3234352c35352c3231362c3235322c342c3137302c31362c3137382c3230392c3134392c32302c3230352c39322c3232312c38302c32392c3139302c3131372c3138382c3132382c3234372c3133312c37382c3138372c3232332c3231382c3131362c3235352c34332c3130392c3132362c3233302c3130382c33372c3131342c332c3138362c3136352c33322c3133312c3139332c3139302c34342c3134362c3234315d2c22706f70223a5b3134322c3133352c38352c342c3134362c32342c37382c34332c36332c3233382c3235312c37382c3138312c37342c37302c362c39342c3138372c3137382c3133332c3135352c3233342c3235352c3134352c3139372c3137302c3135352c3132392c3234332c3137332c31322c31392c36382c3132392c3131342c36392c3231312c33372c3233322c3139332c3130372c3233332c32392c3130332c3232382c34392c36392c38362c3137322c35352c39342c3132332c372c36322c3135382c33352c31372c3131332c38312c3136312c34342c3234392c35332c36362c39332c37302c3136392c3133372c3135372c3233342c3234372c3232332c37312c3135302c3231362c3130322c3139302c3137362c3135322c34362c3134332c3233302c31322c3138382c33312c3234362c3137312c3130352c3230392c3133382c35352c32382c3134312c36302c3132312c3132305d7d',
196                    '7b227369676d61223a7b227369676d61223a7b227369676d61223a7b227369676d61223a7b227369676d61223a7b227369676d61223a5b35322c3230342c36342c3136342c3134382c36302c342c3133372c3132312c312c3130302c38352c3133382c3235302c33382c37352c39302c3133352c31392c35322c3135312c3130302c35392c372c33332c3131322c31332c3138342c3138342c32332c35352c37392c3134332c3231382c3231352c3136362c3137392c36312c3234312c3136342c38342c36302c3234312c3234302c3134362c33332c37342c3136372c37342c3230332c37382c34372c3231362c3132352c37382c31302c38312c3233322c3134332c37302c3234362c382c3130352c345d2c226c68735f706b223a5b3130342c3131332c39302c32332c3231392c3235342c382c3137352c3136372c3133312c3136382c3131322c3130392c3137362c31342c35372c38362c3139372c34392c35312c3136382c3131352c3138372c3137382c392c3233322c38362c3139352c3130362c3134322c3232372c3139365d2c227268735f706b223a5b38372c3136392c39392c3133382c3130352c3133312c3133322c3132362c3139382c3138352c3137302c3132352c39372c3137372c38342c3231302c3137362c3134322c3133382c32312c38362c3133312c3135382c3132332c36332c3131322c39392c3133322c3134352c3231322c382c3233315d7d2c226c68735f706b223a5b34302c38382c39392c372c31362c32302c37352c3132302c3234372c3233302c32372c3233392c35332c3235352c3137302c3132302c3131392c33372c3138362c3130322c35302c33362c39382c3139332c3130332c33372c352c3131362c3134322c3134382c3233322c32355d2c227268735f706b223a5b3130372c3135372c3234342c3230342c3133362c3139332c38382c3130322c3234312c3135392c39392c3233342c38342c3139322c3133302c34372c32302c3136362c322c3230302c37392c3133352c3230352c312c3235332c3233352c32382c3134372c3135352c3132322c33392c3234345d7d2c226c68735f706b223a5b3133332c3131362c3131342c36342c3132322c31332c3138362c3130342c39322c3233372c39372c38322c3232312c38322c35342c3132352c34352c3234342c3139322c39332c35362c3132362c34302c3134362c3131312c3132392c3232312c3234382c362c35342c3233372c3138355d2c227268735f706b223a5b3137312c3137302c33382c3232302c3133302c34372c3133362c3233362c37322c35332c36332c33382c32392c34362c3230352c32392c3234382c3235342c37362c32372c322c3132382c3130372c3131392c3132382c3137302c32322c3131322c3137362c3130322c3136332c37335d7d2c226c68735f706b223a5b3231312c3138392c38342c3139372c302c3231382c3134382c312c34332c36332c38342c3234322c3231392c39342c31302c3134302c3134372c3137322c38342c35392c31352c3131342c3230392c3235302c3230372c31342c3134322c33362c3135372c3230332c37382c3137355d2c227268735f706b223a5b3137302c3132392c3132332c342c3131332c3135322c3232392c3133372c392c32342c3133372c3136362c32352c3136352c34332c3132322c3132332c3230312c3234322c3231302c3137382c3234382c31342c3233302c33392c3231322c31382c33362c34382c38372c39302c3230365d7d2c226c68735f706b223a5b33332c3234322c3130372c39312c3130372c3130322c3136332c36342c33332c3231372c3233342c3138302c33382c312c3138352c3135382c3230372c3234352c3136372c3130352c3134322c33382c3233342c37362c34322c32322c372c3130342c39362c3139382c3234322c35335d2c227268735f706b223a5b39352c34302c392c3131372c3130382c3135362c3138342c3133392c39302c3138382c31352c32312c3131322c31382c3130302c3134362c3130342c352c3135362c39372c3134392c33322c34322c3234302c3134392c3138382c35322c33312c39312c39392c3131382c365d7d2c226c68735f706b223a5b3139332c3234322c37362c3230392c3134312c33372c3130312c36382c37302c392c3134312c33392c3230372c39342c3232362c33392c3136302c3131382c32332c3233302c3234342c3231302c31382c38322c3137332c3135382c3233312c3137392c3138322c31392c32322c3134365d2c227268735f706b223a5b3133322c3232392c3130382c3139392c37312c36392c3233362c36352c31382c3131372c39332c3234332c3234332c37342c36392c39382c3134302c3234392c342c33372c37372c38372c35382c31322c3132302c37332c3230332c39362c36312c3233302c39322c3132385d7d',
197                    '5b5b5b3138362c39352c3232362c3137342c3132352c3235302c31302c3232322c3130322c3234302c36352c3235352c34372c3133382c38392c3131302c31342c3131302c32322c3138322c33322c3136362c3231312c392c32302c32302c35352c35382c3232392c3132302c3235302c37315d2c312c3136352c5b3130352c35342c3234352c35362c3231352c3130362c3133392c3231322c3137342c3232332c39302c3234392c3138372c34372c3134382c35302c34302c31352c3131372c3231372c3134392c3132362c3231382c3232352c3133362c36352c3231392c3136302c3134382c39332c3232382c3235312c31392c3231332c3136382c332c3233362c38392c3132302c3135392c3139382c38302c3234342c3138302c33332c3131392c3132382c3230312c3138362c3132302c32312c3130322c36322c3232392c32382c3135352c37362c31392c3235322c3232312c3234372c3137342c3135392c365d5d2c5b3234312c32372c31332c34342c3131342c37382c3138392c3234392c3135302c3135302c35332c3134342c3233362c3135312c38382c3134302c3132382c3136322c36302c3232382c38382c3131312c392c3134342c3233322c38332c39342c3231302c3135362c3136382c33352c3234325d5d',
198                    29,
199                    9497629046,
200                    '2023-08-12T00:03:51.236860002+00:00'
201                );
202            "#,
203            )
204            .unwrap();
205    }
206
207    #[tokio::test]
208    async fn test_golden_master() {
209        let connection = main_db_connection().unwrap();
210        insert_golden_signer_registration(&connection);
211
212        let repository = SignerRegistrationStore::new(Arc::new(connection), None);
213        repository
214            .get_verification_keys(Epoch(292))
215            .await
216            .expect("Getting Golden signer registration should not fail")
217            .expect("Signer registration should exist for this epoch");
218    }
219
220    pub fn init_signer_registration_store(
221        initial_data: HashMap<Epoch, Vec<SignerWithStake>>,
222        verification_key_epoch_retention_limit: Option<u64>,
223    ) -> Arc<SignerRegistrationStore> {
224        let connection = main_db_connection().unwrap();
225
226        let initial_data = initial_data.into_iter().collect();
227        insert_signer_registrations(&connection, initial_data).unwrap();
228
229        Arc::new(SignerRegistrationStore::new(
230            Arc::new(connection),
231            verification_key_epoch_retention_limit,
232        ))
233    }
234
235    #[tokio::test]
236    pub async fn save_key_in_empty_store() {
237        let signers = build_signers(0, 0);
238        let store = init_signer_registration_store(signers, None);
239        let res = store
240            .save_verification_key(
241                Epoch(0),
242                SignerWithStake {
243                    party_id: "0".to_string(),
244                    verification_key: fake_keys::signer_verification_key()[0].try_into().unwrap(),
245                    verification_key_signature: None,
246                    operational_certificate: None,
247                    kes_period: None,
248                    stake: 10,
249                },
250            )
251            .await
252            .unwrap();
253
254        assert!(res.is_none());
255    }
256
257    #[tokio::test]
258    pub async fn update_signer_in_store() {
259        let signers = build_signers(1, 1);
260        let signers_on_epoch = signers.get(&Epoch(1)).unwrap().clone();
261        let first_signer = signers_on_epoch.first().unwrap();
262
263        let store = init_signer_registration_store(signers, None);
264        let res = store
265            .save_verification_key(
266                Epoch(1),
267                SignerWithStake {
268                    party_id: first_signer.party_id.clone(),
269                    verification_key: fake_keys::signer_verification_key()[2].try_into().unwrap(),
270                    verification_key_signature: None,
271                    operational_certificate: None,
272                    kes_period: None,
273                    stake: 10,
274                },
275            )
276            .await
277            .unwrap();
278
279        assert_eq!(
280            Some(SignerWithStake {
281                party_id: first_signer.party_id.clone(),
282                verification_key: fake_keys::signer_verification_key()[2].try_into().unwrap(),
283                verification_key_signature: None,
284                operational_certificate: None,
285                kes_period: None,
286                stake: 10,
287            }),
288            res,
289        );
290    }
291
292    #[tokio::test]
293    pub async fn get_verification_keys_for_empty_epoch() {
294        let signers = build_signers(2, 1);
295        let store = init_signer_registration_store(signers, None);
296        let res = store.get_verification_keys(Epoch(0)).await.unwrap();
297
298        assert!(res.is_none());
299    }
300
301    #[tokio::test]
302    pub async fn get_signers_for_empty_epoch() {
303        let signers = build_signers(2, 1);
304        let store = init_signer_registration_store(signers, None);
305        let res = store.get_signers(Epoch(0)).await.unwrap();
306
307        assert!(res.is_none());
308    }
309
310    #[tokio::test]
311    pub async fn get_verification_keys_for_existing_epoch() {
312        let signers = build_signers(2, 2);
313        let store = init_signer_registration_store(signers.clone(), None);
314
315        let epoch = Epoch(1);
316        let expected_signers = signers
317            .get(&epoch)
318            .unwrap()
319            .iter()
320            .map(|s| (s.party_id.clone(), Signer::from(s.clone())))
321            .collect::<HashMap<PartyId, Signer>>();
322
323        let res = store.get_verification_keys(epoch).await.unwrap().unwrap();
324
325        assert_eq!(expected_signers, res);
326    }
327
328    #[tokio::test]
329    pub async fn get_signers_for_existing_epoch() {
330        let signers = build_signers(2, 2);
331        let store = init_signer_registration_store(signers.clone(), None);
332
333        let epoch = Epoch(1);
334        let mut expected_signers = signers.get(&epoch).unwrap().clone();
335        expected_signers.sort_by(|a, b| a.party_id.cmp(&b.party_id));
336
337        let mut res = store.get_signers(epoch).await.unwrap().unwrap();
338        res.sort_by(|a, b| a.party_id.cmp(&b.party_id));
339
340        assert_eq!(expected_signers, res);
341    }
342
343    #[tokio::test]
344    pub async fn can_prune_keys_from_given_epoch_retention_limit() {
345        let signers = build_signers(6, 2);
346        let store = init_signer_registration_store(signers, None);
347
348        for epoch in 1..6 {
349            assert!(
350                store
351                    .get_verification_keys(Epoch(epoch))
352                    .await
353                    .unwrap()
354                    .is_some(),
355                "Keys should exist before pruning"
356            );
357            store
358                .prune_verification_keys(Epoch(epoch) + 1)
359                .await
360                .expect("Pruning should not fail");
361
362            let pruned_epoch_keys = store.get_verification_keys(Epoch(epoch)).await.unwrap();
363            assert_eq!(None, pruned_epoch_keys);
364        }
365    }
366
367    async fn get_epochs_in_database_until(
368        store: &SignerRegistrationStore,
369        until_epoch: Epoch,
370    ) -> Vec<Epoch> {
371        let mut epochs_in_database = vec![];
372        for epoch_number in 1..=(*until_epoch) {
373            let current_epoch = Epoch(epoch_number);
374            if store
375                .get_verification_keys(current_epoch)
376                .await
377                .unwrap()
378                .is_some()
379            {
380                epochs_in_database.push(current_epoch);
381            }
382        }
383
384        epochs_in_database
385    }
386
387    #[tokio::test]
388    async fn prune_verification_keys_epoch_older_than_given_epoch() {
389        let signers = build_signers(5, 2);
390        let store = init_signer_registration_store(signers, None);
391
392        assert_eq!(
393            vec!(Epoch(1), Epoch(2), Epoch(3), Epoch(4), Epoch(5)),
394            get_epochs_in_database_until(&store, Epoch(8)).await
395        );
396
397        store.prune_verification_keys(Epoch(4)).await.unwrap();
398
399        assert_eq!(
400            vec!(Epoch(4), Epoch(5)),
401            get_epochs_in_database_until(&store, Epoch(8)).await
402        );
403    }
404
405    #[tokio::test]
406    async fn prune_older_than_threshold() {
407        let signers = build_signers(6, 2);
408        let verification_key_epoch_retention_limit = Some(4);
409        let store = init_signer_registration_store(signers, verification_key_epoch_retention_limit);
410
411        assert_eq!(
412            vec!(Epoch(1), Epoch(2), Epoch(3), Epoch(4), Epoch(5), Epoch(6)),
413            get_epochs_in_database_until(&store, Epoch(8)).await
414        );
415
416        store.prune(Epoch(5)).await.unwrap();
417
418        assert_eq!(
419            vec!(Epoch(2), Epoch(3), Epoch(4), Epoch(5), Epoch(6)),
420            get_epochs_in_database_until(&store, Epoch(8)).await
421        );
422    }
423
424    #[tokio::test]
425    async fn without_threshold_nothing_is_pruned() {
426        let signers = build_signers(6, 2);
427        let store = init_signer_registration_store(signers, None);
428
429        assert_eq!(
430            vec!(Epoch(1), Epoch(2), Epoch(3), Epoch(4), Epoch(5), Epoch(6)),
431            get_epochs_in_database_until(&store, Epoch(8)).await
432        );
433
434        store.prune(Epoch(100)).await.unwrap();
435
436        assert_eq!(
437            vec!(Epoch(1), Epoch(2), Epoch(3), Epoch(4), Epoch(5), Epoch(6)),
438            get_epochs_in_database_until(&store, Epoch(8)).await
439        );
440    }
441}