Microservice to bring 2FA to self hosted PDSes
Create account rate limiting
authored by
Bailey Townsend
and committed by
Tangled
94e9b2ee
239f7ad5
+50
-21
+1
examples/Caddyfile
+1
examples/Caddyfile
+25
-21
src/main.rs
+25
-21
src/main.rs
···
1
1
#![warn(clippy::unwrap_used)]
2
2
use crate::oauth_provider::sign_in;
3
-
use crate::xrpc::com_atproto_server::{create_session, get_session, update_email};
3
+
use crate::xrpc::com_atproto_server::{create_account, create_session, get_session, update_email};
4
4
use axum::body::Body;
5
5
use axum::handler::Handler;
6
6
use axum::http::{Method, header};
···
20
20
use std::{env, net::SocketAddr};
21
21
use tower_governor::GovernorLayer;
22
22
use tower_governor::governor::{GovernorConfig, GovernorConfigBuilder};
23
-
use tower_governor::key_extractor::PeerIpKeyExtractor;
24
23
use tower_http::compression::CompressionLayer;
25
24
use tower_http::cors::{Any, CorsLayer};
26
25
use tracing::log;
···
92
91
let pds_env_location =
93
92
env::var("PDS_ENV_LOCATION").unwrap_or_else(|_| "/pds/pds.env".to_string());
94
93
95
-
dotenvy::from_path(Path::new(&pds_env_location))?;
94
+
let result_of_finding_pds_env = dotenvy::from_path(Path::new(&pds_env_location));
95
+
if let Err(e) = result_of_finding_pds_env {
96
+
log::error!(
97
+
"Error loading pds.env file (ignore if you loaded your variables in the environment somehow else): {e}"
98
+
);
99
+
}
96
100
let pds_root = env::var("PDS_DATA_DIRECTORY")?;
97
101
let account_db_url = format!("{pds_root}/account.sqlite");
98
102
···
182
186
env::var("GATEKEEPER_CREATE_ACCOUNT_PER_SECOND").ok();
183
187
let create_account_limiter_burst: Option<String> =
184
188
env::var("GATEKEEPER_CREATE_ACCOUNT_BURST").ok();
185
-
let mut create_account_governor_conf = None;
186
189
187
-
if create_account_governor_conf.is_some() && create_account_limiter_time.is_some() {
190
+
//Default should be 608 requests per 5 minutes, PDS is 300 per 500 so will never hit it ideally
191
+
let mut create_account_governor_conf = GovernorConfigBuilder::default();
192
+
if create_account_limiter_time.is_some() {
188
193
let time = create_account_limiter_time
189
194
.expect("GATEKEEPER_CREATE_ACCOUNT_PER_SECOND not set")
190
195
.parse::<u64>()
191
196
.expect("GATEKEEPER_CREATE_ACCOUNT_PER_SECOND must be a valid integer");
197
+
create_account_governor_conf.per_second(time);
198
+
}
199
+
200
+
if create_account_limiter_burst.is_some() {
192
201
let burst = create_account_limiter_burst
193
202
.expect("GATEKEEPER_CREATE_ACCOUNT_BURST not set")
194
203
.parse::<u32>()
195
204
.expect("GATEKEEPER_CREATE_ACCOUNT_BURST must be a valid integer");
196
-
197
-
create_account_governor_conf = Some(
198
-
GovernorConfigBuilder::default()
199
-
.per_second(time)
200
-
.burst_size(burst)
201
-
.finish()
202
-
.expect("failed to create governor config for create account. this should not happen and is a bug"),
203
-
)
205
+
create_account_governor_conf.burst_size(burst);
204
206
}
205
207
208
+
let create_account_governor_conf = create_account_governor_conf.finish().expect(
209
+
"failed to create governor config for create account. this should not happen and is a bug",
210
+
);
211
+
206
212
let create_session_governor_limiter = create_session_governor_conf.limiter().clone();
207
213
let sign_in_governor_limiter = sign_in_governor_conf.limiter().clone();
208
-
let create_account_governor_limiter = match create_account_governor_conf {
209
-
None => None,
210
-
Some(conf) => Some(conf.limiter().clone()),
211
-
};
214
+
let create_account_governor_limiter = create_account_governor_conf.limiter().clone();
212
215
213
216
let interval = Duration::from_secs(60);
214
217
// a separate background task to clean up
···
217
220
std::thread::sleep(interval);
218
221
create_session_governor_limiter.retain_recent();
219
222
sign_in_governor_limiter.retain_recent();
220
-
if let Some(ref limiter) = create_account_governor_limiter {
221
-
limiter.retain_recent();
222
-
}
223
+
create_account_governor_limiter.retain_recent();
223
224
}
224
225
});
225
226
···
243
244
"/xrpc/com.atproto.server.createSession",
244
245
post(create_session.layer(GovernorLayer::new(create_session_governor_conf))),
245
246
)
246
-
.route("/xrpc/com.atproto.server.createAccount")
247
+
.route(
248
+
"/xrpc/com.atproto.server.createAccount",
249
+
post(create_account).layer(GovernorLayer::new(create_account_governor_conf)),
250
+
)
247
251
.layer(CompressionLayer::new())
248
252
.layer(cors)
249
253
.with_state(state);
+24
src/xrpc/com_atproto_server.rs
+24
src/xrpc/com_atproto_server.rs
···
264
264
ProxiedResult::Passthrough(resp) => Ok(resp),
265
265
}
266
266
}
267
+
268
+
pub async fn create_account(
269
+
State(state): State<AppState>,
270
+
mut req: Request,
271
+
) -> Result<Response<Body>, StatusCode> {
272
+
let uri = format!(
273
+
"{}{}",
274
+
state.pds_base_url, "/xrpc/com.atproto.server.createAccount"
275
+
);
276
+
277
+
// Rewrite the URI to point at the upstream PDS; keep headers, method, and body intact
278
+
*req.uri_mut() = uri
279
+
.parse()
280
+
.map_err(|_| StatusCode::BAD_REQUEST)?;
281
+
282
+
let proxied = state
283
+
.reverse_proxy_client
284
+
.request(req)
285
+
.await
286
+
.map_err(|_| StatusCode::BAD_REQUEST)?
287
+
.into_response();
288
+
289
+
Ok(proxied)
290
+
}