From bcb21da88e74f8d50a7b1ec302e04dd2f2fe1476 Mon Sep 17 00:00:00 2001 From: Lewis Date: Mon, 4 May 2026 16:31:46 +0300 Subject: [PATCH] feat(knot-proxy): pre-work, dns guard, circuit breaker Lewis: May this revision serve well! --- Cargo.lock | 17 +++ Cargo.toml | 2 + crates/knot-proxy/Cargo.toml | 20 +++ crates/knot-proxy/src/breaker.rs | 239 +++++++++++++++++++++++++++++++ crates/knot-proxy/src/dns.rs | 112 +++++++++++++++ 5 files changed, 390 insertions(+) create mode 100644 crates/knot-proxy/Cargo.toml create mode 100644 crates/knot-proxy/src/breaker.rs create mode 100644 crates/knot-proxy/src/dns.rs diff --git a/Cargo.lock b/Cargo.lock index f8f9261..5431e73 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -209,6 +209,7 @@ dependencies = [ "axum", "bobbin-edge-index", "bobbin-ingest", + "bobbin-knot-proxy", "bobbin-record-lru", "bobbin-slingshot-client", "bobbin-types", @@ -253,6 +254,21 @@ dependencies = [ "url", ] +[[package]] +name = "bobbin-knot-proxy" +version = "0.0.1" +dependencies = [ + "bytes", + "futures", + "http", + "reqwest", + "scc", + "thiserror 2.0.18", + "tokio", + "url", + "wiremock", +] + [[package]] name = "bobbin-record-lru" version = "0.0.1" @@ -305,6 +321,7 @@ version = "0.0.1" dependencies = [ "axum", "bobbin-edge-index", + "bobbin-knot-proxy", "bobbin-record-lru", "bobbin-slingshot-client", "bobbin-types", diff --git a/Cargo.toml b/Cargo.toml index 9a8e91b..16a4797 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ members = [ "crates/ingest", "crates/slingshot-client", "crates/record-lru", + "crates/knot-proxy", "crates/xrpc", ] @@ -22,6 +23,7 @@ bobbin-edge-index = { path = "crates/edge-index" } bobbin-ingest = { path = "crates/ingest" } bobbin-slingshot-client = { path = "crates/slingshot-client" } bobbin-record-lru = { path = "crates/record-lru" } +bobbin-knot-proxy = { path = "crates/knot-proxy" } bobbin-xrpc = { path = "crates/xrpc" } jacquard-common = "0.12.0-beta.2" diff --git a/crates/knot-proxy/Cargo.toml b/crates/knot-proxy/Cargo.toml new file mode 100644 index 0000000..027aa93 --- /dev/null +++ b/crates/knot-proxy/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "bobbin-knot-proxy" +version.workspace = true +edition.workspace = true +license.workspace = true +rust-version.workspace = true + +[dependencies] +bytes = { workspace = true } +futures = { workspace = true } +http = { workspace = true } +reqwest = { workspace = true } +scc = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true, features = ["net"] } +url = { workspace = true } + +[dev-dependencies] +tokio = { workspace = true, features = ["io-util", "macros", "net", "rt-multi-thread", "time"] } +wiremock = { workspace = true } diff --git a/crates/knot-proxy/src/breaker.rs b/crates/knot-proxy/src/breaker.rs new file mode 100644 index 0000000..1416d29 --- /dev/null +++ b/crates/knot-proxy/src/breaker.rs @@ -0,0 +1,239 @@ +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use thiserror::Error; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct FailureThreshold(u32); + +#[derive(Clone, Copy, Debug, Error)] +#[error("failure threshold must be at least 1")] +pub struct ThresholdError; + +impl FailureThreshold { + pub const fn new(n: u32) -> Result { + match n { + 0 => Err(ThresholdError), + other => Ok(Self(other)), + } + } + + pub const fn get(self) -> u32 { + self.0 + } +} + +#[derive(Clone, Copy, Debug)] +enum BreakerState { + Closed { failures: u32 }, + Open { until: Instant }, + HalfOpen, +} + +#[derive(Clone, Copy, Debug, Error)] +#[error("circuit breaker open")] +pub struct CircuitOpen; + +#[derive(Debug)] +pub struct Breaker { + state: Mutex, + threshold: FailureThreshold, + cooldown: Duration, +} + +impl Breaker { + pub fn new(threshold: FailureThreshold, cooldown: Duration) -> Self { + Self { + state: Mutex::new(BreakerState::Closed { failures: 0 }), + threshold, + cooldown, + } + } + + pub fn try_acquire(self: &Arc) -> Result { + self.try_acquire_at(Instant::now()) + } + + fn try_acquire_at(self: &Arc, now: Instant) -> Result { + let mut state = self.state.lock().expect("breaker mutex poisoned"); + match *state { + BreakerState::Closed { .. } => Ok(BreakerPermit::new(Arc::clone(self))), + BreakerState::Open { until } if now >= until => { + *state = BreakerState::HalfOpen; + Ok(BreakerPermit::new(Arc::clone(self))) + } + BreakerState::Open { .. } | BreakerState::HalfOpen => Err(CircuitOpen), + } + } + + pub fn record_success(&self) { + let mut state = self.state.lock().expect("breaker mutex poisoned"); + *state = BreakerState::Closed { failures: 0 }; + } + + pub fn record_failure(&self) { + self.record_failure_at(Instant::now()); + } + + fn record_failure_at(&self, now: Instant) { + let mut state = self.state.lock().expect("breaker mutex poisoned"); + let next = match *state { + BreakerState::Closed { failures } => { + let bumped = failures.saturating_add(1); + if bumped >= self.threshold.get() { + BreakerState::Open { + until: now + self.cooldown, + } + } else { + BreakerState::Closed { failures: bumped } + } + } + BreakerState::HalfOpen => BreakerState::Open { + until: now + self.cooldown, + }, + BreakerState::Open { .. } => *state, + }; + *state = next; + } +} + +#[must_use = "permit must outlive the upstream call so failures can be recorded"] +#[derive(Debug)] +pub struct BreakerPermit { + breaker: Arc, + resolved: bool, +} + +impl BreakerPermit { + fn new(breaker: Arc) -> Self { + Self { + breaker, + resolved: false, + } + } + + pub fn record_success(mut self) { + self.resolved = true; + self.breaker.record_success(); + } + + pub fn record_failure(mut self) { + self.resolved = true; + self.breaker.record_failure(); + } +} + +impl Drop for BreakerPermit { + fn drop(&mut self) { + if !self.resolved { + self.breaker.record_success(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn breaker(threshold: u32, cooldown_ms: u64) -> Arc { + Arc::new(Breaker::new( + FailureThreshold::new(threshold).unwrap(), + Duration::from_millis(cooldown_ms), + )) + } + + #[test] + fn closed_breaker_admits_all() { + let b = breaker(3, 100); + b.try_acquire().unwrap().record_success(); + b.try_acquire().unwrap().record_success(); + b.try_acquire().unwrap().record_success(); + } + + #[test] + fn opens_at_threshold() { + let b = breaker(2, 1_000); + b.record_failure(); + b.record_failure(); + assert!(b.try_acquire().is_err(), "must be open after 2 failures"); + } + + #[test] + fn success_resets_failure_count() { + let b = breaker(2, 1_000); + b.record_failure(); + b.record_success(); + b.record_failure(); + b.try_acquire() + .expect("successful run between failures must reset count") + .record_success(); + } + + #[test] + fn cooldown_admits_one_trial_only() { + let b = breaker(1, 50); + b.record_failure(); + assert!(b.try_acquire().is_err()); + std::thread::sleep(Duration::from_millis(60)); + let trial = b.try_acquire().expect("trial admitted after cooldown"); + assert!( + b.try_acquire().is_err(), + "second concurrent half-open call must be rejected", + ); + trial.record_success(); + } + + #[test] + fn half_open_failure_reopens() { + let b = breaker(1, 50); + b.record_failure(); + std::thread::sleep(Duration::from_millis(60)); + b.try_acquire().expect("trial admitted").record_failure(); + assert!( + b.try_acquire().is_err(), + "half-open failure must reopen breaker", + ); + } + + #[test] + fn half_open_success_closes() { + let b = breaker(1, 50); + b.record_failure(); + std::thread::sleep(Duration::from_millis(60)); + b.try_acquire().expect("trial admitted").record_success(); + b.try_acquire().expect("closed").record_success(); + b.try_acquire().expect("closed").record_success(); + } + + #[test] + fn open_failure_does_not_extend_cooldown_indefinitely() { + let b = breaker(1, 50); + b.record_failure(); + let mid = Instant::now(); + b.record_failure(); + b.try_acquire_at(mid + Duration::from_millis(60)) + .expect("second failure while open must not push cooldown out") + .record_success(); + } + + #[test] + fn dropped_permit_counts_as_success() { + let b = breaker(2, 1_000); + b.record_failure(); + drop(b.try_acquire().expect("admitted")); + b.try_acquire() + .expect("dropped permit must reset failure count") + .record_success(); + } + + #[test] + fn dropped_half_open_permit_closes_breaker() { + let b = breaker(1, 50); + b.record_failure(); + std::thread::sleep(Duration::from_millis(60)); + drop(b.try_acquire().expect("trial admitted")); + b.try_acquire() + .expect("dropped half-open permit must close breaker") + .record_success(); + } +} diff --git a/crates/knot-proxy/src/dns.rs b/crates/knot-proxy/src/dns.rs new file mode 100644 index 0000000..7fa5d30 --- /dev/null +++ b/crates/knot-proxy/src/dns.rs @@ -0,0 +1,112 @@ +use std::error::Error as StdError; +use std::fmt; +use std::io; +use std::net::SocketAddr; + +use reqwest::dns::{Addrs, Name, Resolve, Resolving}; + +use crate::host::{PrivateHostReason, classify_ip}; + +pub(crate) struct PrivateAddressFilter { + allow_private: bool, +} + +impl PrivateAddressFilter { + pub(crate) fn new(allow_private: bool) -> Self { + Self { allow_private } + } +} + +impl Resolve for PrivateAddressFilter { + fn resolve(&self, name: Name) -> Resolving { + let allow_private = self.allow_private; + let host = name.as_str().to_owned(); + Box::pin(async move { + let resolved: Vec = tokio::net::lookup_host((host.as_str(), 0)) + .await? + .collect(); + partition_safe(host, allow_private, resolved) + }) + } +} + +fn partition_safe( + host: String, + allow_private: bool, + resolved: Vec, +) -> Result> { + if !allow_private + && let Some(reason) = resolved.iter().find_map(|sa| classify_ip(&sa.ip())) + { + return Err(Box::new(BlockedAddressError { host, reason })); + } + if resolved.is_empty() { + return Err(Box::new(io::Error::other(format!( + "no resolvable addresses for {host}" + )))); + } + Ok(Box::new(resolved.into_iter())) +} + +#[derive(Debug)] +pub(crate) struct BlockedAddressError { + host: String, + reason: PrivateHostReason, +} + +impl fmt::Display for BlockedAddressError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "dns resolution for {} returned {} address", + self.host, self.reason, + ) + } +} + +impl StdError for BlockedAddressError {} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::{IpAddr, Ipv4Addr}; + + fn sa(ip: &str, port: u16) -> SocketAddr { + SocketAddr::new(ip.parse().unwrap(), port) + } + + #[test] + fn blocks_when_any_resolved_address_is_private_under_strict() { + let mixed = vec![sa("8.8.8.8", 0), sa("127.0.0.1", 0)]; + let res = partition_safe("mixed.example".into(), false, mixed); + assert!(res.is_err(), "any private address must fail strict resolve"); + } + + #[test] + fn allows_all_when_permissive() { + let mixed = vec![sa("8.8.8.8", 0), sa("127.0.0.1", 0)]; + let res = partition_safe("mixed.example".into(), true, mixed).expect("permissive"); + let collected: Vec = res.collect(); + assert_eq!(collected.len(), 2); + } + + #[test] + fn permits_public_only_resolution_under_strict() { + let public = vec![sa("8.8.8.8", 0), sa("1.1.1.1", 0)]; + let res = partition_safe("public.example".into(), false, public).expect("public"); + let collected: Vec = res.collect(); + assert_eq!(collected.len(), 2); + } + + #[test] + fn empty_resolution_is_an_error() { + let res = partition_safe("nx.example".into(), false, vec![]); + assert!(res.is_err(), "empty address list must surface an error"); + } + + #[test] + fn classify_ip_matches_url_classifier() { + assert!(classify_ip(&IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1))).is_some()); + assert!(classify_ip(&IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8))).is_none()); + } +} -- 2.51.2