diff --git a/Cargo.lock b/Cargo.lock index 997e389..be69a74 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -19,6 +19,12 @@ version = "1.0.56" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4361135be9122e0870de935d7c439aef945b9f9ddd4199a553b5270b49c82a27" +[[package]] +name = "arc-swap" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5d78ce20460b82d3fa150275ed9d55e21064fc7951177baacf86a145c4a4b1f" + [[package]] name = "async-anyhow-logger" version = "0.1.0" @@ -258,6 +264,24 @@ dependencies = [ "syn", ] +[[package]] +name = "data-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ee2393c4a91429dffb4bedf19f4d6abf27d8a732c8ce4980305d782e5426d57" + +[[package]] +name = "enum-as-inner" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "21cdad81446a7f7dc43f6a77409efeb9733d2fa65553efef6018ef257c959b73" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "fern" version = "0.6.0" @@ -407,6 +431,17 @@ dependencies = [ "libc", ] +[[package]] +name = "hostname" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c731c3e10504cc8ed35cfe2f1db4c9274c3d35fa486e3b31df46f068ef3e867" +dependencies = [ + "libc", + "match_cfg", + "winapi", +] + [[package]] name = "http" version = "0.2.6" @@ -476,6 +511,17 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" +[[package]] +name = "idna" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "418a0a6fab821475f634efe3ccc45c013f742efe03d853e8d3355d5cb850ecf8" +dependencies = [ + "matches", + "unicode-bidi", + "unicode-normalization", +] + [[package]] name = "indexmap" version = "1.8.1" @@ -495,6 +541,24 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "ipconfig" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "723519edce41262b05d4143ceb95050e4c614f483e78e9fd9e39a8275a84ad98" +dependencies = [ + "socket2", + "widestring", + "winapi", + "winreg", +] + +[[package]] +name = "ipnet" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879d54834c8c76457ef4293a689b2a8c59b076067ad77b15efafbb05f92a592b" + [[package]] name = "itoa" version = "1.0.1" @@ -522,6 +586,12 @@ version = "0.2.121" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "efaa7b300f3b5fe8eb6bf21ce3895e1751d9665086af2d64b42f19701015ff4f" +[[package]] +name = "linked-hash-map" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0717cef1bc8b636c6e1c1bbdefc09e6322da8a9321966e8928ef80d20f7f770f" + [[package]] name = "lock_api" version = "0.4.7" @@ -542,6 +612,15 @@ dependencies = [ "serde", ] +[[package]] +name = "lru-cache" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31e24f1ad8321ca0e8a1e0ac13f23cb668e6f5466c2c57319f6a5cf1cc8e3b1c" +dependencies = [ + "linked-hash-map", +] + [[package]] name = "mach" version = "0.3.2" @@ -551,6 +630,12 @@ dependencies = [ "libc", ] +[[package]] +name = "match_cfg" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffbee8634e0d45d258acb448e7eaab3fce7a0a467395d4d9f228e3c1f01fb2e4" + [[package]] name = "matches" version = "0.1.9" @@ -699,9 +784,9 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.10.0" +version = "1.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87f3e037eac156d1775da914196f0f37741a274155e34a0b7e427c35d2a2ecb9" +checksum = "18a6dbe30758c9f83eb00cbea4ac95966305f5a7772f3f42ebfc7fc7eddbd8e1" [[package]] name = "os_str_bytes" @@ -771,6 +856,7 @@ name = "peshming" version = "0.5.1" dependencies = [ "anyhow", + "arc-swap", "async-anyhow-logger", "axum", "chrono", @@ -783,11 +869,14 @@ dependencies = [ "log", "metrics", "metrics-exporter-prometheus", + "once_cell", + "rand", "serde", "serde_with", "tokio", "tokio-icmp-echo", "toml", + "trust-dns-resolver", ] [[package]] @@ -877,6 +966,12 @@ dependencies = [ "winapi", ] +[[package]] +name = "quick-error" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" + [[package]] name = "quote" version = "1.0.17" @@ -934,6 +1029,16 @@ dependencies = [ "bitflags", ] +[[package]] +name = "resolv-conf" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52e44394d2086d010551b14b53b1f24e31647570cd1deb0379e2c21b329aba00" +dependencies = [ + "hostname", + "quick-error", +] + [[package]] name = "roff" version = "0.2.1" @@ -1133,15 +1238,33 @@ dependencies = [ "winapi", ] +[[package]] +name = "tinyvec" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87cc5ceb3875bb20c2890005a4e226a4651264a5c75edb2421b52861a0a0cb50" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cda74da7e1a664f795bb1f8a87ec406fb89a02522cf6e50620d016add6dbbf5c" + [[package]] name = "tokio" version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af73ac49756f3f7c01172e34a23e5d0216f6c32333757c2c61feb2bbff5a5ee" dependencies = [ + "bytes", "libc", + "memchr", "mio", "num_cpus", + "parking_lot 0.12.0", "pin-project-lite", "socket2", "tokio-macros", @@ -1266,18 +1389,90 @@ dependencies = [ "lazy_static", ] +[[package]] +name = "trust-dns-proto" +version = "0.21.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c31f240f59877c3d4bb3b3ea0ec5a6a0cff07323580ff8c7a605cd7d08b255d" +dependencies = [ + "async-trait", + "cfg-if", + "data-encoding", + "enum-as-inner", + "futures-channel", + "futures-io", + "futures-util", + "idna", + "ipnet", + "lazy_static", + "log", + "rand", + "smallvec", + "thiserror", + "tinyvec", + "tokio", + "url", +] + +[[package]] +name = "trust-dns-resolver" +version = "0.21.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4ba72c2ea84515690c9fcef4c6c660bb9df3036ed1051686de84605b74fd558" +dependencies = [ + "cfg-if", + "futures-util", + "ipconfig", + "lazy_static", + "log", + "lru-cache", + "parking_lot 0.12.0", + "resolv-conf", + "smallvec", + "thiserror", + "tokio", + "trust-dns-proto", +] + [[package]] name = "try-lock" version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "59547bce71d9c38b83d9c0e92b6066c4253371f15005def0c30d9657f50c7642" +[[package]] +name = "unicode-bidi" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "099b7128301d285f79ddd55b9a83d5e6b9e97c92e0ea0daebee7263e932de992" + +[[package]] +name = "unicode-normalization" +version = "0.1.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "854cbdc4f7bc6ae19c820d44abdc3277ac3e1b2b93db20a636825d9322fb60e6" +dependencies = [ + "tinyvec", +] + [[package]] name = "unicode-xid" version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ccb82d61f80a663efe1f787a51b16b5a51e3314d6ac365b08639f52387b33f3" +[[package]] +name = "url" +version = "2.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a507c383b2d33b5fc35d1861e77e6b383d158b2da5e14fe51b83dfedf6fd578c" +dependencies = [ + "form_urlencoded", + "idna", + "matches", + "percent-encoding", +] + [[package]] name = "version_check" version = "0.9.4" @@ -1370,6 +1565,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "widestring" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "17882f045410753661207383517a6f62ec3dbeb6a4ed2acce01f0728238d1983" + [[package]] name = "winapi" version = "0.3.9" @@ -1443,3 +1644,12 @@ name = "windows_x86_64_msvc" version = "0.34.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d19538ccc21819d01deaf88d6a17eae6596a12e9aafdbb97916fb49896d89de9" + +[[package]] +name = "winreg" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0120db82e8a1e0b9fb3345a539c478767c0048d842860994d96113d5b667bd69" +dependencies = [ + "winapi", +] diff --git a/Cargo.toml b/Cargo.toml index e36396d..8e1185e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ build = "build.rs" [dependencies] toml = "0.5" futures = "0.3" -tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "parking_lot"] } clap = { version = "3", features = ["derive", "wrap_help"] } fern = "0.6" log = { version = "0.4", features = ["serde"] } @@ -26,6 +26,10 @@ async-anyhow-logger = "0.1" axum = "0.5" metrics = "0.18" metrics-exporter-prometheus = { version = "0.9", default-features = false } +trust-dns-resolver = "0.21.2" +arc-swap = "1.5.0" +once_cell = "1.13.0" +rand = "0.8.5" [build-dependencies] clap = { version = "3", features = ["derive", "wrap_help"] } diff --git a/build.rs b/build.rs index 1505a42..89ddb9e 100644 --- a/build.rs +++ b/build.rs @@ -23,7 +23,8 @@ use clap::{ArgEnum, CommandFactory}; use clap_complete::{generate_to, Shell}; use cli::Cli; -#[path = "src/cli.rs"] mod cli; +#[path = "src/cli.rs"] +mod cli; fn main() -> std::io::Result<()> { let mut cli = Cli::command(); diff --git a/config.toml.sample b/config.toml.sample index dc9a8fd..72b367a 100644 --- a/config.toml.sample +++ b/config.toml.sample @@ -12,8 +12,11 @@ listener = "[::]:9898" # will ping the primary and secondary IP of cloudflare's 1.1.1.1 DNS service # every 500ms, or twice per second. [ping.hosts] -"1.1.1.1" = 500 -"1.0.0.1" = 500 +"192.0.2.142" = 500 +"198.51.100.17" = 500 +"203.0.113.55" = 500 +"2001:DB8::C0:FF:EE" = 500 +"example.org" = 500 # Configure logging is also possible here instead of using the CLI. If both are # specified, the more verbose of the two will be used. diff --git a/src/config.rs b/src/config.rs index 22474fb..7a27f3c 100644 --- a/src/config.rs +++ b/src/config.rs @@ -50,7 +50,23 @@ pub(crate) struct PingConfig { pub(crate) timeout: Duration, #[serde(default = "default_buckets")] pub(crate) bucket_sizes: Vec, - pub(crate) hosts: HashMap, + pub(crate) hosts: HashMap, +} + +#[derive(Deserialize, Clone, Hash, Eq, PartialEq, Debug)] +#[serde(untagged)] +pub(crate) enum Host { + IpAddr(std::net::IpAddr), + Domain(String), +} + +impl std::fmt::Display for Host { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Host::IpAddr(addr) => write!(f, "{}", addr), + Host::Domain(domain) => write!(f, "{}", domain), + } + } } fn default_timeout() -> Duration { diff --git a/src/ping.rs b/src/ping.rs index 6fd9338..b5f6c50 100644 --- a/src/ping.rs +++ b/src/ping.rs @@ -17,22 +17,54 @@ * You should have received a copy of the GNU Affero General Public License * * along with this program. If not, see . * ********************************************************************************/ -use crate::config::Config; +use crate::config::{Config, Host}; use anyhow::{Context, Result}; +use arc_swap::ArcSwap; use async_anyhow_logger::catch; use log::{info, trace}; use metrics::histogram; -use std::net::IpAddr; -use std::time::Duration; +use once_cell::sync::OnceCell as SyncOnceCell; +use tokio::sync::OnceCell as AsyncOnceCell; +use std::{net::IpAddr, sync::Arc, time::Duration, collections::HashMap}; use tokio_icmp_echo::{PingFuture, Pinger}; +use trust_dns_resolver::{TokioAsyncResolver, lookup_ip::LookupIp}; + +static PING_IP_TARGETS: SyncOnceCell>>> = SyncOnceCell::new(); + +async fn pinger() -> Result<&'static Pinger> { + static INSTANCE: AsyncOnceCell = AsyncOnceCell::const_new(); + INSTANCE.get_or_try_init(|| async { + Pinger::new().await.context("Couldn't create pinger!") + }).await +} + +fn resolver() -> Result<&'static TokioAsyncResolver> { + static INSTANCE: SyncOnceCell = SyncOnceCell::new(); + INSTANCE.get_or_try_init(|| { + TokioAsyncResolver::tokio_from_system_conf().context("Couldn't start resolver!") + }) +} + pub(crate) async fn start_pinging_hosts(config: &Config) -> Result<()> { - let pinger = Pinger::new().await.context("Couldn't create pinger")?; let mut handles = vec![]; - for (host, interval) in config.ping.hosts.clone() { + let mut map = HashMap::new(); + for (host, _interval) in config.ping.hosts.clone().into_iter() { + match host.clone() { + Host::IpAddr(addr) => map.insert(host, ArcSwap::new(Arc::new(vec![addr]))), + Host::Domain(domain) => { + let lookup = get_host_addresses(&domain).await?; + map.insert(host, ArcSwap::new(Arc::new(lookup.iter().collect()))) + }, + }; + } + PING_IP_TARGETS.set(map).unwrap(); + for (host, interval) in config.ping.hosts.clone().into_iter() { + if let Host::Domain(domain) = host.clone() { + tokio::spawn(catch(refresh_host_addresses(&PING_IP_TARGETS.get().unwrap()[&host], domain.clone()))); + } info!("Spawn ping task for {}", host); handles.push(tokio::spawn(ping_host( - pinger.clone(), host, interval, config.ping.timeout, @@ -43,31 +75,53 @@ pub(crate) async fn start_pinging_hosts(config: &Config) -> Result<()> { Ok(()) } -async fn ping_host(pinger: Pinger, host: IpAddr, interval: u64, timeout: Duration) -> Result<()> { - let mut pingchain = pinger.chain(host).timeout(timeout); +async fn ping_host(host: Host, interval: u64, timeout: Duration) -> Result<()> { + let pinger = pinger().await?; + let name: String = match host.clone() { + Host::IpAddr(addr) => addr.to_string(), + Host::Domain(name) => name, + }; + let targets = &PING_IP_TARGETS.get().unwrap()[&host]; let mut interval = tokio::time::interval(Duration::from_millis(interval)); - let host_string = host.to_string(); + let ident = rand::random(); + let mut seq_cnt = 0; loop { interval.tick().await; - tokio::spawn(catch(handle_ping_result( - pingchain.send(), - host_string.clone(), - timeout, - ))); + for target in &**targets.load() { + tokio::spawn(catch(handle_ping_result( + pinger.ping(*target, ident, seq_cnt, timeout), + name.clone(), + target.to_string(), + timeout, + ))); + seq_cnt = seq_cnt.wrapping_add(1); + } } } -async fn handle_ping_result(result: PingFuture, host: String, timeout: Duration) -> Result<()> { - let pong = result.await.context(format!("Couldn't ping {}", &host))?; +async fn refresh_host_addresses(targets: &ArcSwap>, name: String) -> Result<()> { + loop { + let lookup = get_host_addresses(&name).await?; + targets.store(Arc::new(lookup.iter().collect())); + tokio::time::sleep_until(lookup.valid_until().into()).await; + } +} + +async fn get_host_addresses(name: &str) -> Result { + Ok(resolver()?.lookup_ip(name).await?) +} + +async fn handle_ping_result(result: PingFuture, name: String, ip: String, timeout: Duration) -> Result<()> { + let pong = result.await.context(format!("Couldn't ping {}", &name))?; match pong { Some(time) => { let ms = time.as_millis(); - trace!("Received pong from {} after {} ms", &host, &ms); - histogram!("ping_rtt_milliseconds", ms as f64, "target" => host); + trace!("Received pong from {} after {} ms", &name, &ms); + histogram!("ping_rtt_milliseconds", ms as f64, "target" => name, "ip" => ip); } None => { - trace!("Received no response from {} within timeout", &host); - histogram!("ping_rtt_milliseconds", timeout.as_millis() as f64, "target" => host); + trace!("Received no response from {} within timeout", &name); + histogram!("ping_rtt_milliseconds", timeout.as_millis() as f64, "target" => name, "ip" => ip); } };