From c663df1d55b8fbad51035dfbd944aaabef4de646 Mon Sep 17 00:00:00 2001 From: Lewis Date: Mon, 04 May 2026 15:29:01 +0000 Subject: [PATCH] feat(knot-proxy): proxy core, dispatch, header forwarding Lewis: May this revision serve well! --- crates/knot-proxy/src/lib.rs | 670 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file(s) changed, 670 insertion(s)(+), 0 deletion(s)(-) diff --git a/crates/knot-proxy/src/lib.rs b/crates/knot-proxy/src/lib.rs new file mode 100644 --- /dev/null +++ b/crates/knot-proxy/src/lib.rs @@ -0,0 +1,670 @@ +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; +use std::time::Duration; + +use bytes::Bytes; +use futures::Stream; +use http::HeaderMap; +use reqwest::{Client, Response, StatusCode, redirect::Policy}; +use scc::HashMap as SccMap; +use thiserror::Error; +use url::Url; + +mod breaker; +mod dns; +mod host; + +pub use breaker::{Breaker, BreakerPermit, CircuitOpen, FailureThreshold, ThresholdError}; +pub use host::{KnotHost, KnotHostError, PrivateHostReason, RepoSlug, RepoSlugError}; + +const USER_AGENT: &str = concat!("bobbin/", env!("CARGO_PKG_VERSION")); +const HTTPS_SCHEME: &str = "https"; + +#[derive(Clone, Debug)] +pub struct KnotProxyConfig { + pub failure_threshold: FailureThreshold, + pub cooldown: Duration, + pub connect_timeout: Duration, + pub read_timeout: Duration, + pub allow_private_hosts: bool, + pub require_https: bool, +} + +impl Default for KnotProxyConfig { + fn default() -> Self { + Self { + failure_threshold: FailureThreshold::new(5).expect("nonzero literal"), + cooldown: Duration::from_secs(30), + connect_timeout: Duration::from_secs(5), + read_timeout: Duration::from_secs(60), + allow_private_hosts: false, + require_https: true, + } + } +} + +#[derive(Debug, Error)] +pub enum KnotProxyError { + #[error("circuit breaker open")] + CircuitOpen, + #[error("blocked: host {host} resolves to {reason} address space")] + BlockedHost { + host: String, + reason: PrivateHostReason, + }, + #[error("blocked: knot {host} requires https, got plaintext http")] + PlaintextHttp { host: String }, + #[error("connect failed: {0}")] + Connect(String), + #[error("upstream read timed out: {0}")] + Timeout(String), + #[error("redirect refused: {0}")] + Redirect(String), + #[error("transport: {0}")] + Transport(String), + #[error("upstream returned status {0}")] + Upstream(StatusCode), +} + +pub struct KnotProxy { + http: Client, + breakers: SccMap>, + threshold: FailureThreshold, + cooldown: Duration, + allow_private_hosts: bool, + require_https: bool, +} + +impl KnotProxy { + pub fn new(config: KnotProxyConfig) -> Result { + let resolver = Arc::new(dns::PrivateAddressFilter::new(config.allow_private_hosts)); + let http = Client::builder() + .user_agent(USER_AGENT) + .connect_timeout(config.connect_timeout) + .read_timeout(config.read_timeout) + .redirect(Policy::none()) + .no_gzip() + .no_brotli() + .no_deflate() + .dns_resolver(resolver) + .build()?; + Ok(Self { + http, + breakers: SccMap::new(), + threshold: config.failure_threshold, + cooldown: config.cooldown, + allow_private_hosts: config.allow_private_hosts, + require_https: config.require_https, + }) + } + + pub fn allows_private_hosts(&self) -> bool { + self.allow_private_hosts + } + + pub fn requires_https(&self) -> bool { + self.require_https + } + + pub async fn forward( + &self, + host: &KnotHost, + nsid: &str, + query: &[(&str, &str)], + headers: HeaderMap, + ) -> Result { + self.guard_host(host)?; + let breaker = self.breaker_for(host).await; + let permit = breaker + .try_acquire() + .map_err(|_: CircuitOpen| KnotProxyError::CircuitOpen)?; + let url = build_xrpc_url(host, nsid, query); + let outcome = self.http.get(url).headers(headers).send().await; + classify(outcome, permit) + } + + fn guard_host(&self, host: &KnotHost) -> Result<(), KnotProxyError> { + let host_str = || host.url().host_str().unwrap_or_default().to_owned(); + if self.require_https && host.url().scheme() != HTTPS_SCHEME { + return Err(KnotProxyError::PlaintextHttp { host: host_str() }); + } + if self.allow_private_hosts { + return Ok(()); + } + match host.private_literal_reason() { + None => Ok(()), + Some(reason) => Err(KnotProxyError::BlockedHost { + host: host_str(), + reason, + }), + } + } + + async fn breaker_for(&self, host: &KnotHost) -> Arc { + if let Some(existing) = self.breakers.read_async(host, |_, v| Arc::clone(v)).await { + return existing; + } + let entry = self.breakers.entry_async(host.clone()).await; + Arc::clone( + entry + .or_insert_with(|| Arc::new(Breaker::new(self.threshold, self.cooldown))) + .get(), + ) + } +} + +#[derive(Debug)] +pub struct ProxyResponse { + inner: Response, + permit: BreakerPermit, +} + +impl ProxyResponse { + pub fn status(&self) -> StatusCode { + self.inner.status() + } + + pub fn headers(&self) -> &HeaderMap { + self.inner.headers() + } + + pub fn into_body_stream(self) -> BodyStream { + BodyStream::new(self.inner, self.permit) + } +} + +type ChunkStream = Pin> + Send>>; + +pub struct BodyStream { + inner: ChunkStream, + permit: Option, +} + +impl BodyStream { + fn new(response: Response, permit: BreakerPermit) -> Self { + Self { + inner: Box::pin(response.bytes_stream()), + permit: Some(permit), + } + } + + fn resolve(&mut self, success: bool) { + if let Some(permit) = self.permit.take() { + if success { + permit.record_success(); + } else { + permit.record_failure(); + } + } + } +} + +impl Stream for BodyStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let next = match self.inner.as_mut().poll_next(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(item) => item, + }; + match &next { + Some(Ok(_)) => {} + Some(Err(_)) => self.resolve(false), + None => self.resolve(true), + } + Poll::Ready(next) + } +} + +fn build_xrpc_url(host: &KnotHost, nsid: &str, query: &[(&str, &str)]) -> Url { + let mut url = host.xrpc_url(nsid); + { + let mut pairs = url.query_pairs_mut(); + query.iter().for_each(|(k, v)| { + pairs.append_pair(k, v); + }); + } + url +} + +fn classify( + outcome: Result, + permit: BreakerPermit, +) -> Result { + match outcome { + Ok(resp) if is_upstream_failure(resp.status()) => { + let status = resp.status(); + permit.record_failure(); + Err(KnotProxyError::Upstream(status)) + } + Ok(resp) => Ok(ProxyResponse { + inner: resp, + permit, + }), + Err(err) => { + permit.record_failure(); + Err(map_transport(err)) + } + } +} + +fn is_upstream_failure(status: StatusCode) -> bool { + status.is_server_error() || is_unfollowable_redirect(status) +} + +fn is_unfollowable_redirect(status: StatusCode) -> bool { + matches!(status.as_u16(), 301 | 302 | 303 | 307 | 308) +} + +fn map_transport(err: reqwest::Error) -> KnotProxyError { + let msg = err.to_string(); + if err.is_timeout() { + KnotProxyError::Timeout(msg) + } else if err.is_connect() { + KnotProxyError::Connect(msg) + } else if err.is_redirect() { + KnotProxyError::Redirect(msg) + } else { + KnotProxyError::Transport(msg) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::stream::TryStreamExt; + use tokio::io::AsyncWriteExt; + use wiremock::matchers::{method, path, query_param}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + pub(crate) fn config_for_test() -> KnotProxyConfig { + KnotProxyConfig { + failure_threshold: FailureThreshold::new(2).unwrap(), + cooldown: Duration::from_millis(80), + connect_timeout: Duration::from_millis(500), + read_timeout: Duration::from_secs(2), + allow_private_hosts: true, + require_https: false, + } + } + + async fn server() -> MockServer { + MockServer::start().await + } + + fn host_of(server: &MockServer) -> KnotHost { + KnotHost::parse(&server.uri()).unwrap() + } + + pub(crate) async fn drain(stream: BodyStream) -> Result { + let chunks: Vec = stream.try_collect().await?; + let total: usize = chunks.iter().map(|b| b.len()).sum(); + let mut buf = bytes::BytesMut::with_capacity(total); + chunks.iter().for_each(|c| buf.extend_from_slice(c)); + Ok(buf.freeze()) + } + + #[tokio::test] + async fn forwards_query_params_and_returns_body() { + let server = server().await; + Mock::given(method("GET")) + .and(path("/xrpc/sh.tangled.repo.blob")) + .and(query_param("repo", "did:plc:abalone/barnacle")) + .and(query_param("ref", "main")) + .and(query_param("path", "README.md")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .set_body_string(r#"{"path":"README.md"}"#), + ) + .mount(&server) + .await; + + let proxy = KnotProxy::new(config_for_test()).unwrap(); + let resp = proxy + .forward( + &host_of(&server), + "sh.tangled.repo.blob", + &[ + ("repo", "did:plc:abalone/barnacle"), + ("ref", "main"), + ("path", "README.md"), + ], + HeaderMap::new(), + ) + .await + .expect("happy path"); + assert_eq!(resp.status(), 200); + let body = drain(resp.into_body_stream()).await.unwrap(); + assert_eq!(&body[..], br#"{"path":"README.md"}"#); + } + + #[tokio::test] + async fn five_hundreds_open_breaker() { + let server = server().await; + Mock::given(method("GET")) + .and(path("/xrpc/sh.tangled.repo.blob")) + .respond_with(ResponseTemplate::new(503)) + .mount(&server) + .await; + + let proxy = KnotProxy::new(config_for_test()).unwrap(); + let host = host_of(&server); + let r1 = proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + assert!(matches!(r1, Err(KnotProxyError::Upstream(_)))); + let r2 = proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + assert!(matches!(r2, Err(KnotProxyError::Upstream(_)))); + let r3 = proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + assert!(matches!(r3, Err(KnotProxyError::CircuitOpen))); + } + + #[tokio::test] + async fn four_hundreds_do_not_open_breaker() { + let server = server().await; + Mock::given(method("GET")) + .and(path("/xrpc/sh.tangled.repo.blob")) + .respond_with(ResponseTemplate::new(404).set_body_string("not found")) + .mount(&server) + .await; + + let proxy = KnotProxy::new(config_for_test()).unwrap(); + let host = host_of(&server); + let r1 = proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + assert_eq!(r1.unwrap().status(), 404); + let r2 = proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + assert_eq!(r2.unwrap().status(), 404); + let r3 = proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + assert_eq!( + r3.unwrap().status(), + 404, + "client errors must not trip breaker", + ); + } + + #[tokio::test] + async fn breaker_recovers_after_cooldown() { + let server = server().await; + Mock::given(method("GET")) + .and(path("/xrpc/sh.tangled.repo.blob")) + .respond_with(ResponseTemplate::new(503)) + .up_to_n_times(2) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/xrpc/sh.tangled.repo.blob")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .set_body_string("ok"), + ) + .mount(&server) + .await; + + let proxy = KnotProxy::new(config_for_test()).unwrap(); + let host = host_of(&server); + let _ = proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + let _ = proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + assert!(matches!( + proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await, + Err(KnotProxyError::CircuitOpen), + )); + tokio::time::sleep(Duration::from_millis(120)).await; + let recovered = proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await + .expect("must recover after cooldown"); + assert_eq!(recovered.status(), 200); + let body = drain(recovered.into_body_stream()).await.unwrap(); + assert_eq!(&body[..], b"ok"); + } + + #[tokio::test] + async fn breakers_are_isolated_per_host() { + let bad = server().await; + let good = server().await; + Mock::given(method("GET")) + .and(path("/xrpc/sh.tangled.repo.blob")) + .respond_with(ResponseTemplate::new(503)) + .mount(&bad) + .await; + Mock::given(method("GET")) + .and(path("/xrpc/sh.tangled.repo.blob")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .set_body_string("ok"), + ) + .mount(&good) + .await; + + let proxy = KnotProxy::new(config_for_test()).unwrap(); + let bad_host = host_of(&bad); + let good_host = host_of(&good); + let _ = proxy + .forward(&bad_host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + let _ = proxy + .forward(&bad_host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + assert!(matches!( + proxy + .forward(&bad_host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await, + Err(KnotProxyError::CircuitOpen), + )); + let resp = proxy + .forward(&good_host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await + .expect("healthy host stays open"); + assert_eq!(resp.status(), 200); + } + + #[tokio::test] + async fn build_xrpc_url_appends_query() { + let host = KnotHost::parse("https://oyster.cafe").unwrap(); + let url = build_xrpc_url( + &host, + "sh.tangled.repo.tree", + &[("repo", "did:plc:abalone/barnacle"), ("ref", "main")], + ); + assert_eq!( + url.as_str(), + "https://oyster.cafe/xrpc/sh.tangled.repo.tree?repo=did%3Aplc%3Aabalone%2Fbarnacle&ref=main", + ); + } + + #[tokio::test] + async fn rejects_private_host_by_default() { + let server = server().await; + let strict = KnotProxyConfig { + allow_private_hosts: false, + ..config_for_test() + }; + let proxy = KnotProxy::new(strict).unwrap(); + let err = proxy + .forward( + &host_of(&server), + "sh.tangled.repo.blob", + &[], + HeaderMap::new(), + ) + .await + .expect_err("loopback must be blocked under strict config"); + assert!(matches!(err, KnotProxyError::BlockedHost { .. })); + } + + #[tokio::test] + async fn rejects_plaintext_when_https_required() { + let server = server().await; + let strict = KnotProxyConfig { + require_https: true, + ..config_for_test() + }; + let proxy = KnotProxy::new(strict).unwrap(); + let err = proxy + .forward( + &host_of(&server), + "sh.tangled.repo.blob", + &[], + HeaderMap::new(), + ) + .await + .expect_err("plaintext http must be rejected under https-required"); + assert!( + matches!(err, KnotProxyError::PlaintextHttp { .. }), + "got {err:?}", + ); + } + + #[tokio::test] + async fn transport_error_trips_breaker() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + drop(listener); + let dead = KnotHost::parse(&format!("http://{addr}")).unwrap(); + + let proxy = KnotProxy::new(config_for_test()).unwrap(); + let r1 = proxy + .forward(&dead, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + assert!( + r1.is_err(), + "transport must fail against closed port: {r1:?}" + ); + let r2 = proxy + .forward(&dead, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + assert!(r2.is_err(), "second transport must fail: {r2:?}"); + let r3 = proxy + .forward(&dead, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + assert!( + matches!(r3, Err(KnotProxyError::CircuitOpen)), + "transport failures must trip breaker, got {r3:?}", + ); + } + + #[tokio::test] + async fn redirects_surface_as_upstream_failure() { + let primary = server().await; + let secondary = server().await; + Mock::given(method("GET")) + .and(path("/xrpc/sh.tangled.repo.blob")) + .respond_with( + ResponseTemplate::new(302) + .insert_header("location", &format!("{}/secret", secondary.uri())), + ) + .mount(&primary) + .await; + Mock::given(method("GET")) + .and(path("/secret")) + .respond_with(ResponseTemplate::new(200).set_body_string("leaked")) + .mount(&secondary) + .await; + + let proxy = KnotProxy::new(config_for_test()).unwrap(); + let err = proxy + .forward( + &host_of(&primary), + "sh.tangled.repo.blob", + &[], + HeaderMap::new(), + ) + .await + .expect_err("302 must surface as upstream failure"); + assert!( + matches!(err, KnotProxyError::Upstream(s) if s.as_u16() == 302), + "got {err:?}", + ); + let received = secondary.received_requests().await.unwrap(); + assert!(received.is_empty(), "secondary must never be dialled"); + } + + #[tokio::test] + async fn not_modified_passes_through() { + let server = server().await; + Mock::given(method("GET")) + .and(path("/xrpc/sh.tangled.repo.blob")) + .respond_with(ResponseTemplate::new(304).insert_header("etag", "\"v1\"")) + .mount(&server) + .await; + let proxy = KnotProxy::new(config_for_test()).unwrap(); + let resp = proxy + .forward( + &host_of(&server), + "sh.tangled.repo.blob", + &[], + HeaderMap::new(), + ) + .await + .expect("304 is a cache validator, not a redirect"); + assert_eq!(resp.status(), 304); + } + + #[tokio::test] + async fn mid_stream_drop_records_breaker_failure() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + async fn drop_after_partial(mut socket: tokio::net::TcpStream) { + let _ = socket + .write_all( + b"HTTP/1.1 200 OK\r\nContent-Length: 1024\r\nContent-Type: application/octet-stream\r\n\r\nabcd", + ) + .await; + drop(socket); + } + let admit = || async { + let (socket, _) = listener.accept().await.ok()?; + drop_after_partial(socket).await; + Some(()) + }; + admit().await; + admit().await; + }); + + let host = KnotHost::parse(&format!("http://{addr}")).unwrap(); + let proxy = KnotProxy::new(config_for_test()).unwrap(); + + let r1 = proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await + .expect("headers arrive even when body is truncated"); + assert_eq!(r1.status(), 200); + let _ = drain(r1.into_body_stream()).await; + + let r2 = proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await + .expect("second call still gets headers"); + let _ = drain(r2.into_body_stream()).await; + + let r3 = proxy + .forward(&host, "sh.tangled.repo.blob", &[], HeaderMap::new()) + .await; + assert!( + matches!(r3, Err(KnotProxyError::CircuitOpen)), + "two truncated streams must trip the breaker, got {r3:?}", + ); + server.abort(); + } +} -- tangled.sh