use std::collections::HashMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};
use axum::Json;
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode};
use serde::Deserialize;
use serde_json::Value;
use super::auth::{AppState, sha256_hex};
const MAX_ATTEMPTS: usize = 10;
const WINDOW: Duration = Duration::from_hours(1);
static ATTEMPTS: Mutex<Option<HashMap<String, Vec<Instant>>>> = Mutex::new(None);
fn rate_limited(ip_hash: Option<String>, now: Instant) -> bool {
let Some(key) = ip_hash else {
return false;
};
let mut guard = match ATTEMPTS.lock() {
Ok(g) => g,
Err(poisoned) => poisoned.into_inner(),
};
let map = guard.get_or_insert_with(HashMap::new);
map.retain(|_, hits| {
hits.retain(|t| now.duration_since(*t) < WINDOW);
!hits.is_empty()
});
let hits = map.entry(key).or_default();
if hits.len() >= MAX_ATTEMPTS {
return true;
}
hits.push(now);
false
}
fn client_ip_hash(headers: &HeaderMap, salt: &str) -> Option<String> {
let ip = headers
.get("x-forwarded-for")
.and_then(|v| v.to_str().ok())
.and_then(|s| s.split(',').next())
.map(str::trim)
.filter(|s| !s.is_empty())
.or_else(|| {
headers
.get("x-real-ip")
.and_then(|v| v.to_str().ok())
.map(str::trim)
.filter(|s| !s.is_empty())
})?;
Some(sha256_hex(&format!("{salt}:{ip}")))
}
fn valid_code_shape(code: &str) -> bool {
code.len() == 64 && code.bytes().all(|b| b.is_ascii_hexdigit())
}
#[derive(Deserialize)]
pub(super) struct JoinBody {
code: String,
}
pub(super) async fn post_team_join(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<JoinBody>,
) -> Result<Json<Value>, (StatusCode, String)> {
let code = body.code.trim().to_ascii_lowercase();
if !valid_code_shape(&code) {
return Err((
StatusCode::BAD_REQUEST,
"that does not look like an invite code".into(),
));
}
if rate_limited(
client_ip_hash(&headers, &state.cfg.ip_hash_salt),
Instant::now(),
) {
return Err((
StatusCode::TOO_MANY_REQUESTS,
"too many attempts — try again later".into(),
));
}
let (status, json) = super::billing_edge::forward_invite_redeem(&state.cfg, &code).await?;
if status.is_success() {
return Ok(Json(json));
}
if status == StatusCode::NOT_FOUND {
return Err((
StatusCode::NOT_FOUND,
"this invite link is invalid, expired, or already used".into(),
));
}
let msg = json
.get("error")
.and_then(Value::as_str)
.or_else(|| json.get("message").and_then(Value::as_str))
.unwrap_or("could not redeem the invite")
.to_string();
Err((status, msg))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn code_shape_is_64_hex() {
assert!(valid_code_shape(&"a".repeat(64)));
assert!(valid_code_shape(&"A".repeat(64)));
assert!(!valid_code_shape(&"a".repeat(63)));
assert!(!valid_code_shape(&"a".repeat(65)));
assert!(!valid_code_shape(&"g".repeat(64)));
assert!(!valid_code_shape(""));
}
#[test]
fn rate_limit_counts_per_ip_within_window() {
let now = Instant::now();
let ip = Some("test-ip-hash-rate-limit".to_string());
for _ in 0..MAX_ATTEMPTS {
assert!(!rate_limited(ip.clone(), now));
}
assert!(rate_limited(ip.clone(), now));
assert!(!rate_limited(Some("other-ip".into()), now));
let later = now + WINDOW + Duration::from_secs(1);
assert!(!rate_limited(ip, later));
}
#[test]
fn missing_ip_is_never_limited() {
let now = Instant::now();
for _ in 0..(MAX_ATTEMPTS * 3) {
assert!(!rate_limited(None, now));
}
}
}