Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405#![allow(dead_code)]
use std::{hash::Hash, time::Duration};
use jacquard_common::{deps::fluent_uri, types::string::Handle};use rand::RngExt;use reqwest::StatusCode;use serde::{Deserialize, Deserializer, Serializer, de::Error as _};use tokio::sync::watch;use tracing::info;use url::Url;
use crate::{db::types::DidKey, types::RepoStatus};
pub mod throttle;
#[allow(dead_code)]/// checks if the error contains a hyper / std io timeout errorpub fn is_timeout(err: &dyn std::error::Error) -> bool { let mut source = err.source();
while let Some(err) = source { if let Some(hyper_err) = err.downcast_ref::<hyper::Error>() && hyper_err.is_timeout() { return true; } if let Some(io) = err.downcast_ref::<std::io::Error>() && io.kind() == std::io::ErrorKind::TimedOut { return true; } source = err.source(); }
false}
pub fn is_tls_cert_error(io_err: &std::io::Error) -> bool { let Some(inner) = io_err.get_ref() else { return false; }; if let Some(rustls_err) = inner.downcast_ref::<rustls::Error>() { return is_tls_error_their_fault(rustls_err); } if let Some(nested_io) = inner.downcast_ref::<std::io::Error>() { return is_tls_cert_error(nested_io); } false}
pub fn is_io_error_their_fault(e: &std::io::Error) -> bool { use std::io::ErrorKind::*; matches!( e.kind(), // some of these maybe our fault, but lets assume we have working networking // if its our fault chances are most of the other hosts will also fail, which will be in the logs // we log the error anyway so it should be easy to tell if something is going bad ConnectionRefused | HostUnreachable | NetworkUnreachable | ConnectionReset | ConnectionAborted | TimedOut | UnexpectedEof )}
pub fn is_tls_error_their_fault(e: &rustls::Error) -> bool { use rustls::Error::*;
matches!( *e, InvalidCertificate(_) | PeerMisbehaved(_) | InconsistentKeys(_) | InappropriateMessage { .. } | InappropriateHandshakeMessage { .. } | InvalidMessage(_) | NoCertificatesPresented | UnsupportedNameType | DecryptError | AlertReceived(_) | PeerIncompatible(_) | InvalidCertRevocationList(_) | InvalidEncryptedClientHello(_) | PeerSentOversizedRecord | NoApplicationProtocol // this is not exhaustive, so remember to look at rustls::Error on version changes )}
// use this for public (unauth) xrpc errorspub fn is_status_their_fault(status: u16) -> bool { (100..200).contains(&status) // informational, why are we here? || (500..600).contains(&status) // server error :> || (300..400).contains(&status) // any 3xx error doesnt make sense in the context of a pds / relay || matches!( status, 404 // NOT FOUND: we know its not our fault because we use known xrpcs.. | 436 // some stupid ass error code idk, some domain park uses this i think??? | 403 // FORBIDDEN: sob | 401 // UNAUTHORIZED: sob // note: 410 GONE is intentionally excluded — it has explicit application-level meaning // (repo deleted/tombstoned) and is handled as NoSignal in the crawler. including it // here would throttle the whole PDS for a single deleted repo. )}
/// outcome of [`RetryWithBackoff::retry`] when the operation does not succeed.pub enum RetryOutcome<E> { /// ratelimited after exhausting all retries Ratelimited, /// non-ratelimit failure, carrying the last error Failed(E),}
/// extension trait that adds `.retry()` to async `FnMut` closures.////// `on_ratelimit` receives the error and current attempt number./// returning `Some(duration)` signals a transient failure and provides the backoff;/// returning `None` signals a terminal failure.pub trait RetryWithBackoff<T, E, Fut>: FnMut() -> Futwhere Fut: Future<Output = Result<T, E>>,{ #[allow(async_fn_in_trait)] async fn retry( &mut self, max_retries: u32, on_ratelimit: impl Fn(&E, u32) -> Option<Duration>, ) -> Result<T, RetryOutcome<E>> { let mut attempt = 0u32; loop { match self().await { Ok(val) => return Ok(val), Err(e) => match on_ratelimit(&e, attempt) { Some(_) if attempt >= max_retries => return Err(RetryOutcome::Ratelimited), Some(backoff) => { // jitter the backoff let backoff = rand::rng().random_range((backoff / 2)..backoff); tokio::time::sleep(backoff).await; attempt += 1; } None => return Err(RetryOutcome::Failed(e)), }, } } }}
impl<T, E, F, Fut> RetryWithBackoff<T, E, Fut> for Fwhere F: FnMut() -> Fut, Fut: Future<Output = Result<T, E>>,{}
/// extension trait that adds `.wait_enabled()` to `watch::Receiver<bool>`.////// waits until the value becomes `true`, logging once when paused and once when resumed.pub trait WatchEnabledExt { #[allow(async_fn_in_trait)] async fn wait_enabled(&mut self, component: &'static str);}
impl WatchEnabledExt for watch::Receiver<bool> { async fn wait_enabled(&mut self, component: &'static str) { if !*self.borrow() { info!("{component} paused"); while !*self.borrow() { let _ = self.changed().await; } info!("{component} resumed"); } }}
/// extension trait that adds `.error_for_status()` to futures returning a reqwest `Response`.pub trait ErrorForStatus: Future<Output = Result<reqwest::Response, reqwest::Error>> { fn error_for_status(self) -> impl Future<Output = Result<reqwest::Response, reqwest::Error>> where Self: Sized, { futures::FutureExt::map(self, |r| r.and_then(|r| r.error_for_status())) }}
impl<F: Future<Output = Result<reqwest::Response, reqwest::Error>>> ErrorForStatus for F {}
/// extracts a retry delay in seconds from rate limit response headers.////// checks in priority order:/// - `retry-after: <seconds>` (relative)/// - `ratelimit-reset: <unix timestamp>` (absolute) (ref pds sends this)pub fn parse_retry_after(resp: &reqwest::Response) -> Option<u64> { let headers = resp.headers();
let retry_after = headers .get(reqwest::header::RETRY_AFTER) .and_then(|v| v.to_str().ok()) .and_then(|s| s.parse::<u64>().ok());
let rate_limit_reset = headers .get("ratelimit-reset") .and_then(|v| v.to_str().ok()) .and_then(|s| s.parse::<i64>().ok()) .map(|ts| { let now = chrono::Utc::now().timestamp(); (ts - now).max(1) as u64 });
retry_after.or(rate_limit_reset)}
// cloudflare-specific status codespub const CONNECTION_TIMEOUT: StatusCode = unsafe { match StatusCode::from_u16(522) { Ok(s) => s, _ => std::hint::unreachable_unchecked(), }};pub const SITE_FROZEN: StatusCode = unsafe { match StatusCode::from_u16(530) { Ok(s) => s, _ => std::hint::unreachable_unchecked(), }};pub const SSL_HANDSHAKE_FAILURE: StatusCode = unsafe { match StatusCode::from_u16(525) { Ok(s) => s, _ => std::hint::unreachable_unchecked(), }};
pub fn ser_status_code<S: Serializer>(s: &Option<StatusCode>, ser: S) -> Result<S::Ok, S::Error> { match s { Some(code) => ser.serialize_some(&code.as_u16()), None => ser.serialize_none(), }}
pub fn deser_status_code<'de, D: Deserializer<'de>>( deser: D,) -> Result<Option<StatusCode>, D::Error> { Option::<u16>::deserialize(deser)? .map(StatusCode::from_u16) .transpose() .map_err(serde::de::Error::custom)}
pub fn opt_cid_serialize_str<S: Serializer>(v: &Option<cid::Cid>, s: S) -> Result<S::Ok, S::Error> { match v { Some(cid) => s.serialize_some(cid.to_string().as_str()), None => s.serialize_none(), }}
pub fn did_key_serialize_str<S: Serializer>(v: &DidKey<'_>, s: S) -> Result<S::Ok, S::Error> { s.serialize_str(&v.encode())}
pub fn did_key_deserialize_str<'de, D: Deserializer<'de>>( d: D,) -> Result<DidKey<'static>, D::Error> { let s = <&str>::deserialize(d)?; if let Ok(k) = DidKey::from_did_key(s) { return Ok(k); } let (_base, bytes) = multibase::decode(s).map_err(|err| D::Error::custom(err.to_string()))?; Ok(DidKey(std::borrow::Cow::Owned(bytes)))}
pub fn opt_did_key_serialize_str<S: Serializer>( v: &Option<DidKey<'_>>, s: S,) -> Result<S::Ok, S::Error> { match v { Some(k) => s.serialize_some(k.encode().as_str()), None => s.serialize_none(), }}
pub fn repo_status_serialize_str<S: Serializer>(v: &RepoStatus, s: S) -> Result<S::Ok, S::Error> { s.serialize_str(&v.to_string())}
pub fn url_to_fluent_uri_ref(url: &Url) -> fluent_uri::Uri<&str> { fluent_uri::Uri::parse(url.as_str()).expect("that url is validated")}
pub(crate) fn invalid_handle() -> Handle { unsafe { Handle::unchecked(smol_str::SmolStr::new_static("handle.invalid")) }}
/// returns hash of value using ahashpub fn hash<T: Hash>(val: &T) -> u64 { use std::hash::Hasher; let mut hasher = ahash::AHasher::default(); val.hash(&mut hasher); hasher.finish()}
pub async fn read_limited_bytes( resp: reqwest::Response, limit: usize,) -> miette::Result<bytes::Bytes> { use bytes::BytesMut; use futures::StreamExt as _; use miette::Context as _; use miette::IntoDiagnostic as _; let mut stream = resp.bytes_stream(); let mut buf = BytesMut::new(); while let Some(chunk) = stream.next().await { let chunk = chunk .into_diagnostic() .context("failed to read response chunk")?; if buf.len() + chunk.len() > limit { miette::bail!("response body too large (exceeds {limit} bytes)"); } buf.extend_from_slice(&chunk); } Ok(buf.freeze())}
#[cfg(test)]mod tests { use super::*; use serde::Serialize;
#[derive(Deserialize, Serialize)] struct StatusFixture { #[serde( serialize_with = "ser_status_code", deserialize_with = "deser_status_code" )] status: Option<StatusCode>, }
#[derive(Serialize)] struct OutputFixture<'a> { #[serde(serialize_with = "opt_cid_serialize_str")] cid: &'a Option<cid::Cid>, #[serde(serialize_with = "did_key_serialize_str")] did_key: &'a DidKey<'a>, #[serde(serialize_with = "opt_did_key_serialize_str")] optional_did_key: &'a Option<DidKey<'a>>, #[serde(serialize_with = "repo_status_serialize_str")] status: &'a RepoStatus, }
#[test] fn http_status_serde_contract() { for status in [None, Some(StatusCode::OK), Some(CONNECTION_TIMEOUT)] { let encoded = serde_json::to_vec(&StatusFixture { status }).unwrap(); let decoded: StatusFixture = serde_json::from_slice(&encoded).unwrap(); assert_eq!(decoded.status, status); } assert!(serde_json::from_str::<StatusFixture>(r#"{"status":99}"#).is_err()); assert!(serde_json::from_str::<StatusFixture>(r#"{"status":1000}"#).is_err()); }
#[test] fn xrpc_output_string_serializers_are_stable() { let cid = Some(jacquard_repo::mst::util::compute_cid(b"body").unwrap()); let did_key = DidKey(std::borrow::Cow::Borrowed(&[0xed, 1, 2, 3])); let optional_did_key = Some(did_key.clone()); let encoded = serde_json::to_value(OutputFixture { cid: &cid, did_key: &did_key, optional_did_key: &optional_did_key, status: &RepoStatus::Throttled, }) .unwrap();
assert_eq!(encoded["cid"], cid.unwrap().to_string()); assert_eq!(encoded["did_key"], did_key.encode()); assert_eq!(encoded["optional_did_key"], did_key.encode()); assert_eq!(encoded["status"], "throttled"); }
#[derive(Serialize, Deserialize)] struct DidKeyFixture { #[serde( serialize_with = "did_key_serialize_str", deserialize_with = "did_key_deserialize_str" )] did_key: DidKey<'static>, }
#[test] fn did_key_string_serde_roundtrips() { let did_key = DidKey(std::borrow::Cow::Owned(vec![0xed, 1, 2, 3])); let encoded = serde_json::to_string(&DidKeyFixture { did_key: did_key.clone(), }) .unwrap(); let decoded: DidKeyFixture = serde_json::from_str(&encoded).unwrap(); assert_eq!(decoded.did_key, did_key);
let with_prefix = format!(r#"{{"did_key":"did:key:{}"}}"#, did_key.encode()); let decoded_prefix: DidKeyFixture = serde_json::from_str(&with_prefix).unwrap(); assert_eq!(decoded_prefix.did_key, did_key); }}