From 99837162c136d69494792ef2a2c26dc53292516e Mon Sep 17 00:00:00 2001 From: fbrv Date: Tue, 23 Jul 2024 15:43:54 +0100 Subject: [PATCH] use serde for serialize directly --- src/config.rs | 74 +++++++++++++++++++++++++++++++++++++--- src/forward_service.rs | 11 +++--- src/lookahead/manager.rs | 17 ++++----- 3 files changed, 83 insertions(+), 19 deletions(-) diff --git a/src/config.rs b/src/config.rs index c242dd1..c59b8e2 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,8 +1,10 @@ -use std::{fs, path::Path}; +use std::{fs, path::Path, str::FromStr}; +use alloy::rpc::types::beacon::BlsPublicKey; use eyre::{Result, WrapErr}; use hashbrown::HashMap; -use serde::Deserialize; +use serde::{Deserialize, Deserializer}; +use url::Url; #[derive(Debug, Deserialize)] #[serde(rename_all = "kebab-case")] @@ -13,7 +15,7 @@ pub enum UrlProvider { #[derive(Debug, Deserialize)] pub struct Config { - #[serde(rename = "lookahead-providers-relays")] + #[serde(rename = "lookahead")] pub lookahead_providers_relays: Vec, #[serde(rename = "beacon-nodes")] pub beacon_nodes: Vec, @@ -25,15 +27,77 @@ pub struct Lookahead { pub chain_id: u16, #[serde(rename = "relays")] pub relays: Vec, - #[serde(rename = "registry")] - pub registry: Option>, + #[serde(rename = "registry", deserialize_with = "deserialize_registry")] + pub registry: HashMap, #[serde(rename = "url-provider")] pub url_provider: UrlProvider, } +fn deserialize_registry<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + let temp_registry: HashMap = HashMap::deserialize(deserializer)?; + let mut registry: HashMap = HashMap::new(); + + for (key, value) in temp_registry { + match BlsPublicKey::from_str(key.as_str()) { + Ok(bls_key) => { + registry.insert(bls_key, Url::from_str(&value).unwrap()); + } + Err(_) => { + return Err(serde::de::Error::custom(format!("Failed to convert key: {}", key))); + } + } + } + + Ok(registry) +} + impl Config { pub fn from_file(filepath: &Path) -> Result { let toml_str = fs::read_to_string(filepath)?; toml::from_str(&toml_str).wrap_err("could not parse configuration file") } } + +#[cfg(test)] +mod tests { + use std::str::FromStr; + + use super::*; + + #[test] + fn test_deserialize_config() { + let data = r#" + beacon-nodes = ["node1", "node2"] + [[lookahead]] + chain-id = 1 + url-provider = "lookahead" + relays = ["relay1", "relay2"] + [lookahead.registry] + "0x8248efd1f054fcccd090879c4011ed91ee9f9d0db5ad125ae1af74fdd33de809ddc882400d99b5184ca065d4570df8cc" = "localhost:21009" + "#; + + let expected_registry = { + let mut registry = HashMap::new(); + registry.insert(BlsPublicKey::from_str("0x8248efd1f054fcccd090879c4011ed91ee9f9d0db5ad125ae1af74fdd33de809ddc882400d99b5184ca065d4570df8cc").unwrap(), Url::from_str("localhost:21009").unwrap()); + registry + }; + + let expected_lookahead = Lookahead { + chain_id: 1, + relays: vec!["relay1".to_string(), "relay2".to_string()], + registry: expected_registry, + url_provider: UrlProvider::Lookahead, + }; + + let _expected_config = Config { + lookahead_providers_relays: vec![expected_lookahead], + beacon_nodes: vec!["node1".to_string(), "node2".to_string()], + }; + + let config: Config = toml::from_str(data).unwrap(); + assert!(matches!(config, _expected_config)); + } +} diff --git a/src/forward_service.rs b/src/forward_service.rs index 58af53c..8396da9 100644 --- a/src/forward_service.rs +++ b/src/forward_service.rs @@ -22,6 +22,7 @@ use reqwest_tracing::{ use tokio::task::JoinHandle; use tower_http::trace::TraceLayer; use tracing::Span; +use url::Url; use crate::lookahead::LookaheadManager; @@ -110,7 +111,7 @@ async fn scan_id_forward_request( ) -> Result { if let Some(manager) = state.managers.get(&chain_id) { match manager.get_url() { - Ok(url) => match inner_forward_request(&state.client, &url, body, headers).await { + Ok(url) => match inner_forward_request(&state.client, url, body, headers).await { Ok(res) => Ok(res), Err(_) => Err(( StatusCode::INTERNAL_SERVER_ERROR, @@ -133,7 +134,7 @@ async fn forward_request(State(_state): State>) -> impl IntoRes async fn inner_forward_request( client: &ClientWithMiddleware, - to_addr: &str, + to_addr: Url, bytes: Bytes, headers: HeaderMap, ) -> Result { @@ -146,6 +147,7 @@ async fn inner_forward_request( mod test { use std::{ default::Default, + str::FromStr, sync::{Arc, Mutex}, time::Duration, }; @@ -164,6 +166,7 @@ mod test { use hashbrown::HashMap; use http::{HeaderValue, StatusCode}; use tokio::task::JoinHandle; + use url::Url; use crate::{ forward_service::{router, SharedState}, @@ -308,7 +311,7 @@ mod test { let map = Arc::new(DashMap::new()); let signature: BlsPublicKey = BlsPublicKey::from([42u8; BLS_PUBLIC_KEY_BYTES_LEN]); let mut url_mapping = HashMap::new(); - url_mapping.insert(signature, "http://localhost:12006".into()); + url_mapping.insert(signature, Url::from_str("http://localhost:12006").unwrap()); map.insert(0, LookaheadEntry { url: "".into(), election: SignedPreconferElection { @@ -347,7 +350,7 @@ mod test { let signature: BlsPublicKey = BlsPublicKey::from([42u8; BLS_PUBLIC_KEY_BYTES_LEN]); let map = Arc::new(DashMap::new()); let mut provider = HashMap::new(); - provider.insert(signature, "http:://not-a-valid-http".into()); + provider.insert(signature, Url::from_str("http://localhost:12010/1").unwrap()); map.insert(0, LookaheadEntry { url: "".into(), ..Default::default() }); let manager = LookaheadManager::new( Lookahead { map }, diff --git a/src/lookahead/manager.rs b/src/lookahead/manager.rs index 8a3481f..0ff0138 100644 --- a/src/lookahead/manager.rs +++ b/src/lookahead/manager.rs @@ -5,6 +5,7 @@ use dashmap::DashMap; use eyre::{bail, ContextCompat, Result}; use hashbrown::HashMap; use tokio::sync::broadcast; +use url::Url; use super::{ provider::LookaheadProvider, Lookahead, LookaheadEntry, LookaheadProviderOptions, @@ -22,7 +23,7 @@ enum LookaheadProviderManager { #[derive(Debug, Clone)] pub enum UrlProvider { LookaheadEntry, - UrlMap(HashMap), + UrlMap(HashMap), } #[derive(Debug)] @@ -77,11 +78,13 @@ impl LookaheadManager { self.lookahead.get_next_elected_preconfer() } - pub fn get_url(&self) -> Result { + pub fn get_url(&self) -> Result { match self.get_next_elected_preconfer() { None => bail!("no lookahead provider found"), Some(entry) => match &self.url_provider { - UrlProvider::LookaheadEntry => Ok(entry.url), + UrlProvider::LookaheadEntry => { + Ok(Url::from_str(&entry.url).expect("not a valid url")) + } UrlProvider::UrlMap(m) => { let pub_key = entry.election.preconfer_pubkey(); m.get(&pub_key) @@ -113,13 +116,7 @@ pub fn lookahead_managers_from_config( .build_relay_provider(); let url_provider = match r_c.url_provider { crate::config::UrlProvider::Lookahead => UrlProvider::LookaheadEntry, - crate::config::UrlProvider::UrlMapping => UrlProvider::UrlMap( - r_c.registry - .expect("url mapping mode is enabled, registry must be set.") - .into_iter() - .map(|(k, v)| (BlsPublicKey::from_str(&k).expect("invalid bls public key"), v)) - .collect(), - ), + crate::config::UrlProvider::UrlMapping => UrlProvider::UrlMap(r_c.registry), }; map.insert(r_c.chain_id, LookaheadManager::new(lookahead, provider, url_provider)); }