Something went wrong. Try again.
A downstream fork of Tranquil PDS for use in margin.cafe
Something went wrong. Try again.
12 kB · 349 lines
Rust
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350use axum::http::HeaderName;use reqwest::{Client, ClientBuilder, Url};use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};use std::sync::{LazyLock, OnceLock};use std::time::Duration;use tracing::warn;use tranquil_types::{Did, Nsid, Rkey};
pub const DEFAULT_HEADERS_TIMEOUT: Duration = Duration::from_secs(10);pub const DEFAULT_BODY_TIMEOUT: Duration = Duration::from_secs(30);pub const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);pub const MAX_RESPONSE_SIZE: u64 = 10 * 1024 * 1024;
static PROXY_CLIENT: OnceLock<Client> = OnceLock::new();static DID_RESOLUTION_CLIENT: OnceLock<Client> = OnceLock::new();static HANDLE_RESOLUTION_CLIENT: OnceLock<Client> = OnceLock::new();
pub fn proxy_client() -> &'static Client { PROXY_CLIENT.get_or_init(|| { ClientBuilder::new() .timeout(DEFAULT_BODY_TIMEOUT) .connect_timeout(DEFAULT_CONNECT_TIMEOUT) .pool_max_idle_per_host(10) .pool_idle_timeout(Duration::from_secs(90)) .redirect(reqwest::redirect::Policy::none()) .build() .expect( "Failed to build HTTP client - this indicates a TLS or system configuration issue", ) })}
pub fn did_resolution_client() -> &'static Client { DID_RESOLUTION_CLIENT.get_or_init(|| { ClientBuilder::new() .timeout(Duration::from_secs(5)) .connect_timeout(DEFAULT_CONNECT_TIMEOUT) .pool_max_idle_per_host(10) .pool_idle_timeout(Duration::from_secs(90)) .redirect(reqwest::redirect::Policy::none()) .build() .expect( "Failed to build DID resolution client - this indicates a TLS or system configuration issue", ) })}
pub fn handle_resolution_client() -> &'static Client { HANDLE_RESOLUTION_CLIENT.get_or_init(|| { ClientBuilder::new() .timeout(Duration::from_secs(10)) .connect_timeout(DEFAULT_CONNECT_TIMEOUT) .pool_max_idle_per_host(10) .pool_idle_timeout(Duration::from_secs(90)) .redirect(reqwest::redirect::Policy::limited(5)) .build() .expect( "Failed to build handle resolution client - this indicates a TLS or system configuration issue", ) })}
pub async fn is_ssrf_safe(url: &str) -> Result<(), SsrfError> { let parsed = Url::parse(url).map_err(|_| SsrfError::InvalidUrl)?; let scheme = parsed.scheme(); if scheme != "https" { return Err(SsrfError::InsecureProtocol(scheme.to_string())); } if !parsed.username().is_empty() || parsed.password().is_some() { return Err(SsrfError::CredentialsNotAllowed); } let host = parsed.host_str().ok_or(SsrfError::NoHost)?; if let Ok(ip) = host.parse::<IpAddr>() { if !is_public_ip(&ip) { return Err(SsrfError::NonUnicastIp(ip.to_string())); } return Ok(()); } let port = parsed .port() .unwrap_or(if scheme == "https" { 443 } else { 80 }); let socket_addrs: Vec<std::net::SocketAddr> = match tokio::net::lookup_host((host, port)).await { Ok(addrs) => addrs.collect(), Err(_) => return Err(SsrfError::DnsResolutionFailed(host.to_string())), }; if socket_addrs.is_empty() { return Err(SsrfError::DnsResolutionFailed(host.to_string())); } if let Some(addr) = socket_addrs.iter().find(|addr| !is_public_ip(&addr.ip())) { warn!( "DNS resolution for {} returned non-public IP: {}", host, addr.ip() ); return Err(SsrfError::NonUnicastIp(addr.ip().to_string())); } Ok(())}
fn is_public_ip(ip: &IpAddr) -> bool { match ip { IpAddr::V4(v4) => is_public_ipv4(v4), IpAddr::V6(v6) => is_public_ipv6(v6), }}
fn is_public_ipv4(ip: &Ipv4Addr) -> bool { let [a, b, c, _] = ip.octets(); !(a == 0 || a == 10 || (a == 100 && (64..=127).contains(&b)) || a == 127 || (a == 169 && b == 254) || (a == 172 && (16..=31).contains(&b)) || (a == 192 && b == 0 && c == 0) || (a == 192 && b == 0 && c == 2) || (a == 192 && b == 88 && c == 99) || (a == 192 && b == 168) || (a == 198 && (b == 18 || b == 19)) || (a == 198 && b == 51 && c == 100) || (a == 203 && b == 0 && c == 113) || a >= 224)}
fn is_public_ipv6(ip: &Ipv6Addr) -> bool { let segments = ip.segments(); if segments[..6] == [0, 0, 0, 0, 0, 0xffff] { let high = segments[6].to_be_bytes(); let low = segments[7].to_be_bytes(); return is_public_ipv4(&Ipv4Addr::new(high[0], high[1], low[0], low[1])); }
(segments[0] & 0xe000) == 0x2000 && !(segments[0] == 0x2001 && segments[1] < 0x0200) && !(segments[0] == 0x2001 && segments[1] == 0x0db8) && segments[0] != 0x2002 && !(segments[0] == 0x3fff && segments[1] < 0x1000)}
#[derive(Debug, Clone)]pub enum SsrfError { InvalidUrl, InsecureProtocol(String), CredentialsNotAllowed, NoHost, NonUnicastIp(String), DnsResolutionFailed(String),}
impl std::fmt::Display for SsrfError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { SsrfError::InvalidUrl => write!(f, "Invalid URL"), SsrfError::InsecureProtocol(p) => write!(f, "Insecure protocol: {}", p), SsrfError::CredentialsNotAllowed => write!(f, "URL credentials are not allowed"), SsrfError::NoHost => write!(f, "No host in URL"), SsrfError::NonUnicastIp(ip) => write!(f, "Non-public IP address: {}", ip), SsrfError::DnsResolutionFailed(host) => { write!(f, "DNS resolution failed for: {}", host) } } }}
impl std::error::Error for SsrfError {}// TODO: update to match https://github.com/bluesky-social/atproto/blob/main/packages/pds/src/pipethrough.ts// Currently spec says nothing at all!! about forwarding headers during proxying and https://github.com/bluesky-social/atproto/discussions/2350#[cfg(feature = "bsky-support")]pub static HEADERS_TO_FORWARD: LazyLock<[HeaderName; 4]> = LazyLock::new(|| { [ HeaderName::from_static("accept-language"), crate::util::HEADER_ATPROTO_ACCEPT_LABELERS, crate::util::HEADER_X_BSKY_TOPICS, http::header::CONTENT_TYPE, ]});#[cfg(not(feature = "bsky-support"))]pub static HEADERS_TO_FORWARD: LazyLock<[HeaderName; 3]> = LazyLock::new(|| { [ HeaderName::from_static("accept-language"), crate::util::HEADER_ATPROTO_ACCEPT_LABELERS, http::header::CONTENT_TYPE, ]});pub static RESPONSE_HEADERS_TO_FORWARD: LazyLock<[HeaderName; 6]> = LazyLock::new(|| { [ crate::util::HEADER_ATPROTO_REPO_REV, crate::util::HEADER_ATPROTO_CONTENT_LABELERS, HeaderName::from_static("retry-after"), http::header::CONTENT_TYPE, http::header::CACHE_CONTROL, http::header::ETAG, ]});
pub fn validate_at_uri(uri: &str) -> Result<AtUriParts, &'static str> { if !uri.starts_with("at://") { return Err("URI must start with at://"); } let path = uri.trim_start_matches("at://"); let parts: Vec<&str> = path.split('/').collect(); if parts.is_empty() { return Err("URI missing DID"); } let did: Did = parts[0].parse().map_err(|_| "Invalid DID in URI")?; let collection = parts .get(1) .map(|s| s.parse::<Nsid>()) .transpose() .map_err(|_| "Invalid collection NSID")?; let rkey = parts .get(2) .map(|s| s.parse::<Rkey>()) .transpose() .map_err(|_| "Invalid rkey")?; Ok(AtUriParts { did, collection, rkey, })}
#[derive(Debug, Clone)]pub struct AtUriParts { pub did: Did, pub collection: Option<Nsid>, pub rkey: Option<Rkey>,}
pub fn validate_limit(limit: Option<u32>, default: u32, max: u32) -> u32 { match limit { Some(0) => default, Some(l) if l > max => max, Some(l) => l, None => default, }}
#[cfg(test)]mod tests { use super::*;
#[tokio::test] async fn ssrf_allows_public_https() { assert!(is_ssrf_safe("https://1.1.1.1/xrpc/test").await.is_ok()); }
#[tokio::test] async fn ssrf_blocks_http_by_default() { let result = is_ssrf_safe("http://93.184.216.34/xrpc/test").await; assert!(matches!(result, Err(SsrfError::InsecureProtocol(_)))); }
#[tokio::test] async fn ssrf_blocks_non_http_schemes() { assert!(matches!( is_ssrf_safe("file:///etc/passwd").await, Err(SsrfError::InsecureProtocol(_)) )); assert!(matches!( is_ssrf_safe("ftp://1.1.1.1/test").await, Err(SsrfError::InsecureProtocol(_)) )); }
#[tokio::test] async fn ssrf_blocks_url_credentials() { assert!(matches!( is_ssrf_safe("https://user:password@1.1.1.1/test").await, Err(SsrfError::CredentialsNotAllowed) )); }
#[tokio::test] async fn ssrf_blocks_non_public_ipv4_addresses() { for ip in [ "0.0.0.0", "10.0.0.1", "100.64.0.1", "127.0.0.1", "169.254.169.254", "172.16.0.1", "192.168.0.1", "192.0.2.1", "198.18.0.1", "198.51.100.1", "203.0.113.1", "224.0.0.1", "240.0.0.1", "255.255.255.255", ] { assert!( matches!( is_ssrf_safe(&format!("https://{ip}/test")).await, Err(SsrfError::NonUnicastIp(_)) ), "expected {ip} to be rejected" ); } }
#[tokio::test] async fn ssrf_blocks_non_public_ipv6_addresses() { for ip in [ "::", "::1", "fe80::1", "fc00::1", "fd00::1", "ff02::1", "2001:db8::1", "::ffff:10.0.0.1", "::ffff:169.254.169.254", ] { assert!( matches!( is_ssrf_safe(&format!("https://[{ip}]/test")).await, Err(SsrfError::NonUnicastIp(_)) ), "expected {ip} to be rejected" ); } } #[test] fn test_validate_at_uri() { let result = validate_at_uri("at://did:plc:test/app.bsky.feed.post/abc123"); assert!(result.is_ok()); let parts = result.unwrap(); assert_eq!(parts.did, "did:plc:test".parse::<Did>().unwrap()); assert_eq!( parts.collection, Some("app.bsky.feed.post".parse::<Nsid>().unwrap()) ); assert_eq!(parts.rkey, Some("abc123".parse::<Rkey>().unwrap())); } #[test] fn test_validate_at_uri_invalid() { assert!(validate_at_uri("https://example.com").is_err()); assert!(validate_at_uri("at://notadid/collection/rkey").is_err()); } #[test] fn test_validate_limit() { assert_eq!(validate_limit(None, 50, 100), 50); assert_eq!(validate_limit(Some(0), 50, 100), 50); assert_eq!(validate_limit(Some(200), 50, 100), 100); assert_eq!(validate_limit(Some(75), 50, 100), 75); }}