Skip to content

Commit

Permalink
Refine key load from Web3Signer messages
Browse files Browse the repository at this point in the history
  • Loading branch information
povi committed Mar 28, 2024
1 parent 8651b93 commit badb5b3
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 53 deletions.
58 changes: 21 additions & 37 deletions signer/src/signer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use futures::{
try_join,
};
use itertools::Itertools as _;
use log::info;
use log::{info, warn};
use prometheus_metrics::Metrics;
use rayon::iter::{IntoParallelIterator as _, ParallelIterator as _};
use reqwest::{Client, Url};
Expand Down Expand Up @@ -137,10 +137,26 @@ impl Signer {

pub async fn load_keys_from_web3signer(&mut self) -> Result<()> {
for (url, remote_keys) in self.web3signer.load_public_keys().await {
for public_key in remote_keys {
self.sign_methods
.entry(public_key)
.or_insert_with(|| SignMethod::Web3Signer(url.clone()));
match remote_keys {
Some(keys) => {
info!(
"fetched {} validator key(s) from Web3Signer at {}",
keys.len(),
url,
);

for public_key in keys {
self.sign_methods
.entry(public_key)
.or_insert_with(|| SignMethod::Web3Signer(url.clone()));
}
}
None => {
warn!(
"Web3Signer at {} did not return any validator keys. It will retry to fetch keys again in the next epoch.",
url,
);
}
}
}

Expand All @@ -152,38 +168,6 @@ impl Signer {
self.sign_methods.is_empty()
}

pub async fn refresh_keys_from_web3signer(&mut self) {
for (url, remote_keys) in self.web3signer.load_public_keys().await {
self.sign_methods
.retain(|public_key, sign_method| match sign_method {
SignMethod::SecretKey(_, _) => true,
SignMethod::Web3Signer(api_url) => {
let retain = url != api_url || remote_keys.contains(public_key);

if !retain {
info!(
"Validator credentials with public key {:?} were removed from Web3Signer at {}",
public_key, url,
);
}

retain
}
});

for public_key in remote_keys {
self.sign_methods.entry(public_key).or_insert_with(|| {
info!(
"Validator credentials with public key {:?} were added to Web3Signer at {}",
public_key, url,
);

SignMethod::Web3Signer(url.clone())
});
}
}
}

pub async fn sign<'block, P: Preset>(
&self,
message: SigningMessage<'block, P>,
Expand Down
24 changes: 15 additions & 9 deletions signer/src/web3signer/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ impl Web3Signer {
&self.client
}

pub async fn load_public_keys(&mut self) -> HashMap<&Url, HashSet<PublicKeyBytes>> {
pub async fn load_public_keys(&mut self) -> HashMap<&Url, Option<HashSet<PublicKeyBytes>>> {
let _timer = self
.metrics
.as_ref()
Expand All @@ -60,12 +60,16 @@ impl Web3Signer {

match self.load_public_keys_from_url(url).await {
Ok(mut remote_keys) => {
if !self.config.public_keys.is_empty() {
remote_keys.retain(|pubkey| self.config.public_keys.contains(pubkey));
if remote_keys.is_empty() {
keys.insert(url, None);
} else {
if !self.config.public_keys.is_empty() {
remote_keys.retain(|pubkey| self.config.public_keys.contains(pubkey));
}

keys.insert(url, Some(remote_keys));
self.keys_loaded.insert(url.clone());
}

keys.insert(url, remote_keys);
self.keys_loaded.insert(url.clone());
}
Err(error) => warn!("failed to load Web3Signer keys from {url}: {error:?}"),
}
Expand Down Expand Up @@ -158,7 +162,8 @@ mod tests {
let mut web3signer = Web3Signer::new(Client::new(), config, None);

let response = web3signer.load_public_keys().await;
let expected = HashMap::from([(&url, HashSet::from([SAMPLE_PUBKEY, SAMPLE_PUBKEY_2]))]);
let expected =
HashMap::from([(&url, Some(HashSet::from([SAMPLE_PUBKEY, SAMPLE_PUBKEY_2])))]);

assert_eq!(response, expected);

Expand Down Expand Up @@ -190,7 +195,8 @@ mod tests {
let mut web3signer = Web3Signer::new(Client::new(), config, None);

let response = web3signer.load_public_keys().await;
let expected = HashMap::from([(&url, HashSet::from([SAMPLE_PUBKEY, SAMPLE_PUBKEY_2]))]);
let expected =
HashMap::from([(&url, Some(HashSet::from([SAMPLE_PUBKEY, SAMPLE_PUBKEY_2])))]);

assert_eq!(response, expected);

Expand Down Expand Up @@ -220,7 +226,7 @@ mod tests {
let mut web3signer = Web3Signer::new(Client::new(), config, None);

let response = web3signer.load_public_keys().await;
let expected = HashMap::from([(&url, HashSet::from([SAMPLE_PUBKEY_2]))]);
let expected = HashMap::from([(&url, Some(HashSet::from([SAMPLE_PUBKEY_2])))]);

assert_eq!(response, expected);

Expand Down
8 changes: 1 addition & 7 deletions validator/src/validator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2960,13 +2960,7 @@ impl<P: Preset, W: Wait + Sync> Validator<P, W> {
fn refresh_signer_keys(&self) {
let signer_arc = self.signer.clone_arc();

tokio::spawn(async move {
signer_arc
.write()
.await
.refresh_keys_from_web3signer()
.await
});
tokio::spawn(async move { signer_arc.write().await.load_keys_from_web3signer().await });
}

fn get_execution_payload_header(
Expand Down

0 comments on commit badb5b3

Please sign in to comment.