use smol_str::SmolStr;
use url::Url;
use super::*;
#[derive(Deserialize)]
pub struct RequestCrawl {
pub hostname: SmolStr,
}
pub const PATH: &str = "/xrpc/com.atproto.sync.requestCrawl";
pub async fn handle(
State(hydrant): State<Hydrant>,
ExtractXrpc(req): ExtractXrpc<RequestCrawl>,
) -> XrpcResult<StatusCode> {
let url = Url::parse(&format!("wss://{}/", req.hostname)).map_err(bad_request)?;
match hydrant
.firehose
.add_source_validated(url)
.await
.map_err(internal_error)?
{
Ok(_) => Ok(StatusCode::OK),
Err(crate::control::firehose::PdsAdmissionRejection::InvalidHost(error)) => {
Err(bad_request(error))
}
Err(crate::control::firehose::PdsAdmissionRejection::HostBanned) => Err(xrpc_error(
StatusCode::BAD_REQUEST,
"HostBanned",
Some("host is banned".into()),
)),
Err(crate::control::firehose::PdsAdmissionRejection::DailyLimitReached) => Err(
rate_limited("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(RequestCrawl {
hostname: "blocked.example".into(),
}),
)
.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(RequestCrawl {
hostname: "limited.example".into(),
}),
)
.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(RequestCrawl {
hostname: "127.0.0.1".into(),
}),
)
.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 router_serves_request_crawl() -> 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 app = super::super::router(false).with_state(hydrant);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.map_err(|error| miette::miette!(error))?;
let addr = listener
.local_addr()
.map_err(|error| miette::miette!(error))?;
tokio::spawn(async move { axum::serve(listener, app).await });
// a private host is rejected before admission, so this only proves the route is mounted
let response = reqwest::Client::new()
.post(format!("http://{addr}{PATH}"))
.json(&serde_json::json!({ "hostname": "127.0.0.1" }))
.send()
.await
.map_err(|error| miette::miette!(error))?;
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
assert_eq!(
response
.json::<serde_json::Value>()
.await
.map_err(|error| miette::miette!(error))?["error"],
"InvalidRequest"
);
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(RequestCrawl {
hostname: "cold.example".into(),
}),
)
.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(())
}
}