use std::time::Duration;
pub use dig_constants::RPC_DIG_NET_URL as RPC_DIG_NET;
pub use dig_constants::DIG_LOCAL_HOST;
pub use dig_constants::DIG_NODE_PORT as DEFAULT_LOCAL_NODE_PORT;
pub const DEFAULT_PROBE_TIMEOUT: Duration = Duration::from_millis(600);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ResolvedTier {
Override,
DigLocal,
Localhost,
PublicGateway,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ResolvedNode {
pub base_url: String,
pub tier: ResolvedTier,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LocalRung {
pub url: String,
pub tier: ResolvedTier,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TransportMode {
#[default]
Https,
Mtls,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OverrideSource {
Flag,
Env,
Config,
}
#[async_trait::async_trait]
pub trait HealthProbe: Send + Sync {
async fn probe(&self, base_url: &str, timeout: Duration) -> bool;
}
#[derive(Debug, Clone, Default)]
pub struct OverrideInputs {
pub flag: Option<String>,
pub env_var: Option<String>,
pub config_value: Option<String>,
}
impl OverrideInputs {
fn resolve(&self) -> Option<(&str, OverrideSource)> {
if let Some(v) = self.flag.as_deref() {
return Some((v, OverrideSource::Flag));
}
if let Some(v) = self.env_var.as_deref() {
return Some((v, OverrideSource::Env));
}
if let Some(v) = self.config_value.as_deref() {
return Some((v, OverrideSource::Config));
}
None
}
}
pub async fn resolve_node(
overrides: &OverrideInputs,
local_rungs: &[LocalRung],
probe: &dyn HealthProbe,
timeout: Duration,
) -> ResolvedNode {
if let Some((url, _source)) = overrides.resolve() {
return ResolvedNode {
base_url: url.trim_end_matches('/').to_string(),
tier: ResolvedTier::Override,
};
}
for rung in local_rungs {
if probe.probe(&rung.url, timeout).await {
return ResolvedNode {
base_url: rung.url.trim_end_matches('/').to_string(),
tier: rung.tier,
};
}
}
ResolvedNode {
base_url: RPC_DIG_NET.to_string(),
tier: ResolvedTier::PublicGateway,
}
}
#[must_use]
pub fn local_urls(port: u16) -> Vec<LocalRung> {
vec![
LocalRung {
url: format!("https://{DIG_LOCAL_HOST}"),
tier: ResolvedTier::DigLocal,
},
LocalRung {
url: format!("http://{DIG_LOCAL_HOST}"),
tier: ResolvedTier::DigLocal,
},
LocalRung {
url: format!("http://localhost:{port}"),
tier: ResolvedTier::Localhost,
},
]
}
#[must_use]
pub fn override_source(overrides: &OverrideInputs) -> Option<OverrideSource> {
overrides.resolve().map(|(_, s)| s)
}
#[derive(Debug, Default)]
pub struct CachedResolver {
cached: tokio::sync::OnceCell<ResolvedNode>,
}
impl CachedResolver {
#[must_use]
pub fn new() -> Self {
Self {
cached: tokio::sync::OnceCell::new(),
}
}
pub async fn get_or_resolve(
&self,
overrides: &OverrideInputs,
local_rungs: &[LocalRung],
probe: &dyn HealthProbe,
timeout: Duration,
) -> ResolvedNode {
self.cached
.get_or_init(|| resolve_node(overrides, local_rungs, probe, timeout))
.await
.clone()
}
}
#[cfg(feature = "http-probe")]
pub struct HttpHealthProbe {
http: reqwest::Client,
}
#[cfg(feature = "http-probe")]
impl HttpHealthProbe {
#[must_use]
pub fn new(http: reqwest::Client) -> Self {
Self { http }
}
}
#[cfg(feature = "http-probe")]
impl Default for HttpHealthProbe {
fn default() -> Self {
Self::new(
reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.unwrap_or_else(|_| reqwest::Client::new()),
)
}
}
#[cfg(feature = "http-probe")]
#[async_trait::async_trait]
impl HealthProbe for HttpHealthProbe {
async fn probe(&self, base_url: &str, timeout: Duration) -> bool {
let url = format!("{}/health", base_url.trim_end_matches('/'));
let request = self.http.get(&url).send();
match tokio::time::timeout(timeout, request).await {
Ok(Ok(resp)) => resp.status().is_success(),
Ok(Err(_)) | Err(_) => false,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Mutex;
#[derive(Default)]
struct ScriptedProbe {
answers: std::collections::HashMap<String, bool>,
calls: Mutex<Vec<String>>,
}
impl ScriptedProbe {
fn new(answers: &[(&str, bool)]) -> Self {
Self {
answers: answers.iter().map(|(k, v)| (k.to_string(), *v)).collect(),
calls: Mutex::new(Vec::new()),
}
}
fn calls(&self) -> Vec<String> {
self.calls.lock().unwrap().clone()
}
}
#[async_trait::async_trait]
impl HealthProbe for ScriptedProbe {
async fn probe(&self, base_url: &str, _timeout: Duration) -> bool {
self.calls.lock().unwrap().push(base_url.to_string());
self.answers.get(base_url).copied().unwrap_or(false)
}
}
const DIG_LOCAL_HTTPS: &str = "https://dig.local";
const DIG_LOCAL_HTTP: &str = "http://dig.local";
const LOCALHOST_HTTP: &str = "http://localhost:9778";
const T: Duration = Duration::from_millis(50);
fn rungs() -> Vec<LocalRung> {
local_urls(DEFAULT_LOCAL_NODE_PORT)
}
#[tokio::test]
async fn prefers_dig_local_when_it_answers() {
let probe = ScriptedProbe::new(&[(DIG_LOCAL_HTTPS, true), (LOCALHOST_HTTP, true)]);
let resolved = resolve_node(&OverrideInputs::default(), &rungs(), &probe, T).await;
assert_eq!(resolved.base_url, DIG_LOCAL_HTTPS);
assert_eq!(resolved.tier, ResolvedTier::DigLocal);
assert_eq!(probe.calls(), vec![DIG_LOCAL_HTTPS.to_string()]);
}
#[tokio::test]
async fn falls_soft_to_http_dig_local_when_https_is_unreachable() {
let probe = ScriptedProbe::new(&[(DIG_LOCAL_HTTPS, false), (DIG_LOCAL_HTTP, true)]);
let resolved = resolve_node(&OverrideInputs::default(), &rungs(), &probe, T).await;
assert_eq!(resolved.base_url, DIG_LOCAL_HTTP);
assert_eq!(resolved.tier, ResolvedTier::DigLocal);
assert_eq!(
probe.calls(),
vec![DIG_LOCAL_HTTPS.to_string(), DIG_LOCAL_HTTP.to_string()]
);
}
#[tokio::test]
async fn falls_through_to_localhost_when_dig_local_is_unreachable() {
let probe = ScriptedProbe::new(&[
(DIG_LOCAL_HTTPS, false),
(DIG_LOCAL_HTTP, false),
(LOCALHOST_HTTP, true),
]);
let resolved = resolve_node(&OverrideInputs::default(), &rungs(), &probe, T).await;
assert_eq!(resolved.base_url, LOCALHOST_HTTP);
assert_eq!(resolved.tier, ResolvedTier::Localhost);
assert_eq!(
probe.calls(),
vec![
DIG_LOCAL_HTTPS.to_string(),
DIG_LOCAL_HTTP.to_string(),
LOCALHOST_HTTP.to_string(),
]
);
}
#[tokio::test]
async fn falls_through_to_public_gateway_as_final_fallback() {
let probe = ScriptedProbe::new(&[]);
let resolved = resolve_node(&OverrideInputs::default(), &rungs(), &probe, T).await;
assert_eq!(resolved.base_url, RPC_DIG_NET);
assert_eq!(resolved.tier, ResolvedTier::PublicGateway);
}
#[tokio::test]
async fn timeout_behaves_as_no_response_and_falls_through() {
struct NeverRespondsProbe;
#[async_trait::async_trait]
impl HealthProbe for NeverRespondsProbe {
async fn probe(&self, _base_url: &str, _timeout: Duration) -> bool {
false
}
}
let resolved = resolve_node(
&OverrideInputs::default(),
&rungs(),
&NeverRespondsProbe,
Duration::from_millis(5),
)
.await;
assert_eq!(resolved.tier, ResolvedTier::PublicGateway);
}
#[tokio::test]
async fn explicit_override_wins_without_probing_anything() {
let probe = ScriptedProbe::new(&[(DIG_LOCAL_HTTPS, true), (LOCALHOST_HTTP, true)]);
let overrides = OverrideInputs {
flag: Some("https://custom.example:9999".to_string()),
..Default::default()
};
let resolved = resolve_node(&overrides, &rungs(), &probe, T).await;
assert_eq!(resolved.base_url, "https://custom.example:9999");
assert_eq!(resolved.tier, ResolvedTier::Override);
assert!(probe.calls().is_empty());
}
#[tokio::test]
async fn override_trailing_slash_is_normalized() {
let probe = ScriptedProbe::new(&[]);
let overrides = OverrideInputs {
flag: Some("https://custom.example/".to_string()),
..Default::default()
};
let resolved = resolve_node(&overrides, &rungs(), &probe, T).await;
assert_eq!(resolved.base_url, "https://custom.example");
}
#[test]
fn flag_wins_over_env_and_config() {
let overrides = OverrideInputs {
flag: Some("flag-url".into()),
env_var: Some("env-url".into()),
config_value: Some("config-url".into()),
};
assert_eq!(override_source(&overrides), Some(OverrideSource::Flag));
}
#[test]
fn env_wins_over_config_when_no_flag() {
let overrides = OverrideInputs {
flag: None,
env_var: Some("env-url".into()),
config_value: Some("config-url".into()),
};
assert_eq!(override_source(&overrides), Some(OverrideSource::Env));
}
#[test]
fn config_used_when_no_flag_or_env() {
let overrides = OverrideInputs {
config_value: Some("config-url".into()),
..Default::default()
};
assert_eq!(override_source(&overrides), Some(OverrideSource::Config));
}
#[test]
fn no_override_when_all_absent() {
assert_eq!(override_source(&OverrideInputs::default()), None);
}
#[test]
fn local_rungs_target_real_dig_node_listeners() {
struct Listener {
scheme: &'static str,
host: &'static str,
port: Option<u16>,
tier: ResolvedTier,
}
let expected = [
Listener {
scheme: "https",
host: "dig.local",
port: None,
tier: ResolvedTier::DigLocal,
},
Listener {
scheme: "http",
host: "dig.local",
port: None,
tier: ResolvedTier::DigLocal,
},
Listener {
scheme: "http",
host: "localhost",
port: Some(9778),
tier: ResolvedTier::Localhost,
},
];
fn decompose(url: &str) -> (&str, &str, Option<u16>) {
let (scheme, rest) = url.split_once("://").expect("rung URL has a scheme");
match rest.split_once(':') {
Some((host, port)) => (scheme, host, Some(port.parse().expect("numeric port"))),
None => (scheme, rest, None),
}
}
let rungs = local_urls(DEFAULT_LOCAL_NODE_PORT);
assert_eq!(rungs.len(), expected.len(), "one rung per SPEC listener");
for (rung, want) in rungs.iter().zip(expected.iter()) {
let (scheme, host, port) = decompose(&rung.url);
assert_eq!(scheme, want.scheme, "scheme for {}", rung.url);
assert_eq!(host, want.host, "host for {}", rung.url);
assert_eq!(port, want.port, "port for {}", rung.url);
assert_eq!(rung.tier, want.tier, "tier for {}", rung.url);
}
}
#[tokio::test]
async fn cached_resolver_probes_only_once() {
struct CountingProbe {
calls: AtomicUsize,
}
#[async_trait::async_trait]
impl HealthProbe for CountingProbe {
async fn probe(&self, _base_url: &str, _timeout: Duration) -> bool {
self.calls.fetch_add(1, Ordering::SeqCst);
true
}
}
let probe = CountingProbe {
calls: AtomicUsize::new(0),
};
let cache = CachedResolver::new();
let overrides = OverrideInputs::default();
let first = cache.get_or_resolve(&overrides, &rungs(), &probe, T).await;
let second = cache.get_or_resolve(&overrides, &rungs(), &probe, T).await;
assert_eq!(first, second);
assert_eq!(probe.calls.load(Ordering::SeqCst), 1);
}
#[test]
fn default_transport_is_https() {
assert_eq!(TransportMode::default(), TransportMode::Https);
}
}