src/storage/memory.rs
Rama main · réplica de lectura del repositorio original.
use std::collections::HashMap;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::net::IpAddr;
use std::path::Path;
use std::sync::{Mutex, RwLock};
use chrono::{DateTime, Utc};
use rusqlite::{Connection, params};
use uuid::Uuid;
use crate::error::CaptchaError;
use crate::puzzle::types::Challenge;
use crate::risk::behavior::BEHAVIOR_DUPLICATE_WINDOW_SECS;
use crate::risk::signals::RATE_SUSTAINED_WINDOW_SECS;
use crate::site::types::{Site, SitePolicy};
use super::Store;
struct RateWindow {
count: u32,
window_start: i64,
}
/// Number of lock stripes for each rate map. Every `GET /v1/puzzle` bumps a
/// per-IP and a per-site counter; a single `RwLock<HashMap>` per map would
/// serialise every request on one writer. Striping by key hash spreads that
/// contention across independent locks. Power of two so the modulo is cheap.
const RATE_SHARDS: usize = 16;
/// Rate-window map split across `RATE_SHARDS` independently-locked stripes.
/// Each key lands in one stripe by hash, so concurrent requests for different
/// keys rarely contend. The critical section stays tiny (an entry lookup + a
/// counter bump), matching the previous single-lock behaviour per key.
struct ShardedRates<K> {
shards: [RwLock<HashMap<K, RateWindow>>; RATE_SHARDS],
/// Length of one counting window, in seconds. Per-instance because the
/// maps count different things: issuance rate over 60s, repeated
/// behaviour blobs over `BEHAVIOR_DUPLICATE_WINDOW_SECS`.
window_secs: i64,
}
impl<K: Eq + Hash + Clone> ShardedRates<K> {
fn new(window_secs: i64) -> Self {
Self {
shards: std::array::from_fn(|_| RwLock::new(HashMap::new())),
window_secs,
}
}
fn shard(&self, key: &K) -> &RwLock<HashMap<K, RateWindow>> {
let mut hasher = DefaultHasher::new();
key.hash(&mut hasher);
&self.shards[(hasher.finish() as usize) % RATE_SHARDS]
}
/// Increment the key's counter within the current window, resetting it
/// when the window has rolled over. Returns the new count.
fn increment(&self, key: K) -> Result<u32, CaptchaError> {
let mut map = self.shard(&key).write().map_err(lock_err)?;
let now = Utc::now().timestamp();
let entry = map.entry(key).or_insert(RateWindow {
count: 0,
window_start: now,
});
if now - entry.window_start >= self.window_secs {
entry.count = 0;
entry.window_start = now;
}
entry.count += 1;
Ok(entry.count)
}
/// Drop windows older than twice this map's window across every stripe.
fn retain_fresh(&self, now_ts: i64) -> Result<(), CaptchaError> {
for shard in &self.shards {
let mut map = shard.write().map_err(lock_err)?;
map.retain(|_, w| now_ts - w.window_start < self.window_secs * 2);
}
Ok(())
}
}
pub struct InMemoryStore {
challenges: RwLock<HashMap<Uuid, Challenge>>,
sites_by_key: RwLock<HashMap<Uuid, Site>>,
sites_by_secret: RwLock<HashMap<String, Uuid>>,
ip_rates: ShardedRates<IpAddr>,
/// The same per-IP key over the 15 min window. Entries live for two
/// windows before the sweeper reclaims them, so this map holds every
/// distinct source seen in the last half hour — a few dozen bytes each.
ip_rates_sustained: ShardedRates<IpAddr>,
site_rates: ShardedRates<Uuid>,
/// Occurrences of one verify-time behaviour blob per site, keyed on
/// `(site_key, behavior_fingerprint)`. Same shape as the counters above,
/// on a longer window. Entries are only created after a *valid* proof of
/// work, so growing this map costs an attacker one solve per entry, and
/// `cleanup_expired` reclaims it on the same sweep as the rate windows.
behavior_blobs: ShardedRates<(Uuid, u64)>,
/// Optional write-through persistence for the sites table. Challenges
/// and rate windows stay in-memory — they're cheap to lose on restart.
/// Sites can't be: integrators store the secret_key client-side, so a
/// restart that wipes them means every embedded captcha breaks.
site_persistence: Option<Mutex<Connection>>,
}
const RATE_WINDOW_SECS: i64 = 60;
const SITES_SCHEMA: &str = r"
CREATE TABLE IF NOT EXISTS sites (
site_key TEXT PRIMARY KEY,
secret_key TEXT NOT NULL UNIQUE,
name TEXT NOT NULL,
created_at TEXT NOT NULL,
allowed_origins TEXT NOT NULL DEFAULT '',
policy TEXT
);
";
impl InMemoryStore {
pub fn new() -> Self {
Self {
challenges: RwLock::new(HashMap::new()),
sites_by_key: RwLock::new(HashMap::new()),
sites_by_secret: RwLock::new(HashMap::new()),
ip_rates: ShardedRates::new(RATE_WINDOW_SECS),
ip_rates_sustained: ShardedRates::new(RATE_SUSTAINED_WINDOW_SECS),
site_rates: ShardedRates::new(RATE_WINDOW_SECS),
behavior_blobs: ShardedRates::new(BEHAVIOR_DUPLICATE_WINDOW_SECS),
site_persistence: None,
}
}
/// Open or create a SQLite database at `path` for site persistence,
/// load any existing sites into the in-memory map, and keep the
/// connection for write-through inserts.
pub fn with_site_persistence(path: impl AsRef<Path>) -> Result<Self, CaptchaError> {
let conn = Connection::open(path.as_ref())
.map_err(|e| CaptchaError::Storage(format!("open site db: {e}")))?;
conn.execute_batch(SITES_SCHEMA)
.map_err(|e| CaptchaError::Storage(format!("init site schema: {e}")))?;
migrate_add_site_column(&conn, "allowed_origins", "TEXT NOT NULL DEFAULT ''")?;
// Nullable: NULL means "no overrides", which is what every site
// predating per-site policy has and what an emptied policy writes back.
migrate_add_site_column(&conn, "policy", "TEXT")?;
let sites = load_all_sites(&conn)?;
let store = Self::new();
{
let mut by_key = store.sites_by_key.write().map_err(lock_err)?;
let mut by_secret = store.sites_by_secret.write().map_err(lock_err)?;
for site in sites {
by_secret.insert(site.secret_key.clone(), site.site_key);
by_key.insert(site.site_key, site);
}
}
Ok(Self {
site_persistence: Some(Mutex::new(conn)),
..store
})
}
/// Number of challenges currently held in memory. Used by the puzzle
/// handler as a global issuance backstop (`MAX_ACTIVE_CHALLENGES`) so a
/// distributed flood can't grow the map unbounded between cleanup sweeps.
pub fn challenge_count(&self) -> usize {
self.challenges.read().map(|m| m.len()).unwrap_or_default()
}
/// Number of sites currently loaded. Useful at boot to log how many
/// were restored from disk.
pub fn site_count(&self) -> usize {
self.sites_by_key
.read()
.map(|m| m.len())
.unwrap_or_default()
}
/// Snapshot of every registered site, sorted by `created_at` ascending.
/// Used by the admin dashboard. Cloning is cheap (a few strings per site)
/// and keeps the read lock held for as little time as possible.
pub fn list_sites(&self) -> Result<Vec<Site>, CaptchaError> {
let map = self.sites_by_key.read().map_err(lock_err)?;
let mut out: Vec<Site> = map.values().cloned().collect();
out.sort_by_key(|s| s.created_at);
Ok(out)
}
}
/// Add a column to the `sites` table when a database created by an older
/// build doesn't have it yet. SQLite has no `ADD COLUMN IF NOT EXISTS`, so
/// probe `pragma_table_info` first and only ALTER when the column is absent —
/// a fresh database created from the current `SITES_SCHEMA` already has every
/// column, so this is a no-op there. Mirrors the in-place migration pattern in
/// `dashboard/log.rs`.
fn migrate_add_site_column(
conn: &Connection,
column: &str,
definition: &str,
) -> Result<(), CaptchaError> {
let has_column: bool = conn
.prepare("SELECT 1 FROM pragma_table_info('sites') WHERE name = ?1")
.and_then(|mut stmt| stmt.exists([column]))
.map_err(|e| CaptchaError::Storage(format!("probe sites schema: {e}")))?;
if !has_column {
// `column`/`definition` are compile-time literals from the call sites
// below, never user input — SQLite can't bind an identifier anyway.
conn.execute(
&format!("ALTER TABLE sites ADD COLUMN {column} {definition}"),
[],
)
.map_err(|e| CaptchaError::Storage(format!("add {column} column: {e}")))?;
}
Ok(())
}
fn load_all_sites(conn: &Connection) -> Result<Vec<Site>, CaptchaError> {
let mut stmt = conn
.prepare(
"SELECT site_key, secret_key, name, created_at, allowed_origins, policy FROM sites",
)
.map_err(|e| CaptchaError::Storage(format!("prepare load sites: {e}")))?;
let rows = stmt
.query_map([], |row| {
let key_str: String = row.get(0)?;
let secret: String = row.get(1)?;
let name: String = row.get(2)?;
let created_str: String = row.get(3)?;
let origins_str: String = row.get(4)?;
let policy_json: Option<String> = row.get(5)?;
Ok((key_str, secret, name, created_str, origins_str, policy_json))
})
.map_err(|e| CaptchaError::Storage(format!("query sites: {e}")))?;
let mut out = Vec::new();
for row in rows {
let (key_str, secret_key, name, created_str, origins_str, policy_json) =
row.map_err(|e| CaptchaError::Storage(format!("read site row: {e}")))?;
let site_key = Uuid::parse_str(&key_str)
.map_err(|e| CaptchaError::Storage(format!("bad site_key in db: {e}")))?;
let created_at = DateTime::parse_from_rfc3339(&created_str)
.map_err(|e| CaptchaError::Storage(format!("bad created_at in db: {e}")))?
.with_timezone(&Utc);
// A policy row we can't parse must not be silently downgraded to
// "no overrides": that would quietly restore the global (looser)
// thresholds for a site whose operator deliberately tightened them.
// Refuse to boot instead, the same stance `AppConfig::validated` takes.
let policy = match policy_json {
Some(json) => serde_json::from_str(&json).map_err(|e| {
CaptchaError::Storage(format!("bad policy for site {key_str} in db: {e}"))
})?,
None => SitePolicy::default(),
};
out.push(Site {
site_key,
secret_key,
name,
created_at,
// Space-joined; entries never contain whitespace (normalized at
// provisioning time), so splitting on whitespace is lossless.
allowed_origins: origins_str.split_whitespace().map(String::from).collect(),
policy,
});
}
Ok(out)
}
/// Serialize a policy for the `policy` column: `None` (SQL NULL) when nothing
/// is overridden, so an untouched site stays visibly untouched on disk and the
/// column reads the same as it does for sites created before policies existed.
fn policy_column(policy: &SitePolicy) -> Result<Option<String>, CaptchaError> {
if policy.is_empty() {
return Ok(None);
}
serde_json::to_string(policy)
.map(Some)
.map_err(|e| CaptchaError::Storage(format!("serialize site policy: {e}")))
}
fn persist_site(conn: &Connection, site: &Site) -> Result<(), CaptchaError> {
conn.execute(
"INSERT INTO sites (site_key, secret_key, name, created_at, allowed_origins, policy) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
params![
site.site_key.to_string(),
site.secret_key,
site.name,
site.created_at.to_rfc3339(),
site.allowed_origins.join(" "),
policy_column(&site.policy)?,
],
)
.map_err(|e| CaptchaError::Storage(format!("insert site: {e}")))?;
Ok(())
}
impl Default for InMemoryStore {
fn default() -> Self {
Self::new()
}
}
impl Store for InMemoryStore {
async fn store_challenge(&self, challenge: &Challenge) -> Result<(), CaptchaError> {
let mut map = self.challenges.write().map_err(lock_err)?;
map.insert(challenge.id, challenge.clone());
Ok(())
}
async fn get_challenge(&self, id: &Uuid) -> Result<Option<Challenge>, CaptchaError> {
let map = self.challenges.read().map_err(lock_err)?;
Ok(map.get(id).cloned())
}
async fn delete_challenge(&self, id: &Uuid) -> Result<(), CaptchaError> {
let mut map = self.challenges.write().map_err(lock_err)?;
map.remove(id);
Ok(())
}
async fn consume_challenge(&self, challenge_id: &Uuid) -> Result<(), CaptchaError> {
// Single-use is enforced by removal: exactly one caller removes the
// challenge and passes; concurrent/replayed submits get `None` →
// `ChallengeNotFound` (404).
let mut map = self.challenges.write().map_err(lock_err)?;
match map.remove(challenge_id) {
Some(_) => Ok(()),
None => Err(CaptchaError::ChallengeNotFound),
}
}
async fn register_failed_attempt(
&self,
challenge_id: &Uuid,
max_attempts: u32,
) -> Result<bool, CaptchaError> {
if max_attempts == 0 {
return Ok(false);
}
let mut map = self.challenges.write().map_err(lock_err)?;
if let Some(challenge) = map.get_mut(challenge_id) {
challenge.attempts += 1;
if challenge.attempts >= max_attempts {
map.remove(challenge_id);
return Ok(true);
}
}
Ok(false)
}
async fn store_site(&self, site: &Site) -> Result<(), CaptchaError> {
// Hold the in-memory write locks across the SQLite write so persistence
// and memory can't diverge under concurrent mutations of the same site
// (two admin writes committing to disk in one order but to memory in
// the other). Lock order is by_key → by_secret everywhere. Persist
// inside the critical section, first, so a failed disk write (e.g.
// UNIQUE collision on a regenerated secret) leaves memory untouched.
let mut by_key = self.sites_by_key.write().map_err(lock_err)?;
let mut by_secret = self.sites_by_secret.write().map_err(lock_err)?;
if let Some(persist) = &self.site_persistence {
let conn = persist.lock().map_err(lock_err)?;
persist_site(&conn, site)?;
}
by_key.insert(site.site_key, site.clone());
by_secret.insert(site.secret_key.clone(), site.site_key);
Ok(())
}
async fn get_site_by_key(&self, site_key: &Uuid) -> Result<Option<Site>, CaptchaError> {
let map = self.sites_by_key.read().map_err(lock_err)?;
Ok(map.get(site_key).cloned())
}
async fn get_site_by_secret(&self, secret: &str) -> Result<Option<Site>, CaptchaError> {
let site_key = {
let by_secret = self.sites_by_secret.read().map_err(lock_err)?;
by_secret.get(secret).copied()
};
match site_key {
Some(key) => self.get_site_by_key(&key).await,
None => Ok(None),
}
}
async fn rotate_site_secret(
&self,
site_key: &Uuid,
new_secret: String,
) -> Result<(), CaptchaError> {
// One critical section (see store_site): resolve NotFound from memory,
// persist inside the lock, then swap in-memory. On UNIQUE collision
// (vanishingly rare for 32-byte random secrets) the persist fails and
// memory is left unchanged.
let mut by_key = self.sites_by_key.write().map_err(lock_err)?;
let mut by_secret = self.sites_by_secret.write().map_err(lock_err)?;
let Some(site) = by_key.get_mut(site_key) else {
return Err(CaptchaError::NotFound);
};
if let Some(persist) = &self.site_persistence {
let conn = persist.lock().map_err(lock_err)?;
let updated = conn
.execute(
"UPDATE sites SET secret_key = ?1 WHERE site_key = ?2",
params![new_secret, site_key.to_string()],
)
.map_err(|e| CaptchaError::Storage(format!("rotate site secret: {e}")))?;
if updated == 0 {
// Memory has the site but the DB row is missing — a pre-existing
// divergence we should never hit. Surface it; memory is untouched.
return Err(CaptchaError::Storage(
"site in memory but missing from db".into(),
));
}
}
let old_secret = std::mem::replace(&mut site.secret_key, new_secret.clone());
by_secret.remove(&old_secret);
by_secret.insert(new_secret, *site_key);
Ok(())
}
async fn update_site_origins(
&self,
site_key: &Uuid,
origins: Vec<String>,
) -> Result<(), CaptchaError> {
// One critical section (see store_site): resolve NotFound from memory,
// persist inside the lock, then swap in-memory.
let mut by_key = self.sites_by_key.write().map_err(lock_err)?;
let Some(site) = by_key.get_mut(site_key) else {
return Err(CaptchaError::NotFound);
};
if let Some(persist) = &self.site_persistence {
let conn = persist.lock().map_err(lock_err)?;
let updated = conn
.execute(
"UPDATE sites SET allowed_origins = ?1 WHERE site_key = ?2",
params![origins.join(" "), site_key.to_string()],
)
.map_err(|e| CaptchaError::Storage(format!("update site origins: {e}")))?;
if updated == 0 {
return Err(CaptchaError::Storage(
"site in memory but missing from db".into(),
));
}
}
site.allowed_origins = origins;
Ok(())
}
async fn update_site_name(&self, site_key: &Uuid, name: String) -> Result<(), CaptchaError> {
// One critical section (see store_site): resolve NotFound from memory,
// persist inside the lock, then swap in-memory. The name is not indexed
// — nothing looks a site up by it — so only `sites_by_key` is touched.
let mut by_key = self.sites_by_key.write().map_err(lock_err)?;
let Some(site) = by_key.get_mut(site_key) else {
return Err(CaptchaError::NotFound);
};
if let Some(persist) = &self.site_persistence {
let conn = persist.lock().map_err(lock_err)?;
let updated = conn
.execute(
"UPDATE sites SET name = ?1 WHERE site_key = ?2",
params![name, site_key.to_string()],
)
.map_err(|e| CaptchaError::Storage(format!("update site name: {e}")))?;
if updated == 0 {
return Err(CaptchaError::Storage(
"site in memory but missing from db".into(),
));
}
}
site.name = name;
Ok(())
}
async fn update_site_policy(
&self,
site_key: &Uuid,
policy: SitePolicy,
) -> Result<(), CaptchaError> {
// One critical section (see store_site): resolve NotFound from memory,
// persist inside the lock, then swap in-memory.
let mut by_key = self.sites_by_key.write().map_err(lock_err)?;
let Some(site) = by_key.get_mut(site_key) else {
return Err(CaptchaError::NotFound);
};
if let Some(persist) = &self.site_persistence {
let conn = persist.lock().map_err(lock_err)?;
let updated = conn
.execute(
"UPDATE sites SET policy = ?1 WHERE site_key = ?2",
params![policy_column(&policy)?, site_key.to_string()],
)
.map_err(|e| CaptchaError::Storage(format!("update site policy: {e}")))?;
if updated == 0 {
return Err(CaptchaError::Storage(
"site in memory but missing from db".into(),
));
}
}
site.policy = policy;
Ok(())
}
async fn delete_site(&self, site_key: &Uuid) -> Result<(), CaptchaError> {
// One critical section (see store_site): resolve NotFound from memory,
// delete from disk inside the lock, then remove from memory.
let mut by_key = self.sites_by_key.write().map_err(lock_err)?;
let mut by_secret = self.sites_by_secret.write().map_err(lock_err)?;
if !by_key.contains_key(site_key) {
return Err(CaptchaError::NotFound);
}
if let Some(persist) = &self.site_persistence {
let conn = persist.lock().map_err(lock_err)?;
let deleted = conn
.execute(
"DELETE FROM sites WHERE site_key = ?1",
params![site_key.to_string()],
)
.map_err(|e| CaptchaError::Storage(format!("delete site: {e}")))?;
if deleted == 0 {
return Err(CaptchaError::Storage(
"site in memory but missing from db".into(),
));
}
}
let site = by_key
.remove(site_key)
.expect("presence checked under the same lock");
by_secret.remove(&site.secret_key);
Ok(())
}
async fn increment_ip_count(&self, ip: &IpAddr) -> Result<u32, CaptchaError> {
self.ip_rates.increment(*ip)
}
async fn increment_ip_count_sustained(&self, ip: &IpAddr) -> Result<u32, CaptchaError> {
self.ip_rates_sustained.increment(*ip)
}
async fn increment_site_count(&self, site_key: &Uuid) -> Result<u32, CaptchaError> {
self.site_rates.increment(*site_key)
}
async fn increment_behavior_blob_count(
&self,
site_key: &Uuid,
fingerprint: u64,
) -> Result<u32, CaptchaError> {
self.behavior_blobs.increment((*site_key, fingerprint))
}
async fn cleanup_expired(&self) -> Result<(), CaptchaError> {
let now = Utc::now();
// Clean expired challenges
{
let mut map = self.challenges.write().map_err(lock_err)?;
map.retain(|_, c| c.expires_at > now);
}
// Clean stale rate windows across every stripe
let ts = now.timestamp();
self.ip_rates.retain_fresh(ts)?;
self.ip_rates_sustained.retain_fresh(ts)?;
self.site_rates.retain_fresh(ts)?;
self.behavior_blobs.retain_fresh(ts)?;
Ok(())
}
}
fn lock_err<T>(_: T) -> CaptchaError {
CaptchaError::Storage("Lock poisoned".into())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::puzzle::types::Algorithm;
use chrono::Duration;
fn make_challenge(site_key: Uuid, ttl_secs: i64) -> Challenge {
let now = Utc::now();
Challenge {
id: Uuid::new_v4(),
site_key,
algorithm: Algorithm::Sha256,
prefix: "deadbeef".into(),
difficulty: 8,
created_at: now,
dwell_since: now,
expires_at: now + Duration::seconds(ttl_secs),
attempts: 0,
issued_to: None,
}
}
fn make_site() -> Site {
// Random secret so two sites in the same test don't collide on
// the UNIQUE constraint when persistence is enabled.
let secret = format!("secret-{}", Uuid::new_v4().simple());
Site {
site_key: Uuid::new_v4(),
secret_key: secret,
name: "Test Site".into(),
created_at: Utc::now(),
allowed_origins: Vec::new(),
policy: SitePolicy::default(),
}
}
#[tokio::test]
async fn test_challenge_crud() {
let store = InMemoryStore::new();
let challenge = make_challenge(Uuid::new_v4(), 300);
store.store_challenge(&challenge).await.unwrap();
let fetched = store.get_challenge(&challenge.id).await.unwrap().unwrap();
assert_eq!(fetched.id, challenge.id);
store.delete_challenge(&challenge.id).await.unwrap();
assert!(store.get_challenge(&challenge.id).await.unwrap().is_none());
}
#[tokio::test]
async fn test_consume_challenge_removes_once() {
let store = InMemoryStore::new();
let challenge = make_challenge(Uuid::new_v4(), 300);
store.store_challenge(&challenge).await.unwrap();
store.consume_challenge(&challenge.id).await.unwrap();
assert!(store.get_challenge(&challenge.id).await.unwrap().is_none());
let second = store.consume_challenge(&challenge.id).await;
assert!(matches!(second, Err(CaptchaError::ChallengeNotFound)));
}
#[tokio::test]
async fn test_register_failed_attempt_evicts_at_cap() {
let store = InMemoryStore::new();
let challenge = make_challenge(Uuid::new_v4(), 300);
store.store_challenge(&challenge).await.unwrap();
// Below the cap the challenge stays live.
assert!(
!store
.register_failed_attempt(&challenge.id, 3)
.await
.unwrap()
);
assert!(
!store
.register_failed_attempt(&challenge.id, 3)
.await
.unwrap()
);
assert!(store.get_challenge(&challenge.id).await.unwrap().is_some());
// The third failed attempt reaches the cap and evicts it.
assert!(
store
.register_failed_attempt(&challenge.id, 3)
.await
.unwrap()
);
assert!(store.get_challenge(&challenge.id).await.unwrap().is_none());
}
#[tokio::test]
async fn test_register_failed_attempt_disabled_when_zero() {
let store = InMemoryStore::new();
let challenge = make_challenge(Uuid::new_v4(), 300);
store.store_challenge(&challenge).await.unwrap();
// max_attempts = 0 disables the cap: never evicts, never increments.
for _ in 0..1000 {
assert!(
!store
.register_failed_attempt(&challenge.id, 0)
.await
.unwrap()
);
}
assert!(store.get_challenge(&challenge.id).await.unwrap().is_some());
}
#[tokio::test]
async fn test_register_failed_attempt_missing_challenge() {
let store = InMemoryStore::new();
// A never-existed / already-consumed challenge is a no-op, not an error.
assert!(
!store
.register_failed_attempt(&Uuid::new_v4(), 3)
.await
.unwrap()
);
}
#[tokio::test]
async fn test_challenge_count() {
let store = InMemoryStore::new();
assert_eq!(store.challenge_count(), 0);
let a = make_challenge(Uuid::new_v4(), 300);
let b = make_challenge(Uuid::new_v4(), 300);
store.store_challenge(&a).await.unwrap();
store.store_challenge(&b).await.unwrap();
assert_eq!(store.challenge_count(), 2);
store.consume_challenge(&a.id).await.unwrap();
assert_eq!(store.challenge_count(), 1);
}
#[tokio::test]
async fn test_site_crud() {
let store = InMemoryStore::new();
let site = make_site();
store.store_site(&site).await.unwrap();
let by_key = store
.get_site_by_key(&site.site_key)
.await
.unwrap()
.unwrap();
assert_eq!(by_key.name, site.name);
let by_secret = store
.get_site_by_secret(&site.secret_key)
.await
.unwrap()
.unwrap();
assert_eq!(by_secret.site_key, site.site_key);
}
#[tokio::test]
async fn test_site_not_found() {
let store = InMemoryStore::new();
assert!(
store
.get_site_by_key(&Uuid::new_v4())
.await
.unwrap()
.is_none()
);
assert!(
store
.get_site_by_secret("nonexistent")
.await
.unwrap()
.is_none()
);
}
#[tokio::test]
async fn test_sustained_ip_counter_is_independent_of_the_minute_counter() {
let store = InMemoryStore::new();
let ip: IpAddr = "127.0.0.1".parse().unwrap();
for _ in 0..3 {
store.increment_ip_count(&ip).await.unwrap();
}
// The sustained map has seen nothing yet; it starts its own count.
assert_eq!(store.increment_ip_count_sustained(&ip).await.unwrap(), 1);
assert_eq!(store.increment_ip_count_sustained(&ip).await.unwrap(), 2);
// And the minute counter was not touched by those increments.
assert_eq!(store.increment_ip_count(&ip).await.unwrap(), 4);
}
#[tokio::test]
async fn test_ip_rate_counter() {
let store = InMemoryStore::new();
let ip: IpAddr = "127.0.0.1".parse().unwrap();
assert_eq!(store.increment_ip_count(&ip).await.unwrap(), 1);
assert_eq!(store.increment_ip_count(&ip).await.unwrap(), 2);
assert_eq!(store.increment_ip_count(&ip).await.unwrap(), 3);
}
#[tokio::test]
async fn test_site_rate_counter() {
let store = InMemoryStore::new();
let site_key = Uuid::new_v4();
assert_eq!(store.increment_site_count(&site_key).await.unwrap(), 1);
assert_eq!(store.increment_site_count(&site_key).await.unwrap(), 2);
}
#[tokio::test]
async fn test_behavior_blob_counter_is_per_site_and_per_fingerprint() {
let store = InMemoryStore::new();
let site = Uuid::new_v4();
let other_site = Uuid::new_v4();
assert_eq!(
store
.increment_behavior_blob_count(&site, 42)
.await
.unwrap(),
1
);
assert_eq!(
store
.increment_behavior_blob_count(&site, 42)
.await
.unwrap(),
2
);
// A different blob on the same site is a separate count...
assert_eq!(
store
.increment_behavior_blob_count(&site, 43)
.await
.unwrap(),
1
);
// ...and so is the same blob on a different site: one tenant's traffic
// must never push another tenant's visitors over the threshold.
assert_eq!(
store
.increment_behavior_blob_count(&other_site, 42)
.await
.unwrap(),
1
);
}
#[tokio::test]
async fn test_cleanup_reclaims_stale_behavior_blob_windows() {
let store = InMemoryStore::new();
let site = Uuid::new_v4();
store.increment_behavior_blob_count(&site, 7).await.unwrap();
// Backdate past twice the dedup window, which is what `retain_fresh`
// reclaims on. Without the sweep this map would grow for the process
// lifetime.
{
let shard = store.behavior_blobs.shard(&(site, 7));
let mut map = shard.write().unwrap();
map.get_mut(&(site, 7)).unwrap().window_start =
Utc::now().timestamp() - BEHAVIOR_DUPLICATE_WINDOW_SECS * 3;
}
store.cleanup_expired().await.unwrap();
// Assert on the map, not on a fresh count: `increment` rolls a stale
// window over by itself, so a count of 1 would look identical whether
// or not the sweeper ever ran.
assert!(
!store
.behavior_blobs
.shard(&(site, 7))
.read()
.unwrap()
.contains_key(&(site, 7)),
"the sweeper must reclaim the dedup map, not just the rate windows"
);
}
#[tokio::test]
async fn test_rotate_site_secret_in_memory() {
let store = InMemoryStore::new();
let site = make_site();
store.store_site(&site).await.unwrap();
let new_secret = "rotated-secret".to_string();
store
.rotate_site_secret(&site.site_key, new_secret.clone())
.await
.unwrap();
// Old secret no longer resolves.
assert!(
store
.get_site_by_secret(&site.secret_key)
.await
.unwrap()
.is_none()
);
// New secret resolves.
let by_new = store
.get_site_by_secret(&new_secret)
.await
.unwrap()
.unwrap();
assert_eq!(by_new.site_key, site.site_key);
assert_eq!(by_new.secret_key, new_secret);
}
#[tokio::test]
async fn test_rotate_site_secret_not_found() {
let store = InMemoryStore::new();
let result = store.rotate_site_secret(&Uuid::new_v4(), "x".into()).await;
assert!(matches!(result, Err(CaptchaError::NotFound)));
}
#[tokio::test]
async fn test_delete_site_in_memory() {
let store = InMemoryStore::new();
let site = make_site();
store.store_site(&site).await.unwrap();
store.delete_site(&site.site_key).await.unwrap();
assert!(
store
.get_site_by_key(&site.site_key)
.await
.unwrap()
.is_none()
);
assert!(
store
.get_site_by_secret(&site.secret_key)
.await
.unwrap()
.is_none()
);
}
#[tokio::test]
async fn test_delete_site_not_found() {
let store = InMemoryStore::new();
let result = store.delete_site(&Uuid::new_v4()).await;
assert!(matches!(result, Err(CaptchaError::NotFound)));
}
#[tokio::test]
async fn test_rotate_and_delete_persist_across_reopen() {
let dir = tempdir();
let path = dir.join("sites.db");
let site_a = make_site();
let site_b = make_site();
{
let store = InMemoryStore::with_site_persistence(&path).unwrap();
store.store_site(&site_a).await.unwrap();
store.store_site(&site_b).await.unwrap();
store
.rotate_site_secret(&site_a.site_key, "rotated".into())
.await
.unwrap();
store.delete_site(&site_b.site_key).await.unwrap();
}
let store = InMemoryStore::with_site_persistence(&path).unwrap();
assert_eq!(store.site_count(), 1);
let reloaded = store
.get_site_by_secret("rotated")
.await
.unwrap()
.expect("rotated secret resolves after reopen");
assert_eq!(reloaded.site_key, site_a.site_key);
assert!(
store
.get_site_by_key(&site_b.site_key)
.await
.unwrap()
.is_none(),
"deleted site should not reload"
);
}
#[tokio::test]
async fn test_site_persistence_round_trip() {
let dir = tempdir();
let path = dir.join("sites.db");
// First store: register a site, drop the store.
let site = make_site();
{
let store = InMemoryStore::with_site_persistence(&path).unwrap();
store.store_site(&site).await.unwrap();
}
// Second store: should reload the site from disk.
let store = InMemoryStore::with_site_persistence(&path).unwrap();
assert_eq!(store.site_count(), 1);
let by_key = store
.get_site_by_key(&site.site_key)
.await
.unwrap()
.expect("site reloaded by key");
assert_eq!(by_key.name, site.name);
let by_secret = store
.get_site_by_secret(&site.secret_key)
.await
.unwrap()
.expect("site reloaded by secret");
assert_eq!(by_secret.site_key, site.site_key);
}
#[tokio::test]
async fn test_allowed_origins_round_trip() {
let dir = tempdir();
let path = dir.join("sites.db");
let mut site = make_site();
site.allowed_origins = vec![
"https://a.example".to_string(),
"https://b.example:8443".to_string(),
];
{
let store = InMemoryStore::with_site_persistence(&path).unwrap();
store.store_site(&site).await.unwrap();
}
let store = InMemoryStore::with_site_persistence(&path).unwrap();
let reloaded = store
.get_site_by_key(&site.site_key)
.await
.unwrap()
.expect("site reloaded");
assert_eq!(reloaded.allowed_origins, site.allowed_origins);
}
#[tokio::test]
async fn test_update_site_origins_persists_across_reopen() {
let dir = tempdir();
let path = dir.join("sites.db");
let site = make_site();
{
let store = InMemoryStore::with_site_persistence(&path).unwrap();
store.store_site(&site).await.unwrap();
// Site starts with no restriction; set one.
store
.update_site_origins(&site.site_key, vec!["https://only.example".to_string()])
.await
.unwrap();
}
let store = InMemoryStore::with_site_persistence(&path).unwrap();
let reloaded = store
.get_site_by_key(&site.site_key)
.await
.unwrap()
.expect("site reloaded");
assert_eq!(reloaded.allowed_origins, vec!["https://only.example"]);
}
#[tokio::test]
async fn test_update_site_origins_not_found() {
let store = InMemoryStore::new();
let result = store
.update_site_origins(&Uuid::new_v4(), vec!["https://x.example".to_string()])
.await;
assert!(matches!(result, Err(CaptchaError::NotFound)));
}
#[tokio::test]
async fn test_update_site_name_persists_across_reopen() {
let dir = tempdir();
let path = dir.join("sites.db");
let site = make_site();
{
let store = InMemoryStore::with_site_persistence(&path).unwrap();
store.store_site(&site).await.unwrap();
store
.update_site_name(&site.site_key, "renamed".to_string())
.await
.unwrap();
}
let store = InMemoryStore::with_site_persistence(&path).unwrap();
let reloaded = store
.get_site_by_key(&site.site_key)
.await
.unwrap()
.expect("site reloaded");
assert_eq!(reloaded.name, "renamed");
}
#[tokio::test]
async fn test_update_site_name_leaves_the_secret_lookup_intact() {
// The name is not an index, but it shares the write path with the two
// maps that are. A rename that disturbed `sites_by_secret` would break
// every `/v1/verify` for that site with nothing in the response saying
// why, so it gets a test rather than a comment.
let store = InMemoryStore::new();
let site = make_site();
store.store_site(&site).await.unwrap();
store
.update_site_name(&site.site_key, " renamed ".to_string())
.await
.unwrap();
let found = store
.get_site_by_secret(&site.secret_key)
.await
.unwrap()
.expect("the secret still resolves");
assert_eq!(found.site_key, site.site_key);
// Trimming is the handler's job (`normalize_site_name`), not the
// store's — this pins that the store writes what it is handed.
assert_eq!(found.name, " renamed ");
}
#[tokio::test]
async fn test_update_site_name_not_found() {
let store = InMemoryStore::new();
let result = store
.update_site_name(&Uuid::new_v4(), "whatever".to_string())
.await;
assert!(matches!(result, Err(CaptchaError::NotFound)));
}
#[tokio::test]
async fn test_old_schema_db_migrates_allowed_origins_column() {
let dir = tempdir();
let path = dir.join("sites.db");
// Hand-create the pre-feature table (no allowed_origins column) and
// insert a row, exactly as an older binary would have left it.
{
let conn = Connection::open(&path).unwrap();
conn.execute_batch(
r"CREATE TABLE sites (
site_key TEXT PRIMARY KEY,
secret_key TEXT NOT NULL UNIQUE,
name TEXT NOT NULL,
created_at TEXT NOT NULL
);",
)
.unwrap();
conn.execute(
"INSERT INTO sites (site_key, secret_key, name, created_at) VALUES (?1, ?2, ?3, ?4)",
params![
Uuid::new_v4().to_string(),
"legacy-secret",
"Legacy Site",
Utc::now().to_rfc3339(),
],
)
.unwrap();
}
// Opening with the current code must migrate the column in place and
// load the legacy row with an empty (unrestricted) origin list.
let store = InMemoryStore::with_site_persistence(&path).unwrap();
assert_eq!(store.site_count(), 1);
let reloaded = store
.get_site_by_secret("legacy-secret")
.await
.unwrap()
.expect("legacy site reloaded");
assert!(reloaded.allowed_origins.is_empty());
// And the migrated column is usable going forward.
store
.update_site_origins(&reloaded.site_key, vec!["https://new.example".to_string()])
.await
.unwrap();
// The same legacy row must also gain the `policy` column, and load
// with no overrides — a site provisioned before policies existed keeps
// tracking the global config exactly as it did.
assert!(reloaded.policy.is_empty());
store
.update_site_policy(
&reloaded.site_key,
SitePolicy {
tier_block_min: Some(70),
..Default::default()
},
)
.await
.unwrap();
}
#[tokio::test]
async fn test_policy_round_trips_through_persistence() {
let dir = tempdir();
let path = dir.join("sites.db");
let policy = SitePolicy {
tier_checkbox_min: Some(10),
tier_hard_pow_min: Some(25),
tier_block_min: Some(60),
..Default::default()
};
let site_key = {
let store = InMemoryStore::with_site_persistence(&path).unwrap();
let mut site = make_site();
site.policy = policy;
store.store_site(&site).await.unwrap();
site.site_key
};
// Reopen: the policy must come back byte-identical, or a restart
// silently reverts a tenant to the looser global thresholds.
let store = InMemoryStore::with_site_persistence(&path).unwrap();
let reloaded = store.get_site_by_key(&site_key).await.unwrap().unwrap();
assert_eq!(reloaded.policy, policy);
}
#[tokio::test]
async fn test_update_site_policy_persists_and_clears() {
let dir = tempdir();
let path = dir.join("sites.db");
let site = make_site();
{
let store = InMemoryStore::with_site_persistence(&path).unwrap();
store.store_site(&site).await.unwrap();
store
.update_site_policy(
&site.site_key,
SitePolicy {
verify_block_min: Some(30),
..Default::default()
},
)
.await
.unwrap();
}
{
let store = InMemoryStore::with_site_persistence(&path).unwrap();
let reloaded = store
.get_site_by_key(&site.site_key)
.await
.unwrap()
.unwrap();
assert_eq!(reloaded.policy.verify_block_min, Some(30));
// Clearing writes SQL NULL back, not "{}" — an emptied policy must
// be indistinguishable from a site that never had one.
store
.update_site_policy(&site.site_key, SitePolicy::default())
.await
.unwrap();
}
let store = InMemoryStore::with_site_persistence(&path).unwrap();
let reloaded = store
.get_site_by_key(&site.site_key)
.await
.unwrap()
.unwrap();
assert!(reloaded.policy.is_empty());
let conn = Connection::open(&path).unwrap();
let stored: Option<String> = conn
.query_row(
"SELECT policy FROM sites WHERE site_key = ?1",
params![site.site_key.to_string()],
|row| row.get(0),
)
.unwrap();
assert_eq!(stored, None);
}
#[tokio::test]
async fn test_update_site_policy_unknown_site_is_not_found() {
let store = InMemoryStore::new();
let result = store
.update_site_policy(&Uuid::new_v4(), SitePolicy::default())
.await;
assert!(matches!(result, Err(CaptchaError::NotFound)));
}
#[tokio::test]
async fn test_unparseable_policy_column_refuses_to_load() {
// A corrupt policy must not degrade to "no overrides" — that would
// silently restore the looser global thresholds for a site whose
// operator deliberately tightened them.
let dir = tempdir();
let path = dir.join("sites.db");
let site = make_site();
{
let store = InMemoryStore::with_site_persistence(&path).unwrap();
store.store_site(&site).await.unwrap();
}
let conn = Connection::open(&path).unwrap();
conn.execute(
"UPDATE sites SET policy = ?1 WHERE site_key = ?2",
params!["not json", site.site_key.to_string()],
)
.unwrap();
drop(conn);
assert!(InMemoryStore::with_site_persistence(&path).is_err());
}
fn tempdir() -> std::path::PathBuf {
let mut dir = std::env::temp_dir();
dir.push(format!(
"bollwark-test-{}-{}",
std::process::id(),
Uuid::new_v4()
));
std::fs::create_dir_all(&dir).unwrap();
dir
}
#[tokio::test]
async fn test_cleanup_expired() {
let store = InMemoryStore::new();
let expired = make_challenge(Uuid::new_v4(), -10); // already expired
let valid = make_challenge(Uuid::new_v4(), 300);
store.store_challenge(&expired).await.unwrap();
store.store_challenge(&valid).await.unwrap();
store.cleanup_expired().await.unwrap();
assert!(store.get_challenge(&expired.id).await.unwrap().is_none());
assert!(store.get_challenge(&valid.id).await.unwrap().is_some());
}
}