nel.pet
forked from
lewis.moe/bspds-sandbox
I've been saying "PDSes seem easy enough, they're what, some CRUD to a db? I can do that in my sleep". well i'm sleeping rn so let's go
1mod common;
2use common::{base_url, client};
3use reqwest::StatusCode;
4use serde_json::json;
5
6#[tokio::test]
7#[ignore = "rate limiting is disabled in test environment"]
8async fn test_login_rate_limiting() {
9 let client = client();
10 let url = format!("{}/xrpc/com.atproto.server.createSession", base_url().await);
11 let payload = json!({
12 "identifier": "nonexistent-user-for-rate-limit-test",
13 "password": "wrongpassword"
14 });
15 let mut rate_limited_count = 0;
16 let mut auth_failed_count = 0;
17 for _ in 0..15 {
18 let res = client
19 .post(&url)
20 .json(&payload)
21 .send()
22 .await
23 .expect("Request failed");
24 match res.status() {
25 StatusCode::TOO_MANY_REQUESTS => {
26 rate_limited_count += 1;
27 }
28 StatusCode::UNAUTHORIZED => {
29 auth_failed_count += 1;
30 }
31 status => {
32 panic!("Unexpected status: {}", status);
33 }
34 }
35 }
36 assert!(
37 rate_limited_count > 0,
38 "Expected at least one rate-limited response after 15 login attempts. Got {} auth failures and {} rate limits.",
39 auth_failed_count,
40 rate_limited_count
41 );
42}
43
44#[tokio::test]
45#[ignore = "rate limiting is disabled in test environment"]
46async fn test_password_reset_rate_limiting() {
47 let client = client();
48 let url = format!(
49 "{}/xrpc/com.atproto.server.requestPasswordReset",
50 base_url().await
51 );
52 let mut rate_limited_count = 0;
53 let mut success_count = 0;
54 for i in 0..8 {
55 let payload = json!({
56 "email": format!("ratelimit-test_{}@example.com", i)
57 });
58 let res = client
59 .post(&url)
60 .json(&payload)
61 .send()
62 .await
63 .expect("Request failed");
64 match res.status() {
65 StatusCode::TOO_MANY_REQUESTS => {
66 rate_limited_count += 1;
67 }
68 StatusCode::OK => {
69 success_count += 1;
70 }
71 status => {
72 panic!("Unexpected status: {} - {:?}", status, res.text().await);
73 }
74 }
75 }
76 assert!(
77 rate_limited_count > 0,
78 "Expected rate limiting after {} password reset requests. Got {} successes.",
79 success_count + rate_limited_count,
80 success_count
81 );
82}
83
84#[tokio::test]
85#[ignore = "rate limiting is disabled in test environment"]
86async fn test_account_creation_rate_limiting() {
87 let client = client();
88 let url = format!("{}/xrpc/com.atproto.server.createAccount", base_url().await);
89 let mut rate_limited_count = 0;
90 let mut other_count = 0;
91 for i in 0..15 {
92 let unique_id = uuid::Uuid::new_v4();
93 let payload = json!({
94 "handle": format!("ratelimit-{}_{}", i, unique_id),
95 "email": format!("ratelimit-{}_{}@example.com", i, unique_id),
96 "password": "Testpass123!"
97 });
98 let res = client
99 .post(&url)
100 .json(&payload)
101 .send()
102 .await
103 .expect("Request failed");
104 match res.status() {
105 StatusCode::TOO_MANY_REQUESTS => {
106 rate_limited_count += 1;
107 }
108 _ => {
109 other_count += 1;
110 }
111 }
112 }
113 assert!(
114 rate_limited_count > 0,
115 "Expected rate limiting after account creation attempts. Got {} other responses and {} rate limits.",
116 other_count,
117 rate_limited_count
118 );
119}
120
121#[tokio::test]
122async fn test_valkey_connection() {
123 if std::env::var("VALKEY_URL").is_err() {
124 println!("VALKEY_URL not set, skipping Valkey connection test");
125 return;
126 }
127 let valkey_url = std::env::var("VALKEY_URL").unwrap();
128 let client = redis::Client::open(valkey_url.as_str()).expect("Failed to create Redis client");
129 let mut conn = client
130 .get_multiplexed_async_connection()
131 .await
132 .expect("Failed to connect to Valkey");
133 let pong: String = redis::cmd("PING")
134 .query_async(&mut conn)
135 .await
136 .expect("PING failed");
137 assert_eq!(pong, "PONG");
138 let _: () = redis::cmd("SET")
139 .arg("test_key")
140 .arg("test_value")
141 .arg("EX")
142 .arg(10)
143 .query_async(&mut conn)
144 .await
145 .expect("SET failed");
146 let value: String = redis::cmd("GET")
147 .arg("test_key")
148 .query_async(&mut conn)
149 .await
150 .expect("GET failed");
151 assert_eq!(value, "test_value");
152 let _: () = redis::cmd("DEL")
153 .arg("test_key")
154 .query_async(&mut conn)
155 .await
156 .expect("DEL failed");
157}
158
159#[tokio::test]
160async fn test_distributed_rate_limiter_directly() {
161 if std::env::var("VALKEY_URL").is_err() {
162 println!("VALKEY_URL not set, skipping distributed rate limiter test");
163 return;
164 }
165 use tranquil_pds::cache::{DistributedRateLimiter, RedisRateLimiter};
166 let valkey_url = std::env::var("VALKEY_URL").unwrap();
167 let client = redis::Client::open(valkey_url.as_str()).expect("Failed to create Redis client");
168 let conn = client
169 .get_connection_manager()
170 .await
171 .expect("Failed to get connection manager");
172 let rate_limiter = RedisRateLimiter::new(conn);
173 let test_key = format!("test_rate_limit_{}", uuid::Uuid::new_v4());
174 let limit = 5;
175 let window_ms = 60_000;
176 for i in 0..limit {
177 let allowed = rate_limiter
178 .check_rate_limit(&test_key, limit, window_ms)
179 .await;
180 assert!(
181 allowed,
182 "Request {} should have been allowed (limit: {})",
183 i + 1,
184 limit
185 );
186 }
187 let allowed = rate_limiter
188 .check_rate_limit(&test_key, limit, window_ms)
189 .await;
190 assert!(
191 !allowed,
192 "Request {} should have been rate limited (limit: {})",
193 limit + 1,
194 limit
195 );
196 let allowed = rate_limiter
197 .check_rate_limit(&test_key, limit, window_ms)
198 .await;
199 assert!(!allowed, "Subsequent request should also be rate limited");
200}