mirrors/tranquil-pds
1
0
Fork 0
mirror of https://tangled.org/tranquil.farm/tranquil-pds synced 2026-08-17 14:56:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

My misc TODOs

Browse Source
This commit is contained in:
lewis
2025-12-12 20:28:34 +02:00
parent 9b75db2ede
commit 36e31d7ef9
30 changed files with 1000 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.sqlx/query-76c6ef1d5395105a0cdedb27ca321c9e3eae1ce87c223b706ed81ebf973875f3.json → .sqlx/query-08c08b0644d79d5de72f3500dd7dbb8827af340e3c04fec9a5c28aeff46e0c97.json
Generated
+8 -2
View File
@@ -1,6 +1,6 @@
{
"db_name": "PostgreSQL",
"query": "SELECT id, password_hash FROM users WHERE did = $1",
"query": "SELECT id, password_hash, handle FROM users WHERE did = $1",
"describe": {
"columns": [
{
@@ -12,6 +12,11 @@
"ordinal": 1,
"name": "password_hash",
"type_info": "Text"
},
{
"ordinal": 2,
"name": "handle",
"type_info": "Text"
}
],
"parameters": {
@@ -20,9 +25,10 @@
]
},
"nullable": [
false,
false,
false
]
},
"hash": "76c6ef1d5395105a0cdedb27ca321c9e3eae1ce87c223b706ed81ebf973875f3"
"hash": "08c08b0644d79d5de72f3500dd7dbb8827af340e3c04fec9a5c28aeff46e0c97"
}
.sqlx/query-e223898d53602c1c8b23eb08a4b96cf20ac349d1fa4e91334b225d3069209dcf.json
Generated
+22
View File
@@ -0,0 +1,22 @@
{
"db_name": "PostgreSQL",
"query": "SELECT handle FROM users WHERE id = $1",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "handle",
"type_info": "Text"
}
],
"parameters": {
"Left": [
"Uuid"
]
},
"nullable": [
false
]
},
"hash": "e223898d53602c1c8b23eb08a4b96cf20ac349d1fa4e91334b225d3069209dcf"
}
Cargo.lock
Generated
+73 -2
View File
@@ -98,6 +98,12 @@ version = "1.0.100"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61"
[[package]]
name = "arc-swap"
version = "1.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "69f7f8c3906b62b754cd5326047894316021dcfe5a194c8ea52bdd94934a3457"
[[package]]
name = "assert-json-diff"
version = "2.0.2"
@@ -688,6 +694,15 @@ dependencies = [
"syn 2.0.111",
]
[[package]]
name = "backon"
version = "1.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef"
dependencies = [
"fastrand",
]
[[package]]
name = "base-x"
version = "0.2.11"
@@ -931,6 +946,7 @@ dependencies = [
"p256 0.13.2",
"p384",
"rand 0.8.5",
"redis",
"reqwest",
"serde",
"serde_bytes",
@@ -1177,6 +1193,20 @@ version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b"
[[package]]
name = "combine"
version = "4.6.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd"
dependencies = [
"bytes",
"futures-core",
"memchr",
"pin-project-lite",
"tokio",
"tokio-util",
]
[[package]]
name = "compression-codecs"
version = "0.4.33"
@@ -2971,6 +3001,15 @@ dependencies = [
"unsigned-varint 0.7.2",
]
[[package]]
name = "itertools"
version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186"
dependencies = [
"either",
]
[[package]]
name = "itertools"
version = "0.14.0"
@@ -4241,7 +4280,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425"
dependencies = [
"anyhow",
"itertools",
"itertools 0.14.0",
"proc-macro2",
"quote",
"syn 2.0.111",
@@ -4441,6 +4480,32 @@ dependencies = [
"bitflags",
]
[[package]]
name = "redis"
version = "0.27.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "09d8f99a4090c89cc489a94833c901ead69bfbf3877b4867d5482e321ee875bc"
dependencies = [
"arc-swap",
"async-trait",
"backon",
"bytes",
"combine",
"futures",
"futures-util",
"itertools 0.13.0",
"itoa",
"num-bigint",
"percent-encoding",
"pin-project-lite",
"ryu",
"sha1_smol",
"socket2 0.5.10",
"tokio",
"tokio-util",
"url",
]
[[package]]
name = "redox_syscall"
version = "0.5.18"
@@ -5054,6 +5119,12 @@ dependencies = [
"digest",
]
[[package]]
name = "sha1_smol"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d"
[[package]]
name = "sha2"
version = "0.10.9"
@@ -5646,7 +5717,7 @@ dependencies = [
"etcetera 0.11.0",
"ferroid",
"futures",
"itertools",
"itertools 0.14.0",
"log",
"memchr",
"parse-display",
Cargo.toml
+1
View File
@@ -49,6 +49,7 @@ urlencoding = "2.1"
uuid = { version = "1.19.0", features = ["v4", "fast-rng"] }
iroh-car = "0.5.1"
image = { version = "0.25", default-features = false, features = ["jpeg", "png", "gif", "webp"] }
redis = { version = "0.27", features = ["tokio-comp", "connection-manager"] }
[features]
external-infra = []
TODO.md
+17 -7
View File
@@ -198,9 +198,10 @@ These are implemented at PDS level to enable local-first reads (read-after-write
- [x] Implement Atomic Repo Transactions.
- [x] Ensure `blocks` write, `repo_root` update, `records` index update, and `sequencer` event are committed in a single transaction.
- [x] Implement concurrency control (row-level locking via FOR UPDATE).
- [ ] DID Cache
- [ ] Implement caching layer for DID resolution (Redis or in-memory).
- [ ] Handle cache invalidation/expiry.
- [x] DID Cache
- [x] Implement caching layer for DID resolution (valkey).
- [x] Handle cache invalidation/expiry.
- [x] Graceful fallback to no-cache when Valkey unavailable.
- [x] Crawlers Service
- [x] Implement `Crawlers` service (debounce notifications to relays).
- [x] 20-minute notification debounce.
@@ -229,6 +230,14 @@ These are implemented at PDS level to enable local-first reads (read-after-write
- [x] Per-IP rate limiting on OAuth token endpoint (30/min).
- [x] Per-IP rate limiting on password reset (5/hour).
- [x] Per-IP rate limiting on account creation (10/hour).
- [x] Per-IP rate limiting on refreshSession (60/min).
- [x] Per-IP rate limiting on OAuth authorize POST (10/min).
- [x] Per-IP rate limiting on OAuth 2FA POST (10/min).
- [x] Per-IP rate limiting on OAuth PAR (30/min).
- [x] Per-IP rate limiting on OAuth revoke/introspect (30/min).
- [x] Per-IP rate limiting on createAppPassword (10/min).
- [x] Per-IP rate limiting on email endpoints (5/hour).
- [x] Distributed rate limiting via Valkey/Redis (with in-memory fallback).
- [x] Circuit Breakers
- [x] PLC directory circuit breaker (5 failures → open, 60s timeout).
- [x] Relay notification circuit breaker (10 failures → open, 30s timeout).
@@ -237,12 +246,13 @@ These are implemented at PDS level to enable local-first reads (read-after-write
- [x] Signal command injection prevention (phone number validation).
- [x] Constant-time signature comparison.
- [x] SSRF protection for outbound requests.
- [x] Timing attack protection (dummy bcrypt on user-not-found prevents account enumeration).
## Lewis' fabulous mini-list of remaining TODOs
- [ ] The OAuth authorize POST endpoint has no rate limiting, allowing password brute-forcing. Fix this and audit all oauth and 2fa surface again.
- [ ] DID resolution caching (valkey).
- [ ] Record schema validation (generic validation framework).
- [ ] Fix any remaining TODOs in the code.
- [x] The OAuth authorize POST endpoint has no rate limiting, allowing password brute-forcing. Fix this and audit all oauth and 2fa surface again.
- [x] DID resolution caching (valkey).
- [x] Record schema validation (generic validation framework).
- [x] Fix any remaining TODOs in the code.
## Future: Web Management UI
A single-page web app for account management. The frontend (JS framework) calls existing ATProto XRPC endpoints - no server-side rendering or bespoke HTML form handlers.
docker-compose.yaml
+10
View File
@@ -11,9 +11,11 @@ services:
environment:
DATABASE_URL: postgres://postgres:postgres@db:5432/pds
S3_ENDPOINT: http://objsto:9000
VALKEY_URL: redis://cache:6379
depends_on:
- db
- objsto
- cache
db:
image: postgres:latest
@@ -38,6 +40,14 @@ services:
- minio_data:/data
command: server /data --console-address ":9001"
cache:
image: valkey/valkey:8-alpine
ports:
- "6379:6379"
volumes:
- valkey_data:/data
volumes:
postgres_data:
minio_data:
valkey_data:
scripts/test-infra.sh
+20 -3
View File
@@ -38,7 +38,7 @@ start_infra() {
rm -f "$INFRA_FILE"
fi
$CONTAINER_CMD rm -f "${CONTAINER_PREFIX}-postgres" "${CONTAINER_PREFIX}-minio" 2>/dev/null || true
$CONTAINER_CMD rm -f "${CONTAINER_PREFIX}-postgres" "${CONTAINER_PREFIX}-minio" "${CONTAINER_PREFIX}-valkey" 2>/dev/null || true
echo "Starting PostgreSQL..."
$CONTAINER_CMD run -d \
@@ -59,11 +59,19 @@ start_infra() {
--label bspds_test=true \
minio/minio:latest server /data >/dev/null
echo "Starting Valkey..."
$CONTAINER_CMD run -d \
--name "${CONTAINER_PREFIX}-valkey" \
-P \
--label bspds_test=true \
valkey/valkey:8-alpine >/dev/null
echo "Waiting for services to be ready..."
sleep 2
PG_PORT=$($CONTAINER_CMD port "${CONTAINER_PREFIX}-postgres" 5432 | head -1 | cut -d: -f2)
MINIO_PORT=$($CONTAINER_CMD port "${CONTAINER_PREFIX}-minio" 9000 | head -1 | cut -d: -f2)
VALKEY_PORT=$($CONTAINER_CMD port "${CONTAINER_PREFIX}-valkey" 6379 | head -1 | cut -d: -f2)
for i in {1..30}; do
if $CONTAINER_CMD exec "${CONTAINER_PREFIX}-postgres" pg_isready -U postgres >/dev/null 2>&1; then
@@ -81,6 +89,14 @@ start_infra() {
sleep 1
done
for i in {1..30}; do
if $CONTAINER_CMD exec "${CONTAINER_PREFIX}-valkey" valkey-cli ping 2>/dev/null | grep -q PONG; then
break
fi
echo "Waiting for Valkey... ($i/30)"
sleep 1
done
echo "Creating MinIO bucket..."
$CONTAINER_CMD run --rm --network host \
-e MC_HOST_minio="http://minioadmin:minioadmin@127.0.0.1:${MINIO_PORT}" \
@@ -94,6 +110,7 @@ export S3_BUCKET="test-bucket"
export AWS_ACCESS_KEY_ID="minioadmin"
export AWS_SECRET_ACCESS_KEY="minioadmin"
export AWS_REGION="us-east-1"
export VALKEY_URL="redis://127.0.0.1:${VALKEY_PORT}"
export BSPDS_TEST_INFRA_READY="1"
export BSPDS_ALLOW_INSECURE_SECRETS="1"
export SKIP_IMPORT_VERIFICATION="true"
@@ -108,7 +125,7 @@ EOF
stop_infra() {
echo "Stopping test infrastructure..."
$CONTAINER_CMD rm -f "${CONTAINER_PREFIX}-postgres" "${CONTAINER_PREFIX}-minio" 2>/dev/null || true
$CONTAINER_CMD rm -f "${CONTAINER_PREFIX}-postgres" "${CONTAINER_PREFIX}-minio" "${CONTAINER_PREFIX}-valkey" 2>/dev/null || true
rm -f "$INFRA_FILE"
echo "Infrastructure stopped."
}
@@ -157,7 +174,7 @@ case "${1:-}" in
echo "Usage: $0 {start|stop|restart|status|env}"
echo ""
echo "Commands:"
echo " start - Start test infrastructure (Postgres, MinIO)"
echo " start - Start test infrastructure (Postgres, MinIO, Valkey)"
echo " stop - Stop and remove test containers"
echo " restart - Stop then start infrastructure"
echo " status - Show infrastructure status"
src/api/admin/account/delete.rs
+5 -3
View File
@@ -37,12 +37,12 @@ pub async fn delete_account(
.into_response();
}
let user = sqlx::query!("SELECT id FROM users WHERE did = $1", did)
let user = sqlx::query!("SELECT id, handle FROM users WHERE did = $1", did)
.fetch_optional(&state.db)
.await;
let user_id = match user {
Ok(Some(row)) => row.id,
let (user_id, handle) = match user {
Ok(Some(row)) => (row.id, row.handle),
Ok(None) => {
return (
StatusCode::NOT_FOUND,
@@ -186,5 +186,7 @@ pub async fn delete_account(
.into_response();
}
let _ = state.cache.delete(&format!("handle:{}", handle)).await;
(StatusCode::OK, Json(json!({}))).into_response()
}
src/api/admin/account/update.rs
+10
View File
@@ -108,6 +108,12 @@ pub async fn update_account_handle(
.into_response();
}
let old_handle = sqlx::query_scalar!("SELECT handle FROM users WHERE did = $1", did)
.fetch_optional(&state.db)
.await
.ok()
.flatten();
let existing = sqlx::query!("SELECT id FROM users WHERE handle = $1 AND did != $2", handle, did)
.fetch_optional(&state.db)
.await;
@@ -133,6 +139,10 @@ pub async fn update_account_handle(
)
.into_response();
}
if let Some(old) = old_handle {
let _ = state.cache.delete(&format!("handle:{}", old)).await;
}
let _ = state.cache.delete(&format!("handle:{}", handle)).await;
(StatusCode::OK, Json(json!({}))).into_response()
}
Err(e) => {
src/api/admin/status.rs
+7
View File
@@ -305,6 +305,13 @@ pub async fn update_subject_status(
.into_response();
}
if let Ok(Some(handle)) = sqlx::query_scalar!("SELECT handle FROM users WHERE did = $1", did)
.fetch_optional(&state.db)
.await
{
let _ = state.cache.delete(&format!("handle:{}", handle)).await;
}
return (
StatusCode::OK,
Json(json!({
src/api/identity/did.rs
+19 -1
View File
@@ -33,12 +33,18 @@ pub async fn resolve_handle(
.into_response();
}
let cache_key = format!("handle:{}", handle);
if let Some(did) = state.cache.get(&cache_key).await {
return (StatusCode::OK, Json(json!({ "did": did }))).into_response();
}
let user = sqlx::query!("SELECT did FROM users WHERE handle = $1", handle)
.fetch_optional(&state.db)
.await;
match user {
Ok(Some(row)) => {
let _ = state.cache.set(&cache_key, &row.did, std::time::Duration::from_secs(300)).await;
(StatusCode::OK, Json(json!({ "did": row.did }))).into_response()
}
Ok(None) => (
@@ -406,6 +412,12 @@ pub async fn update_handle(
.into_response();
}
let old_handle = sqlx::query_scalar!("SELECT handle FROM users WHERE id = $1", user_id)
.fetch_optional(&state.db)
.await
.ok()
.flatten();
let existing = sqlx::query!("SELECT id FROM users WHERE handle = $1 AND id != $2", new_handle, user_id)
.fetch_optional(&state.db)
.await;
@@ -423,7 +435,13 @@ pub async fn update_handle(
.await;
match result {
Ok(_) => (StatusCode::OK, Json(json!({}))).into_response(),
Ok(_) => {
if let Some(old) = old_handle {
let _ = state.cache.delete(&format!("handle:{}", old)).await;
}
let _ = state.cache.delete(&format!("handle:{}", new_handle)).await;
(StatusCode::OK, Json(json!({}))).into_response()
}
Err(e) => {
error!("DB error updating handle: {:?}", e);
(
src/api/repo/record/batch.rs
+11
View File
@@ -1,3 +1,4 @@
use super::validation::validate_record;
use crate::api::repo::record::utils::{commit_and_log, RecordOp};
use crate::repo::tracking::TrackingBlockStore;
use crate::state::AppState;
@@ -211,6 +212,11 @@ pub async fn apply_writes(
rkey,
value,
} => {
if input.validate.unwrap_or(true) {
if let Err(err_response) = validate_record(value, collection) {
return err_response;
}
}
let rkey = rkey
.clone()
.unwrap_or_else(|| Utc::now().format("%Y%m%d%H%M%S%f").to_string());
@@ -249,6 +255,11 @@ pub async fn apply_writes(
rkey,
value,
} => {
if input.validate.unwrap_or(true) {
if let Err(err_response) = validate_record(value, collection) {
return err_response;
}
}
let mut record_bytes = Vec::new();
if serde_ipld_dagcbor::to_writer(&mut record_bytes, value).is_err() {
return (StatusCode::BAD_REQUEST, Json(json!({"error": "InvalidRecord", "message": "Failed to serialize record"}))).into_response();
src/api/repo/record/mod.rs
+1
View File
@@ -2,6 +2,7 @@ pub mod batch;
pub mod delete;
pub mod read;
pub mod utils;
pub mod validation;
pub mod write;
pub use batch::apply_writes;
src/api/repo/record/validation.rs
+38
View File
@@ -0,0 +1,38 @@
use crate::validation::{RecordValidator, ValidationError};
use axum::{
http::StatusCode,
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
pub fn validate_record(record: &serde_json::Value, collection: &str) -> Result<(), Response> {
let validator = RecordValidator::new();
match validator.validate(record, collection) {
Ok(_) => Ok(()),
Err(ValidationError::MissingType) => Err((
StatusCode::BAD_REQUEST,
Json(json!({"error": "InvalidRecord", "message": "Record must have a $type field"})),
).into_response()),
Err(ValidationError::TypeMismatch { expected, actual }) => Err((
StatusCode::BAD_REQUEST,
Json(json!({"error": "InvalidRecord", "message": format!("Record $type '{}' does not match collection '{}'", actual, expected)})),
).into_response()),
Err(ValidationError::MissingField(field)) => Err((
StatusCode::BAD_REQUEST,
Json(json!({"error": "InvalidRecord", "message": format!("Missing required field: {}", field)})),
).into_response()),
Err(ValidationError::InvalidField { path, message }) => Err((
StatusCode::BAD_REQUEST,
Json(json!({"error": "InvalidRecord", "message": format!("Invalid field '{}': {}", path, message)})),
).into_response()),
Err(ValidationError::InvalidDatetime { path }) => Err((
StatusCode::BAD_REQUEST,
Json(json!({"error": "InvalidRecord", "message": format!("Invalid datetime format at '{}'", path)})),
).into_response()),
Err(e) => Err((
StatusCode::BAD_REQUEST,
Json(json!({"error": "InvalidRecord", "message": e.to_string()})),
).into_response()),
}
}
src/api/repo/record/write.rs
+5 -16
View File
@@ -1,3 +1,4 @@
use super::validation::validate_record;
use crate::api::repo::record::utils::{commit_and_log, RecordOp};
use crate::repo::tracking::TrackingBlockStore;
use crate::state::AppState;
@@ -156,14 +157,8 @@ pub async fn create_record(
};
if input.validate.unwrap_or(true) {
if input.collection == "app.bsky.feed.post" {
if input.record.get("text").is_none() || input.record.get("createdAt").is_none() {
return (
StatusCode::BAD_REQUEST,
Json(json!({"error": "InvalidRecord", "message": "Record validation failed"})),
)
.into_response();
}
if let Err(err_response) = validate_record(&input.record, &input.collection) {
return err_response;
}
}
@@ -263,14 +258,8 @@ pub async fn put_record(
let key = format!("{}/{}", collection_nsid, input.rkey);
if input.validate.unwrap_or(true) {
if input.collection == "app.bsky.feed.post" {
if input.record.get("text").is_none() || input.record.get("createdAt").is_none() {
return (
StatusCode::BAD_REQUEST,
Json(json!({"error": "InvalidRecord", "message": "Record validation failed"})),
)
.into_response();
}
if let Err(err_response) = validate_record(&input.record, &input.collection) {
return err_response;
}
}
src/api/server/account_status.rs
+28 -5
View File
@@ -123,12 +123,23 @@ pub async fn activate_account(
Err(e) => return ApiError::from(e).into_response(),
};
let handle = sqlx::query_scalar!("SELECT handle FROM users WHERE did = $1", did)
.fetch_optional(&state.db)
.await
.ok()
.flatten();
let result = sqlx::query!("UPDATE users SET deactivated_at = NULL WHERE did = $1", did)
.execute(&state.db)
.await;
match result {
Ok(_) => (StatusCode::OK, Json(json!({}))).into_response(),
Ok(_) => {
if let Some(h) = handle {
let _ = state.cache.delete(&format!("handle:{}", h)).await;
}
(StatusCode::OK, Json(json!({}))).into_response()
}
Err(e) => {
error!("DB error activating account: {:?}", e);
(
@@ -163,12 +174,23 @@ pub async fn deactivate_account(
Err(e) => return ApiError::from(e).into_response(),
};
let handle = sqlx::query_scalar!("SELECT handle FROM users WHERE did = $1", did)
.fetch_optional(&state.db)
.await
.ok()
.flatten();
let result = sqlx::query!("UPDATE users SET deactivated_at = NOW() WHERE did = $1", did)
.execute(&state.db)
.await;
match result {
Ok(_) => (StatusCode::OK, Json(json!({}))).into_response(),
Ok(_) => {
if let Some(h) = handle {
let _ = state.cache.delete(&format!("handle:{}", h)).await;
}
(StatusCode::OK, Json(json!({}))).into_response()
}
Err(e) => {
error!("DB error deactivating account: {:?}", e);
(
@@ -283,14 +305,14 @@ pub async fn delete_account(
}
let user = sqlx::query!(
"SELECT id, password_hash FROM users WHERE did = $1",
"SELECT id, password_hash, handle FROM users WHERE did = $1",
did
)
.fetch_optional(&state.db)
.await;
let (user_id, password_hash) = match user {
Ok(Some(row)) => (row.id, row.password_hash),
let (user_id, password_hash, handle) = match user {
Ok(Some(row)) => (row.id, row.password_hash, row.handle),
Ok(None) => {
return (
StatusCode::BAD_REQUEST,
@@ -437,6 +459,7 @@ pub async fn delete_account(
)
.into_response();
}
let _ = state.cache.delete(&format!("handle:{}", handle)).await;
info!("Account {} deleted successfully", did);
(StatusCode::OK, Json(json!({}))).into_response()
}
src/api/server/app_password.rs
+21 -1
View File
@@ -5,11 +5,12 @@ use crate::util::get_user_id_by_did;
use axum::{
Json,
extract::State,
http::HeaderMap,
response::{IntoResponse, Response},
};
use serde::{Deserialize, Serialize};
use serde_json::json;
use tracing::error;
use tracing::{error, warn};
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
@@ -76,9 +77,28 @@ pub struct CreateAppPasswordOutput {
pub async fn create_app_password(
State(state): State<AppState>,
headers: HeaderMap,
BearerAuth(auth_user): BearerAuth,
Json(input): Json<CreateAppPasswordInput>,
) -> Response {
let client_ip = crate::rate_limit::extract_client_ip(&headers, None);
if !state.distributed_rate_limiter.check_rate_limit(
&format!("app_password:{}", client_ip),
10,
60_000,
).await {
if state.rate_limiters.app_password.check_key(&client_ip).is_err() {
warn!(ip = %client_ip, "App password creation rate limit exceeded");
return (
axum::http::StatusCode::TOO_MANY_REQUESTS,
Json(json!({
"error": "RateLimitExceeded",
"message": "Too many requests. Please try again later."
})),
).into_response();
}
}
let user_id = match get_user_id_by_did(&state.db, &auth_user.did).await {
Ok(id) => id,
Err(e) => return ApiError::from(e).into_response(),
src/api/server/email.rs
+36
View File
@@ -26,6 +26,24 @@ pub async fn request_email_update(
headers: axum::http::HeaderMap,
Json(input): Json<RequestEmailUpdateInput>,
) -> Response {
let client_ip = crate::rate_limit::extract_client_ip(&headers, None);
if !state.distributed_rate_limiter.check_rate_limit(
&format!("email_update:{}", client_ip),
5,
3_600_000,
).await {
if state.rate_limiters.email_update.check_key(&client_ip).is_err() {
warn!(ip = %client_ip, "Email update rate limit exceeded");
return (
StatusCode::TOO_MANY_REQUESTS,
Json(json!({
"error": "RateLimitExceeded",
"message": "Too many requests. Please try again later."
})),
).into_response();
}
}
let token = match crate::auth::extract_bearer_token_from_header(
headers.get("Authorization").and_then(|h| h.to_str().ok())
) {
@@ -135,6 +153,24 @@ pub async fn confirm_email(
headers: axum::http::HeaderMap,
Json(input): Json<ConfirmEmailInput>,
) -> Response {
let client_ip = crate::rate_limit::extract_client_ip(&headers, None);
if !state.distributed_rate_limiter.check_rate_limit(
&format!("confirm_email:{}", client_ip),
10,
60_000,
).await {
if state.rate_limiters.app_password.check_key(&client_ip).is_err() {
warn!(ip = %client_ip, "Confirm email rate limit exceeded");
return (
StatusCode::TOO_MANY_REQUESTS,
Json(json!({
"error": "RateLimitExceeded",
"message": "Too many requests. Please try again later."
})),
).into_response();
}
}
let token = match crate::auth::extract_bearer_token_from_header(
headers.get("Authorization").and_then(|h| h.to_str().ok())
) {
src/api/server/password.rs
+19
View File
@@ -124,8 +124,27 @@ pub struct ResetPasswordInput {
pub async fn reset_password(
State(state): State<AppState>,
headers: HeaderMap,
Json(input): Json<ResetPasswordInput>,
) -> Response {
let client_ip = extract_client_ip(&headers);
if !state.distributed_rate_limiter.check_rate_limit(
&format!("reset_password:{}", client_ip),
10,
60_000,
).await {
if state.rate_limiters.reset_password.check_key(&client_ip).is_err() {
warn!(ip = %client_ip, "Reset password rate limit exceeded");
return (
StatusCode::TOO_MANY_REQUESTS,
Json(json!({
"error": "RateLimitExceeded",
"message": "Too many requests. Please try again later."
})),
).into_response();
}
}
let token = input.token.trim();
let password = &input.password;
src/api/server/session.rs
+19
View File
@@ -72,6 +72,7 @@ pub async fn create_session(
{
Ok(Some(row)) => row,
Ok(None) => {
let _ = verify(&input.password, "$2b$12$LQv3c1yqBWVHxkd0LHAkCOYz6TtxMQJqhN8/X4.VTtYw1ZzQKZqmK");
warn!("User not found for login attempt");
return ApiError::AuthenticationFailedMsg("Invalid identifier or password".into()).into_response();
}
@@ -196,6 +197,24 @@ pub async fn refresh_session(
State(state): State<AppState>,
headers: axum::http::HeaderMap,
) -> Response {
let client_ip = crate::rate_limit::extract_client_ip(&headers, None);
if !state.distributed_rate_limiter.check_rate_limit(
&format!("refresh_session:{}", client_ip),
60,
60_000,
).await {
if state.rate_limiters.refresh_session.check_key(&client_ip).is_err() {
tracing::warn!(ip = %client_ip, "Refresh session rate limit exceeded");
return (
axum::http::StatusCode::TOO_MANY_REQUESTS,
axum::Json(serde_json::json!({
"error": "RateLimitExceeded",
"message": "Too many requests. Please try again later."
})),
).into_response();
}
}
let refresh_token = match crate::auth::extract_bearer_token_from_header(
headers.get("Authorization").and_then(|h| h.to_str().ok())
) {
src/cache/mod.rs
Vendored
+207
View File
@@ -0,0 +1,207 @@
use async_trait::async_trait;
use std::sync::Arc;
use std::time::Duration;
#[derive(Debug, thiserror::Error)]
pub enum CacheError {
#[error("Cache connection error: {0}")]
Connection(String),
#[error("Serialization error: {0}")]
Serialization(String),
}
#[async_trait]
pub trait Cache: Send + Sync {
async fn get(&self, key: &str) -> Option<String>;
async fn set(&self, key: &str, value: &str, ttl: Duration) -> Result<(), CacheError>;
async fn delete(&self, key: &str) -> Result<(), CacheError>;
}
#[derive(Clone)]
pub struct ValkeyCache {
conn: redis::aio::ConnectionManager,
}
impl ValkeyCache {
pub async fn new(url: &str) -> Result<Self, CacheError> {
let client = redis::Client::open(url)
.map_err(|e| CacheError::Connection(e.to_string()))?;
let manager = client
.get_connection_manager()
.await
.map_err(|e| CacheError::Connection(e.to_string()))?;
Ok(Self { conn: manager })
}
pub fn connection(&self) -> redis::aio::ConnectionManager {
self.conn.clone()
}
}
#[async_trait]
impl Cache for ValkeyCache {
async fn get(&self, key: &str) -> Option<String> {
let mut conn = self.conn.clone();
redis::cmd("GET")
.arg(key)
.query_async::<Option<String>>(&mut conn)
.await
.ok()
.flatten()
}
async fn set(&self, key: &str, value: &str, ttl: Duration) -> Result<(), CacheError> {
let mut conn = self.conn.clone();
redis::cmd("SET")
.arg(key)
.arg(value)
.arg("EX")
.arg(ttl.as_secs() as i64)
.query_async::<()>(&mut conn)
.await
.map_err(|e| CacheError::Connection(e.to_string()))
}
async fn delete(&self, key: &str) -> Result<(), CacheError> {
let mut conn = self.conn.clone();
redis::cmd("DEL")
.arg(key)
.query_async::<()>(&mut conn)
.await
.map_err(|e| CacheError::Connection(e.to_string()))
}
}
pub struct NoOpCache;
#[async_trait]
impl Cache for NoOpCache {
async fn get(&self, _key: &str) -> Option<String> {
None
}
async fn set(&self, _key: &str, _value: &str, _ttl: Duration) -> Result<(), CacheError> {
Ok(())
}
async fn delete(&self, _key: &str) -> Result<(), CacheError> {
Ok(())
}
}
#[async_trait]
pub trait DistributedRateLimiter: Send + Sync {
async fn check_rate_limit(&self, key: &str, limit: u32, window_ms: u64) -> bool;
}
#[derive(Clone)]
pub struct RedisRateLimiter {
conn: redis::aio::ConnectionManager,
}
impl RedisRateLimiter {
pub fn new(conn: redis::aio::ConnectionManager) -> Self {
Self { conn }
}
}
#[async_trait]
impl DistributedRateLimiter for RedisRateLimiter {
async fn check_rate_limit(&self, key: &str, limit: u32, window_ms: u64) -> bool {
let mut conn = self.conn.clone();
let full_key = format!("rl:{}", key);
let window_secs = ((window_ms + 999) / 1000).max(1) as i64;
let count: Result<i64, _> = redis::cmd("INCR")
.arg(&full_key)
.query_async(&mut conn)
.await;
let count = match count {
Ok(c) => c,
Err(e) => {
tracing::warn!("Redis rate limit INCR failed: {}. Allowing request.", e);
return true;
}
};
if count == 1 {
let _: Result<bool, redis::RedisError> = redis::cmd("EXPIRE")
.arg(&full_key)
.arg(window_secs)
.query_async(&mut conn)
.await;
}
count <= limit as i64
}
}
pub struct NoOpRateLimiter;
#[async_trait]
impl DistributedRateLimiter for NoOpRateLimiter {
async fn check_rate_limit(&self, _key: &str, _limit: u32, _window_ms: u64) -> bool {
true
}
}
pub enum CacheBackend {
Valkey(ValkeyCache),
NoOp,
}
impl CacheBackend {
pub fn rate_limiter(&self) -> Arc<dyn DistributedRateLimiter> {
match self {
CacheBackend::Valkey(cache) => {
Arc::new(RedisRateLimiter::new(cache.connection()))
}
CacheBackend::NoOp => Arc::new(NoOpRateLimiter),
}
}
}
#[async_trait]
impl Cache for CacheBackend {
async fn get(&self, key: &str) -> Option<String> {
match self {
CacheBackend::Valkey(c) => c.get(key).await,
CacheBackend::NoOp => None,
}
}
async fn set(&self, key: &str, value: &str, ttl: Duration) -> Result<(), CacheError> {
match self {
CacheBackend::Valkey(c) => c.set(key, value, ttl).await,
CacheBackend::NoOp => Ok(()),
}
}
async fn delete(&self, key: &str) -> Result<(), CacheError> {
match self {
CacheBackend::Valkey(c) => c.delete(key).await,
CacheBackend::NoOp => Ok(()),
}
}
}
pub async fn create_cache() -> (Arc<dyn Cache>, Arc<dyn DistributedRateLimiter>) {
match std::env::var("VALKEY_URL") {
Ok(url) => match ValkeyCache::new(&url).await {
Ok(cache) => {
tracing::info!("Connected to Valkey cache at {}", url);
let rate_limiter = Arc::new(RedisRateLimiter::new(cache.connection()));
(Arc::new(cache), rate_limiter)
}
Err(e) => {
tracing::warn!("Failed to connect to Valkey: {}. Running without cache.", e);
(Arc::new(NoOpCache), Arc::new(NoOpRateLimiter))
}
},
Err(_) => {
tracing::info!("VALKEY_URL not set. Running without cache.");
(Arc::new(NoOpCache), Arc::new(NoOpRateLimiter))
}
}
}
src/lib.rs
+1
View File
@@ -1,5 +1,6 @@
pub mod api;
pub mod auth;
pub mod cache;
pub mod circuit_breaker;
pub mod config;
pub mod crawlers;
src/oauth/endpoints/authorize.rs
+37 -1
View File
@@ -272,6 +272,27 @@ pub async fn authorize_post(
) -> Response {
let json_response = wants_json(&headers);
let client_ip = extract_client_ip(&headers);
if state.rate_limiters.oauth_authorize.check_key(&client_ip).is_err() {
tracing::warn!(ip = %client_ip, "OAuth authorize rate limit exceeded");
if json_response {
return (
axum::http::StatusCode::TOO_MANY_REQUESTS,
Json(serde_json::json!({
"error": "RateLimitExceeded",
"error_description": "Too many login attempts. Please try again later."
})),
).into_response();
}
return (
axum::http::StatusCode::TOO_MANY_REQUESTS,
Html(templates::error_page(
"RateLimitExceeded",
Some("Too many login attempts. Please try again later."),
)),
).into_response();
}
let request_data = match db::get_authorization_request(&state.db, &form.request_uri).await {
Ok(Some(data)) => data,
Ok(None) => {
@@ -357,7 +378,10 @@ pub async fn authorize_post(
.await
{
Ok(Some(u)) => u,
Ok(None) => return show_login_error("Invalid handle/email or password.", json_response),
Ok(None) => {
let _ = bcrypt::verify(&form.password, "$2b$12$LQv3c1yqBWVHxkd0LHAkCOYz6TtxMQJqhN8/X4.VTtYw1ZzQKZqmK");
return show_login_error("Invalid handle/email or password.", json_response);
}
Err(_) => return show_login_error("An error occurred. Please try again.", json_response),
};
@@ -736,6 +760,18 @@ pub async fn authorize_2fa_post(
headers: HeaderMap,
Form(form): Form<Authorize2faSubmit>,
) -> Response {
let client_ip = extract_client_ip(&headers);
if state.rate_limiters.oauth_authorize.check_key(&client_ip).is_err() {
tracing::warn!(ip = %client_ip, "OAuth 2FA rate limit exceeded");
return (
axum::http::StatusCode::TOO_MANY_REQUESTS,
Html(templates::error_page(
"RateLimitExceeded",
Some("Too many attempts. Please try again later."),
)),
).into_response();
}
let challenge = match db::get_2fa_challenge(&state.db, &form.request_uri).await {
Ok(Some(c)) => c,
Ok(None) => {
src/oauth/endpoints/par.rs
+14
View File
@@ -1,6 +1,7 @@
use axum::{
Form, Json,
extract::State,
http::HeaderMap,
};
use chrono::{Duration, Utc};
use serde::{Deserialize, Serialize};
@@ -49,8 +50,21 @@ pub struct ParResponse {
pub async fn pushed_authorization_request(
State(state): State<AppState>,
headers: HeaderMap,
Form(request): Form<ParRequest>,
) -> Result<Json<ParResponse>, OAuthError> {
let client_ip = crate::rate_limit::extract_client_ip(&headers, None);
if !state.distributed_rate_limiter.check_rate_limit(
&format!("oauth_par:{}", client_ip),
30,
60_000,
).await {
if state.rate_limiters.oauth_par.check_key(&client_ip).is_err() {
tracing::warn!(ip = %client_ip, "OAuth PAR rate limit exceeded");
return Err(OAuthError::RateLimited);
}
}
if request.response_type != "code" {
return Err(OAuthError::InvalidRequest(
"response_type must be 'code'".to_string(),
src/oauth/endpoints/token/introspect.rs
+33 -7
View File
@@ -1,6 +1,6 @@
use axum::{Form, Json};
use axum::extract::State;
use axum::http::StatusCode;
use axum::http::{HeaderMap, StatusCode};
use chrono::Utc;
use serde::{Deserialize, Serialize};
@@ -18,8 +18,21 @@ pub struct RevokeRequest {
pub async fn revoke_token(
State(state): State<AppState>,
headers: HeaderMap,
Form(request): Form<RevokeRequest>,
) -> Result<StatusCode, OAuthError> {
let client_ip = crate::rate_limit::extract_client_ip(&headers, None);
if !state.distributed_rate_limiter.check_rate_limit(
&format!("oauth_revoke:{}", client_ip),
30,
60_000,
).await {
if state.rate_limiters.oauth_introspect.check_key(&client_ip).is_err() {
tracing::warn!(ip = %client_ip, "OAuth revoke rate limit exceeded");
return Err(OAuthError::RateLimited);
}
}
if let Some(token) = &request.token {
if let Some((db_id, _)) = db::get_token_by_refresh_token(&state.db, token).await? {
db::delete_token_family(&state.db, db_id).await?;
@@ -67,8 +80,21 @@ pub struct IntrospectResponse {
pub async fn introspect_token(
State(state): State<AppState>,
headers: HeaderMap,
Form(request): Form<IntrospectRequest>,
) -> Json<IntrospectResponse> {
) -> Result<Json<IntrospectResponse>, OAuthError> {
let client_ip = crate::rate_limit::extract_client_ip(&headers, None);
if !state.distributed_rate_limiter.check_rate_limit(
&format!("oauth_introspect:{}", client_ip),
30,
60_000,
).await {
if state.rate_limiters.oauth_introspect.check_key(&client_ip).is_err() {
tracing::warn!(ip = %client_ip, "OAuth introspect rate limit exceeded");
return Err(OAuthError::RateLimited);
}
}
let inactive_response = IntrospectResponse {
active: false,
scope: None,
@@ -86,22 +112,22 @@ pub async fn introspect_token(
let token_info = match extract_token_claims(&request.token) {
Ok(info) => info,
Err(_) => return Json(inactive_response),
Err(_) => return Ok(Json(inactive_response)),
};
let token_data = match db::get_token_by_id(&state.db, &token_info.jti).await {
Ok(Some(data)) => data,
_ => return Json(inactive_response),
_ => return Ok(Json(inactive_response)),
};
if token_data.expires_at < Utc::now() {
return Json(inactive_response);
return Ok(Json(inactive_response));
}
let pds_hostname = std::env::var("PDS_HOSTNAME").unwrap_or_else(|_| "localhost".to_string());
let issuer = format!("https://{}", pds_hostname);
Json(IntrospectResponse {
Ok(Json(IntrospectResponse {
active: true,
scope: token_data.scope,
client_id: Some(token_data.client_id),
@@ -118,5 +144,5 @@ pub async fn introspect_token(
aud: Some(issuer.clone()),
iss: Some(issuer),
jti: Some(token_info.jti),
})
}))
}
src/oauth/error.rs
+4
View File
@@ -19,6 +19,7 @@ pub enum OAuthError {
InvalidDpopProof(String),
ExpiredToken(String),
InvalidToken(String),
RateLimited,
}
#[derive(Serialize)]
@@ -74,6 +75,9 @@ impl IntoResponse for OAuthError {
OAuthError::InvalidToken(msg) => {
(StatusCode::UNAUTHORIZED, "invalid_token", Some(msg))
}
OAuthError::RateLimited => {
(StatusCode::TOO_MANY_REQUESTS, "rate_limited", Some("Too many requests. Please try again later.".to_string()))
}
};
(
src/rate_limit.rs
+36 -1
View File
@@ -24,8 +24,15 @@ pub type GlobalRateLimiter = RateLimiter<NotKeyed, InMemoryState, DefaultClock>;
pub struct RateLimiters {
pub login: Arc<KeyedRateLimiter>,
pub oauth_token: Arc<KeyedRateLimiter>,
pub oauth_authorize: Arc<KeyedRateLimiter>,
pub password_reset: Arc<KeyedRateLimiter>,
pub account_creation: Arc<KeyedRateLimiter>,
pub refresh_session: Arc<KeyedRateLimiter>,
pub reset_password: Arc<KeyedRateLimiter>,
pub oauth_par: Arc<KeyedRateLimiter>,
pub oauth_introspect: Arc<KeyedRateLimiter>,
pub app_password: Arc<KeyedRateLimiter>,
pub email_update: Arc<KeyedRateLimiter>,
}
impl Default for RateLimiters {
@@ -43,12 +50,33 @@ impl RateLimiters {
oauth_token: Arc::new(RateLimiter::keyed(
Quota::per_minute(NonZeroU32::new(30).unwrap())
)),
oauth_authorize: Arc::new(RateLimiter::keyed(
Quota::per_minute(NonZeroU32::new(10).unwrap())
)),
password_reset: Arc::new(RateLimiter::keyed(
Quota::per_hour(NonZeroU32::new(5).unwrap())
)),
account_creation: Arc::new(RateLimiter::keyed(
Quota::per_hour(NonZeroU32::new(10).unwrap())
)),
refresh_session: Arc::new(RateLimiter::keyed(
Quota::per_minute(NonZeroU32::new(60).unwrap())
)),
reset_password: Arc::new(RateLimiter::keyed(
Quota::per_minute(NonZeroU32::new(10).unwrap())
)),
oauth_par: Arc::new(RateLimiter::keyed(
Quota::per_minute(NonZeroU32::new(30).unwrap())
)),
oauth_introspect: Arc::new(RateLimiter::keyed(
Quota::per_minute(NonZeroU32::new(30).unwrap())
)),
app_password: Arc::new(RateLimiter::keyed(
Quota::per_minute(NonZeroU32::new(10).unwrap())
)),
email_update: Arc::new(RateLimiter::keyed(
Quota::per_hour(NonZeroU32::new(5).unwrap())
)),
}
}
@@ -66,6 +94,13 @@ impl RateLimiters {
self
}
pub fn with_oauth_authorize_limit(mut self, per_minute: u32) -> Self {
self.oauth_authorize = Arc::new(RateLimiter::keyed(
Quota::per_minute(NonZeroU32::new(per_minute).unwrap_or(NonZeroU32::new(10).unwrap()))
));
self
}
pub fn with_password_reset_limit(mut self, per_hour: u32) -> Self {
self.password_reset = Arc::new(RateLimiter::keyed(
Quota::per_hour(NonZeroU32::new(per_hour).unwrap_or(NonZeroU32::new(5).unwrap()))
@@ -81,7 +116,7 @@ impl RateLimiters {
}
}
fn extract_client_ip(headers: &HeaderMap, addr: Option<SocketAddr>) -> String {
pub fn extract_client_ip(headers: &HeaderMap, addr: Option<SocketAddr>) -> String {
if let Some(forwarded) = headers.get("x-forwarded-for") {
if let Ok(value) = forwarded.to_str() {
if let Some(first_ip) = value.split(',').next() {
src/state.rs
+6
View File
@@ -1,3 +1,4 @@
use crate::cache::{Cache, DistributedRateLimiter, create_cache};
use crate::circuit_breaker::CircuitBreakers;
use crate::config::AuthConfig;
use crate::rate_limit::RateLimiters;
@@ -16,6 +17,8 @@ pub struct AppState {
pub firehose_tx: broadcast::Sender<SequencedEvent>,
pub rate_limiters: Arc<RateLimiters>,
pub circuit_breakers: Arc<CircuitBreakers>,
pub cache: Arc<dyn Cache>,
pub distributed_rate_limiter: Arc<dyn DistributedRateLimiter>,
}
impl AppState {
@@ -27,6 +30,7 @@ impl AppState {
let (firehose_tx, _) = broadcast::channel(1000);
let rate_limiters = Arc::new(RateLimiters::new());
let circuit_breakers = Arc::new(CircuitBreakers::new());
let (cache, distributed_rate_limiter) = create_cache().await;
Self {
db,
block_store,
@@ -34,6 +38,8 @@ impl AppState {
firehose_tx,
rate_limiters,
circuit_breakers,
cache,
distributed_rate_limiter,
}
}
tests/oauth_security.rs
+64
View File
@@ -1447,3 +1447,67 @@ async fn test_security_revoked_token_rejected() {
let introspect_body: Value = introspect_res.json().await.unwrap();
assert_eq!(introspect_body["active"], false, "Revoked token should be inactive");
}
#[tokio::test]
async fn test_security_oauth_authorize_rate_limiting() {
let url = base_url().await;
let http_client = no_redirect_client();
let ts = Utc::now().timestamp_nanos_opt().unwrap_or(0);
let unique_ip = format!("10.{}.{}.{}", (ts >> 16) & 0xFF, (ts >> 8) & 0xFF, ts & 0xFF);
let redirect_uri = "https://example.com/rate-limit-callback";
let mock_client = setup_mock_client_metadata(redirect_uri).await;
let client_id = mock_client.uri();
let (_, code_challenge) = generate_pkce();
let client_for_par = client();
let par_body: Value = client_for_par
.post(format!("{}/oauth/par", url))
.form(&[
("response_type", "code"),
("client_id", &client_id),
("redirect_uri", redirect_uri),
("code_challenge", &code_challenge),
("code_challenge_method", "S256"),
])
.send()
.await
.unwrap()
.json()
.await
.unwrap();
let request_uri = par_body["request_uri"].as_str().unwrap();
let mut rate_limited_count = 0;
let mut other_count = 0;
for _ in 0..15 {
let res = http_client
.post(format!("{}/oauth/authorize", url))
.header("X-Forwarded-For", &unique_ip)
.form(&[
("request_uri", request_uri),
("username", "nonexistent_user"),
("password", "wrong_password"),
("remember_device", "false"),
])
.send()
.await
.unwrap();
match res.status() {
StatusCode::TOO_MANY_REQUESTS => rate_limited_count += 1,
_ => other_count += 1,
}
}
assert!(
rate_limited_count > 0,
"Expected at least one rate-limited response after 15 OAuth authorize attempts. Got {} other and {} rate limited.",
other_count,
rate_limited_count
);
}
tests/rate_limit.rs
+228
View File
@@ -0,0 +1,228 @@
mod common;
use common::{base_url, client};
use reqwest::StatusCode;
use serde_json::json;
#[tokio::test]
async fn test_login_rate_limiting() {
let client = client();
let url = format!("{}/xrpc/com.atproto.server.createSession", base_url().await);
let payload = json!({
"identifier": "nonexistent_user_for_rate_limit_test",
"password": "wrongpassword"
});
let mut rate_limited_count = 0;
let mut auth_failed_count = 0;
for _ in 0..15 {
let res = client
.post(&url)
.json(&payload)
.send()
.await
.expect("Request failed");
match res.status() {
StatusCode::TOO_MANY_REQUESTS => {
rate_limited_count += 1;
}
StatusCode::UNAUTHORIZED => {
auth_failed_count += 1;
}
status => {
panic!("Unexpected status: {}", status);
}
}
}
assert!(
rate_limited_count > 0,
"Expected at least one rate-limited response after 15 login attempts. Got {} auth failures and {} rate limits.",
auth_failed_count,
rate_limited_count
);
}
#[tokio::test]
async fn test_password_reset_rate_limiting() {
let client = client();
let url = format!(
"{}/xrpc/com.atproto.server.requestPasswordReset",
base_url().await
);
let mut rate_limited_count = 0;
let mut success_count = 0;
for i in 0..8 {
let payload = json!({
"email": format!("ratelimit_test_{}@example.com", i)
});
let res = client
.post(&url)
.json(&payload)
.send()
.await
.expect("Request failed");
match res.status() {
StatusCode::TOO_MANY_REQUESTS => {
rate_limited_count += 1;
}
StatusCode::OK => {
success_count += 1;
}
status => {
panic!("Unexpected status: {} - {:?}", status, res.text().await);
}
}
}
assert!(
rate_limited_count > 0,
"Expected rate limiting after {} password reset requests. Got {} successes.",
success_count + rate_limited_count,
success_count
);
}
#[tokio::test]
async fn test_account_creation_rate_limiting() {
let client = client();
let url = format!(
"{}/xrpc/com.atproto.server.createAccount",
base_url().await
);
let mut rate_limited_count = 0;
let mut other_count = 0;
for i in 0..15 {
let unique_id = uuid::Uuid::new_v4();
let payload = json!({
"handle": format!("ratelimit_{}_{}", i, unique_id),
"email": format!("ratelimit_{}_{}@example.com", i, unique_id),
"password": "testpassword123"
});
let res = client
.post(&url)
.json(&payload)
.send()
.await
.expect("Request failed");
match res.status() {
StatusCode::TOO_MANY_REQUESTS => {
rate_limited_count += 1;
}
_ => {
other_count += 1;
}
}
}
assert!(
rate_limited_count > 0,
"Expected rate limiting after account creation attempts. Got {} other responses and {} rate limits.",
other_count,
rate_limited_count
);
}
#[tokio::test]
async fn test_valkey_connection() {
if std::env::var("VALKEY_URL").is_err() {
println!("VALKEY_URL not set, skipping Valkey connection test");
return;
}
let valkey_url = std::env::var("VALKEY_URL").unwrap();
let client = redis::Client::open(valkey_url.as_str()).expect("Failed to create Redis client");
let mut conn = client
.get_multiplexed_async_connection()
.await
.expect("Failed to connect to Valkey");
let pong: String = redis::cmd("PING")
.query_async(&mut conn)
.await
.expect("PING failed");
assert_eq!(pong, "PONG");
let _: () = redis::cmd("SET")
.arg("test_key")
.arg("test_value")
.arg("EX")
.arg(10)
.query_async(&mut conn)
.await
.expect("SET failed");
let value: String = redis::cmd("GET")
.arg("test_key")
.query_async(&mut conn)
.await
.expect("GET failed");
assert_eq!(value, "test_value");
let _: () = redis::cmd("DEL")
.arg("test_key")
.query_async(&mut conn)
.await
.expect("DEL failed");
}
#[tokio::test]
async fn test_distributed_rate_limiter_directly() {
if std::env::var("VALKEY_URL").is_err() {
println!("VALKEY_URL not set, skipping distributed rate limiter test");
return;
}
use bspds::cache::{DistributedRateLimiter, RedisRateLimiter};
let valkey_url = std::env::var("VALKEY_URL").unwrap();
let client = redis::Client::open(valkey_url.as_str()).expect("Failed to create Redis client");
let conn = client
.get_connection_manager()
.await
.expect("Failed to get connection manager");
let rate_limiter = RedisRateLimiter::new(conn);
let test_key = format!("test_rate_limit_{}", uuid::Uuid::new_v4());
let limit = 5;
let window_ms = 60_000;
for i in 0..limit {
let allowed = rate_limiter
.check_rate_limit(&test_key, limit, window_ms)
.await;
assert!(
allowed,
"Request {} should have been allowed (limit: {})",
i + 1,
limit
);
}
let allowed = rate_limiter
.check_rate_limit(&test_key, limit, window_ms)
.await;
assert!(
!allowed,
"Request {} should have been rate limited (limit: {})",
limit + 1,
limit
);
let allowed = rate_limiter
.check_rate_limit(&test_key, limit, window_ms)
.await;
assert!(!allowed, "Subsequent request should also be rate limited");
}