Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187#[cfg(test)]use jacquard_api::com_atproto::sync::request_crawl::RequestCrawl;use jacquard_api::com_atproto::sync::request_crawl::{ RequestCrawlError, RequestCrawlRequest, RequestCrawlResponse,};use url::Url;
use super::*;
pub async fn handle( State(hydrant): State<Hydrant>, ExtractXrpc(req): ExtractXrpc<RequestCrawlRequest>,) -> XrpcResult<StatusCode, RequestCrawlError> { let nsid = RequestCrawlResponse::NSID; let url = Url::parse(&format!("wss://{}/", req.hostname)) .map_err(|error| bad_request(nsid, error))?;
match hydrant .firehose .add_source_validated(url) .await .map_err(|error| internal_error(nsid, error))? { Ok(_) => Ok(StatusCode::OK), Err(crate::control::firehose::PdsAdmissionRejection::InvalidHost(error)) => { Err(bad_request(nsid, error)) } Err(crate::control::firehose::PdsAdmissionRejection::HostBanned) => { Err(XrpcErrorResponse { status: StatusCode::BAD_REQUEST, error: jacquard_common::xrpc::XrpcError::Xrpc(RequestCrawlError::HostBanned(Some( "host is banned".into(), ))), }) } Err(crate::control::firehose::PdsAdmissionRejection::DailyLimitReached) => Err( rate_limited(nsid, "POST", "daily limit for new PDS sources reached"), ), }}
#[cfg(test)]mod tests { use super::*; use axum::response::IntoResponse;
#[tokio::test] async fn banned_request_crawl_json_contract() -> miette::Result<()> { let tmp = tempfile::tempdir().map_err(|e| miette::miette!(e))?; let hydrant = Hydrant::new(crate::config::Config { database_path: tmp.path().to_path_buf(), ..Default::default() }) .await?; hydrant.pds.ban("blocked.example").await?;
let error = handle( State(hydrant), ExtractXrpc::<RequestCrawlRequest>(RequestCrawl { hostname: "blocked.example".into(), extra_data: Default::default(), }), ) .await .expect_err("banned host must be rejected"); let response = error.into_response(); assert_eq!(response.status(), StatusCode::BAD_REQUEST); let body = axum::body::to_bytes(response.into_body(), 1024) .await .map_err(|e| miette::miette!(e))?; assert_eq!( serde_json::from_slice::<serde_json::Value>(&body).map_err(|e| miette::miette!(e))?, serde_json::json!({ "error": "HostBanned", "message": "host is banned", }) ); Ok(()) }
#[tokio::test] async fn limited_request_crawl_json_contract() -> miette::Result<()> { let tmp = tempfile::tempdir().map_err(|error| miette::miette!(error))?; let hydrant = Hydrant::new(crate::config::Config { database_path: tmp.path().to_path_buf(), new_host_limit: Some(0), ..Default::default() }) .await?;
let error = handle( State(hydrant), ExtractXrpc::<RequestCrawlRequest>(RequestCrawl { hostname: "limited.example".into(), extra_data: Default::default(), }), ) .await .expect_err("exhausted PDS quota must be rejected"); let response = error.into_response(); assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); let body = axum::body::to_bytes(response.into_body(), 1024) .await .map_err(|error| miette::miette!(error))?; assert_eq!( serde_json::from_slice::<serde_json::Value>(&body) .map_err(|error| miette::miette!(error))?, serde_json::json!({ "error": "RateLimitExceeded", "message": "daily limit for new PDS sources reached", }) ); Ok(()) }
#[tokio::test] async fn invalid_request_crawl_json_contract() -> miette::Result<()> { let tmp = tempfile::tempdir().map_err(|error| miette::miette!(error))?; let hydrant = Hydrant::new(crate::config::Config { database_path: tmp.path().to_path_buf(), ..Default::default() }) .await?;
let error = handle( State(hydrant), ExtractXrpc::<RequestCrawlRequest>(RequestCrawl { hostname: "127.0.0.1".into(), extra_data: Default::default(), }), ) .await .expect_err("private PDS host must be rejected"); let response = error.into_response(); assert_eq!(response.status(), StatusCode::BAD_REQUEST); let body = axum::body::to_bytes(response.into_body(), 1024) .await .map_err(|error| miette::miette!(error))?; assert_eq!( serde_json::from_slice::<serde_json::Value>(&body) .map_err(|error| miette::miette!(error))?, serde_json::json!({ "error": "InvalidRequest", "message": "PDS hostname is a non-public IP address", }) ); Ok(()) }
#[tokio::test] async fn cold_request_crawl_persists_its_quota() -> miette::Result<()> { let tmp = tempfile::tempdir().map_err(|error| miette::miette!(error))?; let hydrant = Hydrant::new(crate::config::Config { database_path: tmp.path().to_path_buf(), new_host_limit: Some(1), ..Default::default() }) .await?; hydrant.firehose.startable_for_test(); hydrant.firehose.disable_spawns_for_test();
assert_eq!( handle( State(hydrant.clone()), ExtractXrpc::<RequestCrawlRequest>(RequestCrawl { hostname: "cold.example".into(), extra_data: Default::default(), }), ) .await .map_err(|error| miette::miette!("request crawl failed: {:?}", error.status))?, StatusCode::OK ); assert!(matches!( crate::db::load_pds_daily_adds(&hydrant.state.db)?, Some((_, 1)) )); hydrant .firehose .remove_source( &Url::parse("wss://cold.example/").map_err(|error| miette::miette!(error))?, ) .await?; Ok(()) }}