Skip to content

Commit 9f0fdf3

Browse files
AlenarAlenar
committed
Signer: split the runtime run methods in severals parts
Signed-off-by: DJO <[email protected]>
1 parent 8e7365c commit 9f0fdf3

File tree

2 files changed

+95
-77
lines changed

2 files changed

+95
-77
lines changed

‎mithril-signer/src/main.rs‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,8 @@
11
use clap::Parser;
22
use slog::{o, Drain, Level, Logger};
3-
use slog_scope::{debug, error, info};
3+
use slog_scope::debug;
44
use std::env;
55
use std::sync::Arc;
6-
use std::time::Duration;
7-
use tokio::time::sleep;
86

97
use mithril_signer::{CertificateHandlerHTTPClient, Config, MithrilSingleSigner, Runtime};
108

‎mithril-signer/src/runtime.rs‎

Lines changed: 94 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,10 @@ use std::time::Duration;
44
use thiserror::Error;
55
use tokio::time::sleep;
66

7-
use mithril_common::crypto_helper::key_encode_hex;
8-
use mithril_common::entities::{self, Beacon, SignerWithStake};
7+
use crate::certificate_handler::CertificateHandlerError;
8+
use crate::single_signer::SingleSignerError;
9+
use mithril_common::crypto_helper::{key_encode_hex, Bytes};
10+
use mithril_common::entities::{self, Beacon, CertificatePending, SignerWithStake};
911
use mithril_common::fake_data;
1012

1113
use super::certificate_handler::CertificateHandler;
@@ -17,12 +19,12 @@ pub struct Runtime {
1719
current_beacon: Option<Beacon>,
1820
}
1921

20-
#[derive(Error, Debug, PartialEq)]
22+
#[derive(Error, Debug)]
2123
pub enum RuntimeError {
2224
#[error("single signatures computation failed: `{0}`")]
23-
SingleSignaturesComputeFailed(String),
25+
SingleSignaturesComputeFailed(#[from] SingleSignerError),
2426
#[error("could not retrieve pending certificate: `{0}`")]
25-
RetrievePendingCertificateFailed(String),
27+
RetrievePendingCertificateFailed(#[from] CertificateHandlerError),
2628
#[error("could not retrieve protocol initializer")]
2729
RetrieveProtocolInitializerFailed(),
2830
#[error("register signer failed: `{0}`")]
@@ -58,74 +60,89 @@ impl Runtime {
5860
if let Some(pending_certificate) = self
5961
.certificate_handler
6062
.retrieve_pending_certificate()
61-
.await
62-
.map_err(|e| RuntimeError::RetrievePendingCertificateFailed(e.to_string()))?
63+
.await?
6364
{
6465
let message = fake_data::digest(&pending_certificate.beacon);
65-
let must_register_signature = match &self.current_beacon {
66-
None => true,
67-
Some(beacon) => beacon != &pending_certificate.beacon,
68-
};
69-
70-
let must_register_signer = !self.single_signer.get_is_registered();
71-
if must_register_signer {
72-
if let Some(protocol_initializer) = self.single_signer.get_protocol_initializer() {
73-
let verification_key = protocol_initializer.verification_key();
74-
let verification_key =
75-
key_encode_hex(verification_key).map_err(RuntimeError::Codec)?;
76-
let signer =
77-
entities::Signer::new(self.single_signer.get_party_id(), verification_key);
78-
self.certificate_handler
79-
.register_signer(&signer)
80-
.await
81-
.map_err(|e| RuntimeError::RegisterSignerFailed(e.to_string()))?;
82-
self.single_signer
83-
.update_is_registered(true)
84-
.map_err(|e| RuntimeError::RegisterSignerFailed(e.to_string()))?;
85-
}
86-
}
8766

88-
if must_register_signature {
89-
let stake_distribution = fake_data::signers_with_stakes(5);
90-
let verification_keys = pending_certificate
91-
.signers
92-
.iter()
93-
.map(|signer| (signer.party_id, signer.verification_key.as_str()))
94-
.collect::<HashMap<u64, &str>>();
95-
let stake_distribution_extended = stake_distribution
96-
.into_iter()
97-
.map(|signer| {
98-
let verification_key = match verification_keys.get(&signer.party_id) {
99-
Some(verification_key_found) => *verification_key_found,
100-
None => "",
101-
};
102-
SignerWithStake::new(
103-
signer.party_id,
104-
verification_key.to_string(),
105-
signer.stake,
106-
)
107-
})
108-
.collect::<Vec<SignerWithStake>>();
109-
let signatures = self
110-
.single_signer
111-
.compute_single_signatures(
112-
message,
113-
stake_distribution_extended,
114-
&pending_certificate.protocol_parameters,
115-
)
116-
.map_err(|e| RuntimeError::SingleSignaturesComputeFailed(e.to_string()))?;
117-
if !signatures.is_empty() {
118-
let _ = self
119-
.certificate_handler
120-
.register_signatures(&signatures)
121-
.await;
122-
}
123-
self.current_beacon = Some(pending_certificate.beacon);
67+
self.register_to_aggregator_if_needed().await?;
68+
69+
if self.should_register_signature(&pending_certificate.beacon) {
70+
self.register_signature(message, pending_certificate)
71+
.await?;
12472
}
12573
}
12674

12775
Ok(())
12876
}
77+
78+
fn should_register_signature(&self, new_beacon: &Beacon) -> bool {
79+
match &self.current_beacon {
80+
None => true,
81+
Some(beacon) => beacon != new_beacon,
82+
}
83+
}
84+
85+
async fn register_to_aggregator_if_needed(&mut self) -> Result<(), RuntimeError> {
86+
let must_register_to_aggregator = !self.single_signer.get_is_registered();
87+
if !must_register_to_aggregator {
88+
return Ok(());
89+
}
90+
91+
if let Some(protocol_initializer) = self.single_signer.get_protocol_initializer() {
92+
let verification_key = protocol_initializer.verification_key();
93+
let verification_key = key_encode_hex(verification_key).map_err(RuntimeError::Codec)?;
94+
let signer = entities::Signer::new(self.single_signer.get_party_id(), verification_key);
95+
self.certificate_handler
96+
.register_signer(&signer)
97+
.await
98+
.map_err(|e| RuntimeError::RegisterSignerFailed(e.to_string()))?;
99+
self.single_signer
100+
.update_is_registered(true)
101+
.map_err(|e| RuntimeError::RegisterSignerFailed(e.to_string()))?;
102+
}
103+
104+
Ok(())
105+
}
106+
107+
async fn register_signature(
108+
&mut self,
109+
message: Bytes,
110+
pending_certificate: CertificatePending,
111+
) -> Result<(), RuntimeError> {
112+
let verification_keys = pending_certificate
113+
.signers
114+
.iter()
115+
.map(|signer| (signer.party_id, signer.verification_key.as_str()))
116+
.collect::<HashMap<u64, &str>>();
117+
118+
let stake_distribution = fake_data::signers_with_stakes(5);
119+
let stake_distribution_extended = stake_distribution
120+
.into_iter()
121+
.map(|signer| {
122+
let verification_key = match verification_keys.get(&signer.party_id) {
123+
Some(verification_key_found) => *verification_key_found,
124+
None => "",
125+
};
126+
SignerWithStake::new(signer.party_id, verification_key.to_string(), signer.stake)
127+
})
128+
.collect::<Vec<SignerWithStake>>();
129+
130+
let signatures = self.single_signer.compute_single_signatures(
131+
message,
132+
stake_distribution_extended,
133+
&pending_certificate.protocol_parameters,
134+
)?;
135+
136+
if !signatures.is_empty() {
137+
let _ = self
138+
.certificate_handler
139+
.register_signatures(&signatures)
140+
.await;
141+
}
142+
self.current_beacon = Some(pending_certificate.beacon);
143+
144+
Ok(())
145+
}
129146
}
130147

131148
#[cfg(test)]
@@ -190,9 +207,10 @@ mod tests {
190207
);
191208
assert_eq!(
192209
RuntimeError::RetrievePendingCertificateFailed(
193-
CertificateHandlerError::RemoteServerTechnical("An Error".to_string()).to_string()
194-
),
195-
signer.run().await.unwrap_err()
210+
CertificateHandlerError::RemoteServerTechnical("An Error".to_string())
211+
)
212+
.to_string(),
213+
signer.run().await.unwrap_err().to_string()
196214
);
197215
}
198216

@@ -355,9 +373,10 @@ mod tests {
355373
);
356374
assert_eq!(
357375
RuntimeError::SingleSignaturesComputeFailed(
358-
SingleSignerError::UnregisteredVerificationKey().to_string()
359-
),
360-
signer.run().await.unwrap_err()
376+
SingleSignerError::UnregisteredVerificationKey()
377+
)
378+
.to_string(),
379+
signer.run().await.unwrap_err().to_string()
361380
);
362381
}
363382

@@ -398,10 +417,11 @@ mod tests {
398417
);
399418
assert_eq!(
400419
RuntimeError::RegisterSignerFailed(
401-
CertificateHandlerError::RemoteServerLogical("an error occurred".to_string(),)
420+
CertificateHandlerError::RemoteServerLogical("an error occurred".to_string())
402421
.to_string()
403-
),
404-
signer.run().await.unwrap_err()
422+
)
423+
.to_string(),
424+
signer.run().await.unwrap_err().to_string()
405425
);
406426
}
407427
}

0 commit comments

Comments
 (0)