use serde for serialize directly

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<Lookahead>,
     #[serde(rename = "beacon-nodes")]
     pub beacon_nodes: Vec<String>,
@@ -25,15 +27,77 @@ pub struct Lookahead {
     pub chain_id: u16,
     #[serde(rename = "relays")]
     pub relays: Vec<String>,
-    #[serde(rename = "registry")]
-    pub registry: Option<HashMap<String, String>>,
+    #[serde(rename = "registry", deserialize_with = "deserialize_registry")]
+    pub registry: HashMap<BlsPublicKey, Url>,
     #[serde(rename = "url-provider")]
     pub url_provider: UrlProvider,
 }
 
+fn deserialize_registry<'de, D>(deserializer: D) -> Result<HashMap<BlsPublicKey, Url>, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    let temp_registry: HashMap<String, String> = HashMap::deserialize(deserializer)?;
+    let mut registry: HashMap<BlsPublicKey, Url> = 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<Self> {
         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));
+    }
+}

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<impl IntoResponse, impl IntoResponse> {
     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<Arc<SharedState>>) -> impl IntoRes
 
 async fn inner_forward_request(
     client: &ClientWithMiddleware,
-    to_addr: &str,
+    to_addr: Url,
     bytes: Bytes,
     headers: HeaderMap,
 ) -> Result<impl IntoResponse> {
@@ -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 },

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<BlsPublicKey, String>),
+    UrlMap(HashMap<BlsPublicKey, Url>),
 }
 
 #[derive(Debug)]
@@ -77,11 +78,13 @@ impl LookaheadManager {
         self.lookahead.get_next_elected_preconfer()
     }
 
-    pub fn get_url(&self) -> Result<String> {
+    pub fn get_url(&self) -> Result<Url> {
         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));
     }