#![allow(dead_code)]
use std::{hash::Hash, time::Duration};
use jacquard_common::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 crate::{db::types::DidKey, types::RepoStatus};
pub mod drain;
pub mod stop;
pub mod throttle;
pub(crate) mod hex {
pub(crate) fn encode(bytes: impl AsRef<[u8]>) -> String {
data_encoding::HEXLOWER.encode(bytes.as_ref())
}
#[cfg(test)]
pub(crate) fn decode(s: &str) -> Result<Vec<u8>, data_encoding::DecodeError> {
data_encoding::HEXLOWER.decode(s.as_bytes())
}
}
pub(crate) fn parse_duration(s: &str) -> Option<Duration> {
let mut rest = s.trim();
if rest.is_empty() {
return None;
}
if rest == "0" {
return Some(Duration::ZERO);
}
let mut total = Duration::ZERO;
while !rest.is_empty() {
let num_end = rest.find(|c: char| !c.is_ascii_digit()).unwrap_or(rest.len());
if num_end == 0 {
return None;
}
let val: u64 = rest[..num_end].parse().ok()?;
rest = rest[num_end..].trim_start();
let unit_end = rest
.find(|c: char| !c.is_alphabetic() && c != 'µ')
.unwrap_or(rest.len());
if unit_end == 0 {
return None;
}
let unit = &rest[..unit_end];
rest = rest[unit_end..].trim_start();
let part = match unit {
"ns" | "nsec" | "nanos" | "nanosecond" | "nanoseconds" => Duration::from_nanos(val),
"us" | "µs" | "usec" | "micros" | "microsecond" | "microseconds" => {
Duration::from_micros(val)
}
"ms" | "msec" | "millis" | "millisecond" | "milliseconds" => Duration::from_millis(val),
"s" | "sec" | "secs" | "second" | "seconds" => Duration::from_secs(val),
"m" | "min" | "mins" | "minute" | "minutes" => Duration::from_secs(val.checked_mul(60)?),
"h" | "hr" | "hrs" | "hour" | "hours" => Duration::from_secs(val.checked_mul(3600)?),
"d" | "day" | "days" => Duration::from_secs(val.checked_mul(86_400)?),
"w" | "week" | "weeks" => Duration::from_secs(val.checked_mul(604_800)?),
"M" | "month" | "months" => Duration::from_secs(val.checked_mul(2_630_016)?),
"y" | "year" | "years" => Duration::from_secs(val.checked_mul(31_557_600)?),
_ => return None,
};
total = total.checked_add(part)?;
}
Some(total)
}
pub(crate) fn format_duration(d: Duration) -> impl std::fmt::Display {
struct Formatted(Duration);
impl std::fmt::Display for Formatted {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let secs = self.0.as_secs();
let nanos = self.0.subsec_nanos();
if secs == 0 && nanos == 0 {
return f.write_str("0s");
}
let days = secs / 86_400;
let hours = (secs % 86_400) / 3600;
let mins = (secs % 3600) / 60;
let s = secs % 60;
let ms = nanos / 1_000_000;
let us = (nanos % 1_000_000) / 1_000;
let ns = nanos % 1_000;
let mut wrote = false;
for (val, unit) in [
(days, "d"),
(hours, "h"),
(mins, "m"),
(s, "s"),
(u64::from(ms), "ms"),
(u64::from(us), "us"),
(u64::from(ns), "ns"),
] {
if val > 0 {
if wrote {
f.write_str(" ")?;
}
write!(f, "{val}{unit}")?;
wrote = true;
}
}
Ok(())
}
}
Formatted(d)
}
/// runs `f` under a subscriber with `filter` and returns everything it logged
#[cfg(test)]
pub(crate) fn capture_logs(filter: &str, f: impl FnOnce()) -> String {
#[derive(Clone, Default)]
struct Logs(std::sync::Arc<std::sync::Mutex<Vec<u8>>>);
impl std::io::Write for Logs {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.0.lock().unwrap().extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
let logs = Logs::default();
let subscriber = tracing_subscriber::fmt()
.with_env_filter(filter)
.with_writer({
let logs = logs.clone();
move || logs.clone()
})
.with_ansi(false)
.finish();
tracing::subscriber::with_default(subscriber, f);
String::from_utf8(logs.0.lock().unwrap().clone()).unwrap()
}
#[allow(dead_code)]
/// checks if the error contains a hyper / std io timeout error
pub 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 errors
pub 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() -> Fut
where
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 F
where
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 codes
pub 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(crate) fn invalid_handle() -> Handle {
unsafe { Handle::unchecked(smol_str::SmolStr::new_static("handle.invalid")) }
}
/// returns hash of value using ahash
pub 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(crate::car::compute_cid(b"body"));
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);
}
}