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