Microservice to bring 2FA to self hosted PDSes

Create account rate limiting

authored by baileytownsend.dev and committed by Tangled 94e9b2ee 239f7ad5

Changed files
+50 -21
examples
src
+1
examples/Caddyfile
··· 14 14 path /xrpc/com.atproto.server.getSession 15 15 path /xrpc/com.atproto.server.updateEmail 16 16 path /xrpc/com.atproto.server.createSession 17 + path /xrpc/com.atproto.server.createAccount 17 18 path /@atproto/oauth-provider/~api/sign-in 18 19 } 19 20
+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
··· 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 + }