Something went wrong. Try again.
Monorepo for Tangled forked from tangled.org/core
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373use std::sync::Arc;
use bobbin_runtime::{HttpRequest, HttpResponseHead, HttpTransport, NetworkError};use bytes::{Bytes, BytesMut};use futures::TryStreamExt;use http::{HeaderMap, StatusCode};use jacquard_common::DefaultStr;use jacquard_common::types::did::Did;use serde::Deserialize;use sha2::{Digest, Sha256};use thiserror::Error;use url::Url;
const MAX_BODY_BYTES: u64 = 64 * 1024;
#[derive(Debug, Error)]pub enum HydrantError { #[error("invalid base url scheme: {0}")] BadScheme(String), #[error("network: {0}")] Network(#[from] NetworkError), #[error("upstream returned status {0}")] Upstream(StatusCode), #[error("response body exceeded {limit} bytes")] BodyTooLarge { limit: u64 }, #[error("decode response: {0}")] Decode(#[from] serde_json::Error), #[error("stream chunk {0} doesn't match its hash")] ChunkMismatch(String),}
#[derive(Debug, Deserialize)]struct DescribedRepo { #[serde(rename = "didDoc")] did_doc: DescribedDoc,}
#[derive(Debug, Deserialize)]struct DescribedDoc { #[serde(default)] service: Vec<DescribedService>,}
#[derive(Debug, Deserialize)]struct DescribedService { id: String, #[serde(rename = "serviceEndpoint")] endpoint: serde_json::Value,}
fn knot_host(did: &Did<DefaultStr>, services: &[DescribedService]) -> Option<String> { let find = |id: &str| { services .iter() .find(|service| service.id == id || service.id.strip_prefix(did.as_str()) == Some(id)) }; // TODO(post-1.0): remove legacy `#atproto_pds` fallback let service = find("#tangled_knot").or_else(|| find("#atproto_pds"))?; let url = Url::parse(service.endpoint.as_str()?).ok()?; Some(url[url::Position::BeforeHost..url::Position::AfterPort].to_owned())}
#[derive(Debug, Deserialize)]struct RepoStatus { active: bool,}
#[derive(Debug, Deserialize)]struct StreamHead { id: Option<u64>,}
/// pick up `/stream` after `through` to end where a replay from 0 does#[derive(Clone, Debug, Deserialize)]pub struct StreamManifest { pub stream: String, pub through: Option<u64>, pub chunks: Vec<StreamChunk>,}
#[derive(Clone, Debug, Deserialize)]pub struct StreamChunk { pub first: u64, pub last: u64, pub bytes: u64, pub sha256: String,}
#[derive(Clone)]pub struct HydrantClient { http: Arc<dyn HttpTransport>, base: Url,}
impl std::fmt::Debug for HydrantClient { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("HydrantClient") .field("base", &self.base) .finish_non_exhaustive() }}
impl HydrantClient { // takes the websocket schemes too, callers share the stream's base url pub fn new(base: Url, http: Arc<dyn HttpTransport>) -> Result<Self, HydrantError> { let mut base = base; let scheme = match base.scheme() { "http" | "ws" => "http", "https" | "wss" => "https", other => return Err(HydrantError::BadScheme(other.to_owned())), }; base.set_scheme(scheme) .map_err(|()| HydrantError::BadScheme(scheme.to_owned()))?; Ok(Self { http, base }) }
// segment-wise so the colons in a did escape instead of parsing as a scheme fn endpoint<'a>( &self, segments: impl IntoIterator<Item = &'a str>, ) -> Result<Url, HydrantError> { let mut url = self.base.clone(); url.path_segments_mut() .map_err(|()| HydrantError::BadScheme(self.base.scheme().to_owned()))? .pop_if_empty() .extend(segments); Ok(url) }
pub fn xrpc_url(&self, nsid: &str) -> Result<Url, HydrantError> { self.endpoint(["xrpc", nsid]) }
/// the id of the newest event hydrant has stored, `None` while it stores none. /// a `cursor=0` replay has delivered everything hydrant held at the time of /// asking once it passes this id pub async fn stream_head(&self) -> Result<Option<u64>, HydrantError> { let bytes = self .get(self.endpoint(["stream", "head"])?, MAX_BODY_BYTES) .await?; Ok(serde_json::from_slice::<StreamHead>(&bytes)?.id) }
/// a 404 means hydrant isn't building stream chunks pub async fn stream_manifest(&self) -> Result<Option<StreamManifest>, HydrantError> { let url = self.endpoint(["stream", "manifest"])?; let Some(bytes) = found(self.get(url, MAX_BODY_BYTES).await)? else { return Ok(None); }; Ok(Some(serde_json::from_slice(&bytes)?)) }
/// still zstd compressed pub async fn stream_chunk(&self, chunk: &StreamChunk) -> Result<Bytes, HydrantError> { let url = self.endpoint(["stream", "chunk", &chunk.sha256])?; let bytes = self.get(url, chunk.bytes).await?.freeze(); let hex: String = Sha256::digest(&bytes) .iter() .map(|b| format!("{b:02x}")) .collect(); if hex != chunk.sha256 { return Err(HydrantError::ChunkMismatch(chunk.sha256.clone())); } Ok(bytes) }
/// a 404 means hydrant doesn't track the repo pub async fn repo_active(&self, did: &Did<DefaultStr>) -> Result<Option<bool>, HydrantError> { let mut url = self.xrpc_url("com.atproto.sync.getRepoStatus")?; url.query_pairs_mut().append_pair("did", did.as_str()); let Some(bytes) = found(self.get(url, MAX_BODY_BYTES).await)? else { return Ok(None); }; Ok(Some(serde_json::from_slice::<RepoStatus>(&bytes)?.active)) }
/// the host in a did's `#tangled_knot` service, `None` when its doc names no knot pub async fn repo_knot(&self, did: &Did<DefaultStr>) -> Result<Option<String>, HydrantError> { // the only hydrant endpoint with the raw doc, the rest are minidocs without services let mut url = self.endpoint(["xrpc", "com.atproto.repo.describeRepo"])?; url.query_pairs_mut().append_pair("repo", did.as_str()); let bytes = self.get(url, MAX_BODY_BYTES).await?; let described = serde_json::from_slice::<DescribedRepo>(&bytes)?; Ok(knot_host(did, &described.did_doc.service)) }
async fn get(&self, url: Url, limit: u64) -> Result<BytesMut, HydrantError> { let resp = self .http .execute(HttpRequest { url, headers: HeaderMap::new(), }) .await?; if resp.status != StatusCode::OK { return Err(HydrantError::Upstream(resp.status)); } read_bounded(resp, limit).await }}
fn found<T>(result: Result<T, HydrantError>) -> Result<Option<T>, HydrantError> { match result { Err(HydrantError::Upstream(StatusCode::NOT_FOUND)) => Ok(None), result => result.map(Some), }}
async fn read_bounded(resp: HttpResponseHead, limit: u64) -> Result<BytesMut, HydrantError> { if resp.content_length.is_some_and(|len| len > limit) { return Err(HydrantError::BodyTooLarge { limit }); } resp.body .map_err(HydrantError::Network) .try_fold(BytesMut::new(), |mut acc, chunk| async move { if (acc.len() as u64).saturating_add(chunk.len() as u64) > limit { return Err(HydrantError::BodyTooLarge { limit }); } acc.extend_from_slice(&chunk); Ok(acc) }) .await}
#[cfg(test)]mod tests { use super::*; use bobbin_runtime::ReqwestHttp; use bobbin_slingshot_client::default_http_client; use wiremock::matchers::{method, path, query_param}; use wiremock::{Mock, MockServer, ResponseTemplate};
fn did(value: &str) -> Did<DefaultStr> { Did::new_owned(value).unwrap() }
async fn client(server: &MockServer) -> HydrantClient { HydrantClient::new( Url::parse(&server.uri()).unwrap(), ReqwestHttp::shared(default_http_client().unwrap()), ) .unwrap() }
#[tokio::test] async fn websocket_base_is_accepted() { let client = HydrantClient::new( Url::parse("ws://127.0.0.1:3099").unwrap(), ReqwestHttp::shared(default_http_client().unwrap()), ) .unwrap(); assert_eq!(client.base.scheme(), "http"); }
#[tokio::test] async fn base_forms_all_address_the_same_endpoint() { let url = |base: &str| { HydrantClient::new( Url::parse(base).unwrap(), ReqwestHttp::shared(default_http_client().unwrap()), ) .unwrap() .xrpc_url("com.atproto.repo.describeRepo") .unwrap() .to_string() };
let xrpc = "xrpc/com.atproto.repo.describeRepo"; assert_eq!(url("http://h:3099"), format!("http://h:3099/{xrpc}")); assert_eq!(url("http://h:3099/"), format!("http://h:3099/{xrpc}")); assert_eq!(url("wss://h/hydrant/"), format!("https://h/hydrant/{xrpc}")); }
#[tokio::test] async fn stream_head_reads_the_newest_event_id() { let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/hydrant/stream/head")) .respond_with( ResponseTemplate::new(200).set_body_json(serde_json::json!({"id": 289455})), ) .mount(&server) .await;
let client = HydrantClient::new( Url::parse(&format!("{}/hydrant/", server.uri())).unwrap(), ReqwestHttp::shared(default_http_client().unwrap()), ) .unwrap(); assert_eq!(client.stream_head().await.unwrap(), Some(289455)); }
#[tokio::test] async fn an_empty_stream_has_no_head() { let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/stream/head")) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"id": null}))) .mount(&server) .await;
assert_eq!(client(&server).await.stream_head().await.unwrap(), None); }
#[tokio::test] async fn a_hydrant_without_stream_head_is_an_upstream_error() { let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/stream/head")) .respond_with(ResponseTemplate::new(404)) .mount(&server) .await;
let err = client(&server).await.stream_head().await.unwrap_err(); assert!( matches!(err, HydrantError::Upstream(StatusCode::NOT_FOUND)), "got {err:?}" ); }
fn described(services: serde_json::Value) -> ResponseTemplate { ResponseTemplate::new(200).set_body_json(serde_json::json!({ "did": "did:plc:kelp", "handle": "handle.invalid", "handleIsCorrect": false, "collections": [], "didDoc": { "id": "did:plc:kelp", "service": services }, })) }
#[tokio::test] async fn repo_knot_takes_the_host_of_the_tangled_knot_service() { let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/xrpc/com.atproto.repo.describeRepo")) .and(query_param("repo", "did:plc:kelp")) .respond_with(described(serde_json::json!([ { "id": "#atproto_pds", "type": "AtprotoPersonalDataServer", "serviceEndpoint": "https://knot.example.com", }, { "id": "did:plc:kelp#tangled_knot", "type": "TangledKnot", "serviceEndpoint": "https://knot.example.com:8443/repo/abc", }, ]))) .mount(&server) .await;
let knot = client(&server).await.repo_knot(&did("did:plc:kelp")).await; assert_eq!(knot.unwrap().as_deref(), Some("knot.example.com:8443")); }
#[tokio::test] async fn repo_knot_falls_back_to_the_legacy_pds_service() { let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/xrpc/com.atproto.repo.describeRepo")) .respond_with(described(serde_json::json!([{ "id": "#atproto_pds", "type": "AtprotoPersonalDataServer", "serviceEndpoint": "https://pds.example.com", }]))) .mount(&server) .await;
let knot = client(&server).await.repo_knot(&did("did:plc:kelp")).await; assert_eq!(knot.unwrap().as_deref(), Some("pds.example.com")); }}