diff --git a/.sqlx/query-76c6ef1d5395105a0cdedb27ca321c9e3eae1ce87c223b706ed81ebf973875f3.json b/.sqlx/query-08c08b0644d79d5de72f3500dd7dbb8827af340e3c04fec9a5c28aeff46e0c97.json similarity index 59% rename from .sqlx/query-76c6ef1d5395105a0cdedb27ca321c9e3eae1ce87c223b706ed81ebf973875f3.json rename to .sqlx/query-08c08b0644d79d5de72f3500dd7dbb8827af340e3c04fec9a5c28aeff46e0c97.json index a63e9a8..8b18439 100644 --- a/.sqlx/query-76c6ef1d5395105a0cdedb27ca321c9e3eae1ce87c223b706ed81ebf973875f3.json +++ b/.sqlx/query-08c08b0644d79d5de72f3500dd7dbb8827af340e3c04fec9a5c28aeff46e0c97.json @@ -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" } diff --git a/.sqlx/query-e223898d53602c1c8b23eb08a4b96cf20ac349d1fa4e91334b225d3069209dcf.json b/.sqlx/query-e223898d53602c1c8b23eb08a4b96cf20ac349d1fa4e91334b225d3069209dcf.json new file mode 100644 index 0000000..22bbecc --- /dev/null +++ b/.sqlx/query-e223898d53602c1c8b23eb08a4b96cf20ac349d1fa4e91334b225d3069209dcf.json @@ -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" +} diff --git a/Cargo.lock b/Cargo.lock index 6806171..fe625cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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", diff --git a/Cargo.toml b/Cargo.toml index 336722f..03d00b9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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 = [] diff --git a/TODO.md b/TODO.md index ce5d325..30139bd 100644 --- a/TODO.md +++ b/TODO.md @@ -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. diff --git a/docker-compose.yaml b/docker-compose.yaml index 999d86d..c858f11 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -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: diff --git a/scripts/test-infra.sh b/scripts/test-infra.sh index beb20cd..19163c4 100755 --- a/scripts/test-infra.sh +++ b/scripts/test-infra.sh @@ -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" diff --git a/src/api/admin/account/delete.rs b/src/api/admin/account/delete.rs index e6c0c21..0669772 100644 --- a/src/api/admin/account/delete.rs +++ b/src/api/admin/account/delete.rs @@ -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() } diff --git a/src/api/admin/account/update.rs b/src/api/admin/account/update.rs index 5c7d194..8cdfd67 100644 --- a/src/api/admin/account/update.rs +++ b/src/api/admin/account/update.rs @@ -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) => { diff --git a/src/api/admin/status.rs b/src/api/admin/status.rs index c936353..e4e843e 100644 --- a/src/api/admin/status.rs +++ b/src/api/admin/status.rs @@ -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!({ diff --git a/src/api/identity/did.rs b/src/api/identity/did.rs index 57cfbca..19db902 100644 --- a/src/api/identity/did.rs +++ b/src/api/identity/did.rs @@ -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); ( diff --git a/src/api/repo/record/batch.rs b/src/api/repo/record/batch.rs index 86c87b4..e0c242f 100644 --- a/src/api/repo/record/batch.rs +++ b/src/api/repo/record/batch.rs @@ -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(); diff --git a/src/api/repo/record/mod.rs b/src/api/repo/record/mod.rs index 1eee4c2..0f1ce5b 100644 --- a/src/api/repo/record/mod.rs +++ b/src/api/repo/record/mod.rs @@ -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; diff --git a/src/api/repo/record/validation.rs b/src/api/repo/record/validation.rs new file mode 100644 index 0000000..20efd4b --- /dev/null +++ b/src/api/repo/record/validation.rs @@ -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()), + } +} diff --git a/src/api/repo/record/write.rs b/src/api/repo/record/write.rs index f8d7837..ff827b2 100644 --- a/src/api/repo/record/write.rs +++ b/src/api/repo/record/write.rs @@ -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; } } diff --git a/src/api/server/account_status.rs b/src/api/server/account_status.rs index b877563..0ef4d84 100644 --- a/src/api/server/account_status.rs +++ b/src/api/server/account_status.rs @@ -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() } diff --git a/src/api/server/app_password.rs b/src/api/server/app_password.rs index 8dff8d8..da69b86 100644 --- a/src/api/server/app_password.rs +++ b/src/api/server/app_password.rs @@ -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, + headers: HeaderMap, BearerAuth(auth_user): BearerAuth, Json(input): Json, ) -> 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(), diff --git a/src/api/server/email.rs b/src/api/server/email.rs index 3905632..382124e 100644 --- a/src/api/server/email.rs +++ b/src/api/server/email.rs @@ -26,6 +26,24 @@ pub async fn request_email_update( headers: axum::http::HeaderMap, Json(input): Json, ) -> 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, ) -> 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()) ) { diff --git a/src/api/server/password.rs b/src/api/server/password.rs index 1ae645d..30ff31a 100644 --- a/src/api/server/password.rs +++ b/src/api/server/password.rs @@ -124,8 +124,27 @@ pub struct ResetPasswordInput { pub async fn reset_password( State(state): State, + headers: HeaderMap, Json(input): Json, ) -> 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; diff --git a/src/api/server/session.rs b/src/api/server/session.rs index f61f748..4129e6f 100644 --- a/src/api/server/session.rs +++ b/src/api/server/session.rs @@ -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, 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()) ) { diff --git a/src/cache/mod.rs b/src/cache/mod.rs new file mode 100644 index 0000000..017732b --- /dev/null +++ b/src/cache/mod.rs @@ -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; + 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 { + 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 { + let mut conn = self.conn.clone(); + redis::cmd("GET") + .arg(key) + .query_async::>(&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 { + 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 = 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 = 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 { + 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 { + 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, Arc) { + 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)) + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 15ea6ce..3ced835 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ pub mod api; pub mod auth; +pub mod cache; pub mod circuit_breaker; pub mod config; pub mod crawlers; diff --git a/src/oauth/endpoints/authorize.rs b/src/oauth/endpoints/authorize.rs index 39499bd..d70f15a 100644 --- a/src/oauth/endpoints/authorize.rs +++ b/src/oauth/endpoints/authorize.rs @@ -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, ) -> 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) => { diff --git a/src/oauth/endpoints/par.rs b/src/oauth/endpoints/par.rs index d865cf4..8267258 100644 --- a/src/oauth/endpoints/par.rs +++ b/src/oauth/endpoints/par.rs @@ -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, + headers: HeaderMap, Form(request): Form, ) -> Result, 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(), diff --git a/src/oauth/endpoints/token/introspect.rs b/src/oauth/endpoints/token/introspect.rs index c8a5d97..dc8d2be 100644 --- a/src/oauth/endpoints/token/introspect.rs +++ b/src/oauth/endpoints/token/introspect.rs @@ -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, + headers: HeaderMap, Form(request): Form, ) -> Result { + 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, + headers: HeaderMap, Form(request): Form, -) -> Json { +) -> Result, 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), - }) + })) } diff --git a/src/oauth/error.rs b/src/oauth/error.rs index dd5fd9a..2c00eba 100644 --- a/src/oauth/error.rs +++ b/src/oauth/error.rs @@ -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())) + } }; ( diff --git a/src/rate_limit.rs b/src/rate_limit.rs index ef196b3..e139f99 100644 --- a/src/rate_limit.rs +++ b/src/rate_limit.rs @@ -24,8 +24,15 @@ pub type GlobalRateLimiter = RateLimiter; pub struct RateLimiters { pub login: Arc, pub oauth_token: Arc, + pub oauth_authorize: Arc, pub password_reset: Arc, pub account_creation: Arc, + pub refresh_session: Arc, + pub reset_password: Arc, + pub oauth_par: Arc, + pub oauth_introspect: Arc, + pub app_password: Arc, + pub email_update: Arc, } 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) -> String { +pub fn extract_client_ip(headers: &HeaderMap, addr: Option) -> 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() { diff --git a/src/state.rs b/src/state.rs index 0b366ec..4355dfa 100644 --- a/src/state.rs +++ b/src/state.rs @@ -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, pub rate_limiters: Arc, pub circuit_breakers: Arc, + pub cache: Arc, + pub distributed_rate_limiter: Arc, } 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, } } diff --git a/tests/oauth_security.rs b/tests/oauth_security.rs index 347cd12..86d2566 100644 --- a/tests/oauth_security.rs +++ b/tests/oauth_security.rs @@ -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 + ); +} diff --git a/tests/rate_limit.rs b/tests/rate_limit.rs new file mode 100644 index 0000000..3dc2df1 --- /dev/null +++ b/tests/rate_limit.rs @@ -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"); +}