feat: add rpc endpoint config

crates/chisel/src/dispatcher.rs

@@ -328,7 +328,7 @@ impl ChiselDispatcher {
                     {
                         endpoint.clone()
                     } else {
-                        RpcEndpoint::Env(arg.to_string())
+                        RpcEndpoint::Env(arg.to_string()).into()
                     };
                     let fork_url = match endpoint.resolve() {
                         Ok(fork_url) => fork_url,

crates/config/src/endpoints.rs

@@ -1,7 +1,7 @@
 //! Support for multiple RPC-endpoints
 
 use crate::resolve::{interpolate, UnresolvedEnvVarError, RE_PLACEHOLDER};
-use serde::{Deserialize, Deserializer, Serialize, Serializer};
+use serde::{ser::SerializeMap, Deserialize, Deserializer, Serialize, Serializer};
 use std::{
     collections::BTreeMap,
     fmt,
@@ -12,15 +12,28 @@ use std::{
 #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(transparent)]
 pub struct RpcEndpoints {
-    endpoints: BTreeMap<String, RpcEndpoint>,
+    endpoints: BTreeMap<String, RpcEndpointConfig>,
 }
 
 // === impl RpcEndpoints ===
 
 impl RpcEndpoints {
     /// Creates a new list of endpoints
-    pub fn new(endpoints: impl IntoIterator<Item = (impl Into<String>, RpcEndpoint)>) -> Self {
-        Self { endpoints: endpoints.into_iter().map(|(name, url)| (name.into(), url)).collect() }
+    pub fn new(
+        endpoints: impl IntoIterator<Item = (impl Into<String>, impl Into<RpcEndpointType>)>,
+    ) -> Self {
+        Self {
+            endpoints: endpoints
+                .into_iter()
+                .map(|(name, e)| match e.into() {
+                    RpcEndpointType::String(url) => (
+                        name.into(),
+                        RpcEndpointConfig { endpoint: url.into(), ..Default::default() },
+                    ),
+                    RpcEndpointType::Config(config) => (name.into(), config),
+                })
+                .collect(),
+        }
     }
 
     /// Returns `true` if this type doesn't contain any endpoints
@@ -37,13 +50,73 @@ impl RpcEndpoints {
 }
 
 impl Deref for RpcEndpoints {
-    type Target = BTreeMap<String, RpcEndpoint>;
+    type Target = BTreeMap<String, RpcEndpointConfig>;
 
     fn deref(&self) -> &Self::Target {
         &self.endpoints
     }
 }
 
+/// RPC endpoint wrapper type
+#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
+#[serde(untagged)]
+pub enum RpcEndpointType {
+    /// Raw Endpoint url string
+    String(RpcEndpoint),
+    /// Config object
+    Config(RpcEndpointConfig),
+}
+
+impl RpcEndpointType {
+    /// Returns the string variant
+    pub fn as_endpoint_string(&self) -> Option<&RpcEndpoint> {
+        match self {
+            RpcEndpointType::String(url) => Some(url),
+            RpcEndpointType::Config(_) => None,
+        }
+    }
+
+    /// Returns the config variant
+    pub fn as_endpoint_config(&self) -> Option<&RpcEndpointConfig> {
+        match self {
+            RpcEndpointType::Config(config) => Some(&config),
+            RpcEndpointType::String(_) => None,
+        }
+    }
+
+    /// Returns the url or config this type holds
+    ///
+    /// # Error
+    ///
+    /// Returns an error if the type holds a reference to an env var and the env var is not set
+    pub fn resolve(self) -> Result<String, UnresolvedEnvVarError> {
+        match self {
+            RpcEndpointType::String(url) => url.resolve(),
+            RpcEndpointType::Config(config) => config.endpoint.resolve(),
+        }
+    }
+}
+
+impl fmt::Display for RpcEndpointType {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            RpcEndpointType::String(url) => url.fmt(f),
+            RpcEndpointType::Config(config) => config.fmt(f),
+        }
+    }
+}
+
+impl TryFrom<RpcEndpointType> for String {
+    type Error = UnresolvedEnvVarError;
+
+    fn try_from(value: RpcEndpointType) -> Result<Self, Self::Error> {
+        match value {
+            RpcEndpointType::String(url) => url.resolve(),
+            RpcEndpointType::Config(config) => config.endpoint.resolve(),
+        }
+    }
+}
+
 /// Represents a single endpoint
 ///
 /// This type preserves the value as it's stored in the config. If the value is a reference to an
@@ -134,6 +207,133 @@ impl<'de> Deserialize<'de> for RpcEndpoint {
     }
 }
 
+impl Into<RpcEndpointType> for RpcEndpoint {
+    fn into(self) -> RpcEndpointType {
+        RpcEndpointType::String(self)
+    }
+}
+
+impl Into<RpcEndpointConfig> for RpcEndpoint {
+    fn into(self) -> RpcEndpointConfig {
+        RpcEndpointConfig { endpoint: self, ..Default::default() }
+    }
+}
+
+/// Rpc endpoint configuration variant
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct RpcEndpointConfig {
+    /// endpoint url or env
+    pub endpoint: RpcEndpoint,
+
+    /// The number of retries.
+    pub retries: Option<u32>,
+
+    /// Initial retry backoff.
+    pub retry_backoff: Option<u64>,
+
+    /// The available compute units per second.
+    ///
+    /// See also <https://docs.alchemy.com/reference/compute-units#what-are-cups-compute-units-per-second>
+    pub compute_units_per_second: Option<u64>,
+}
+
+impl RpcEndpointConfig {
+    /// Returns the url this type holds, see [RpcEndpoints::resolve()]
+    pub fn resolve(self) -> Result<String, UnresolvedEnvVarError> {
+        self.endpoint.resolve()
+    }
+}
+
+impl fmt::Display for RpcEndpointConfig {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let RpcEndpointConfig { endpoint, retries, retry_backoff, compute_units_per_second } = self;
+
+        write!(f, "{}", endpoint)?;
+
+        if let Some(retries) = retries {
+            write!(f, ", retries={}", retries)?;
+        }
+
+        if let Some(retry_backoff) = retry_backoff {
+            write!(f, ", retry_backoff={}", retry_backoff)?;
+        }
+
+        if let Some(compute_units_per_second) = compute_units_per_second {
+            write!(f, ", compute_units_per_second={}", compute_units_per_second)?;
+        }
+
+        Ok(())
+    }
+}
+
+impl Serialize for RpcEndpointConfig {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        if self.retries.is_none() &&
+            self.retry_backoff.is_none() &&
+            self.compute_units_per_second.is_none()
+        {
+            // serialize as endpoint if there's no additional config
+            return self.endpoint.serialize(serializer);
+        } else {
+            let mut map = serializer.serialize_map(Some(4))?;
+            map.serialize_entry("endpoint", &self.endpoint)?;
+            map.serialize_entry("retries", &self.retries)?;
+            map.serialize_entry("retry_backoff", &self.retry_backoff)?;
+            map.serialize_entry("compute_units_per_second", &self.compute_units_per_second)?;
+            map.end()
+        }
+    }
+}
+
+impl<'de> Deserialize<'de> for RpcEndpointConfig {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let value = serde_json::Value::deserialize(deserializer)?;
+        if value.is_string() {
+            return Ok(Self {
+                endpoint: serde_json::from_value(value).map_err(serde::de::Error::custom)?,
+                ..Default::default()
+            });
+        }
+
+        #[derive(Deserialize)]
+        struct RpcEndpointConfigInner {
+            #[serde(alias = "url")]
+            endpoint: RpcEndpoint,
+            retries: Option<u32>,
+            retry_backoff: Option<u64>,
+            compute_units_per_second: Option<u64>,
+        }
+
+        let RpcEndpointConfigInner { endpoint, retries, retry_backoff, compute_units_per_second } =
+            serde_json::from_value(value).map_err(serde::de::Error::custom)?;
+
+        Ok(RpcEndpointConfig { endpoint, retries, retry_backoff, compute_units_per_second })
+    }
+}
+
+impl Into<RpcEndpointType> for RpcEndpointConfig {
+    fn into(self) -> RpcEndpointType {
+        RpcEndpointType::Config(self)
+    }
+}
+
+impl Default for RpcEndpointConfig {
+    fn default() -> Self {
+        Self {
+            endpoint: RpcEndpoint::Url("http://localhost:8545".to_string()),
+            retries: None,
+            retry_backoff: None,
+            compute_units_per_second: None,
+        }
+    }
+}
+
 /// Container type for _resolved_ endpoints, see [RpcEndpoints::resolve_all()]
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub struct ResolvedRpcEndpoints {
@@ -164,3 +364,40 @@ impl DerefMut for ResolvedRpcEndpoints {
         &mut self.endpoints
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn serde_rpc_config() {
+        let s = r#"{
+            "endpoint": "http://localhost:8545",
+            "retries": 5,
+            "retry_backoff": 250,
+            "compute_units_per_second": 100
+        }"#;
+        let config: RpcEndpointConfig = serde_json::from_str(s).unwrap();
+        assert_eq!(
+            config,
+            RpcEndpointConfig {
+                endpoint: RpcEndpoint::Url("http://localhost:8545".to_string()),
+                retries: Some(5),
+                retry_backoff: Some(250),
+                compute_units_per_second: Some(100),
+            }
+        );
+
+        let s = "\"http://localhost:8545\"";
+        let config: RpcEndpointConfig = serde_json::from_str(s).unwrap();
+        assert_eq!(
+            config,
+            RpcEndpointConfig {
+                endpoint: RpcEndpoint::Url("http://localhost:8545".to_string()),
+                retries: None,
+                retry_backoff: None,
+                compute_units_per_second: None,
+            }
+        );
+    }
+}

crates/config/src/lib.rs

@@ -2519,7 +2519,7 @@ mod tests {
     use super::*;
     use crate::{
         cache::{CachedChains, CachedEndpoints},
-        endpoints::RpcEndpoint,
+        endpoints::{RpcEndpoint, RpcEndpointConfig, RpcEndpointType},
         etherscan::ResolvedEtherscanConfigs,
         fs_permissions::PathPermission,
     };
@@ -3120,6 +3120,57 @@ mod tests {
         });
     }
 
+    #[test]
+    fn test_resolve_rpc_config() {
+        figment::Jail::expect_with(|jail| {
+            jail.create_file(
+                "foundry.toml",
+                r#"
+                [rpc_endpoints]
+                optimism = "https://example.com/"
+                mainnet = { endpoint = "${_CONFIG_MAINNET}", retries = 3, retry_backoff = 1000, compute_units_per_second = 1000 }
+            "#,
+            )?;
+            jail.set_env("_CONFIG_MAINNET", "https://eth-mainnet.alchemyapi.io/v2/123455");
+
+            let config = Config::load();
+            assert_eq!(
+                RpcEndpoints::new([
+                    (
+                        "optimism",
+                        RpcEndpointType::String(RpcEndpoint::Url(
+                            "https://example.com/".to_string()
+                        ))
+                    ),
+                    (
+                        "mainnet",
+                        RpcEndpointType::Config(RpcEndpointConfig {
+                            endpoint: RpcEndpoint::Env("${_CONFIG_MAINNET}".to_string()),
+                            retries: Some(3),
+                            retry_backoff: Some(1000),
+                            compute_units_per_second: Some(1000),
+                        })
+                    ),
+                ]),
+                config.rpc_endpoints
+            );
+
+            let resolved = config.rpc_endpoints.resolved();
+            assert_eq!(
+                RpcEndpoints::new([
+                    ("optimism", RpcEndpoint::Url("https://example.com/".to_string())),
+                    (
+                        "mainnet",
+                        RpcEndpoint::Url("https://eth-mainnet.alchemyapi.io/v2/123455".to_string())
+                    ),
+                ])
+                .resolved(),
+                resolved
+            );
+            Ok(())
+        })
+    }
+
     #[test]
     fn test_resolve_endpoints() {
         figment::Jail::expect_with(|jail| {