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, 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,
dig_local_url: &str,
localhost_url: &str,
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,
};
}
if probe.probe(dig_local_url, timeout).await {
return ResolvedNode {
base_url: dig_local_url.trim_end_matches('/').to_string(),
tier: ResolvedTier::DigLocal,
};
}
if probe.probe(localhost_url, timeout).await {
return ResolvedNode {
base_url: localhost_url.trim_end_matches('/').to_string(),
tier: ResolvedTier::Localhost,
};
}
ResolvedNode {
base_url: RPC_DIG_NET.to_string(),
tier: ResolvedTier::PublicGateway,
}
}
#[must_use]
pub fn local_urls(port: u16) -> (String, String) {
(
format!("https://{DIG_LOCAL_HOST}:{port}"),
format!("https://localhost:{port}"),
)
}
#[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,
dig_local_url: &str,
localhost_url: &str,
probe: &dyn HealthProbe,
timeout: Duration,
) -> ResolvedNode {
self.cached
.get_or_init(|| resolve_node(overrides, dig_local_url, localhost_url, 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: &str = "https://dig.local:9778";
const LOCALHOST: &str = "https://localhost:9778";
const T: Duration = Duration::from_millis(50);
#[tokio::test]
async fn prefers_dig_local_when_it_answers() {
let probe = ScriptedProbe::new(&[(DIG_LOCAL, true), (LOCALHOST, true)]);
let resolved =
resolve_node(&OverrideInputs::default(), DIG_LOCAL, LOCALHOST, &probe, T).await;
assert_eq!(resolved.base_url, DIG_LOCAL);
assert_eq!(resolved.tier, ResolvedTier::DigLocal);
assert_eq!(probe.calls(), vec![DIG_LOCAL.to_string()]);
}
#[tokio::test]
async fn falls_through_to_localhost_when_dig_local_is_unreachable() {
let probe = ScriptedProbe::new(&[(DIG_LOCAL, false), (LOCALHOST, true)]);
let resolved =
resolve_node(&OverrideInputs::default(), DIG_LOCAL, LOCALHOST, &probe, T).await;
assert_eq!(resolved.base_url, LOCALHOST);
assert_eq!(resolved.tier, ResolvedTier::Localhost);
assert_eq!(
probe.calls(),
vec![DIG_LOCAL.to_string(), LOCALHOST.to_string()]
);
}
#[tokio::test]
async fn falls_through_to_public_gateway_as_final_fallback() {
let probe = ScriptedProbe::new(&[(DIG_LOCAL, false), (LOCALHOST, false)]);
let resolved =
resolve_node(&OverrideInputs::default(), DIG_LOCAL, LOCALHOST, &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(),
DIG_LOCAL,
LOCALHOST,
&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, true), (LOCALHOST, true)]);
let overrides = OverrideInputs {
flag: Some("https://custom.example:9999".to_string()),
..Default::default()
};
let resolved = resolve_node(&overrides, DIG_LOCAL, LOCALHOST, &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, DIG_LOCAL, LOCALHOST, &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_urls_uses_https_host_and_port() {
let (dig_local, localhost) = local_urls(DEFAULT_LOCAL_NODE_PORT);
assert_eq!(dig_local, "https://dig.local:9778");
assert_eq!(localhost, "https://localhost:9778");
}
#[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, DIG_LOCAL, LOCALHOST, &probe, T)
.await;
let second = cache
.get_or_resolve(&overrides, DIG_LOCAL, LOCALHOST, &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);
}
}