use std::net::IpAddr;
use std::num::NonZeroU32;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;
use axum::extract::Request;
use axum::middleware::Next;
use axum::response::Response;
use governor::clock::{Clock, DefaultClock};
use governor::{DefaultKeyedRateLimiter, Quota, RateLimiter};
use tower_governor::key_extractor::KeyExtractor;
use vti_common::rate_limit::TrustedProxyKeyExtractor;
use super::rate_limit::{
ROOM_ADDRESS_LIMITER, ROOM_PRESENTER_LIMITER, ROOM_TRANSFER_LIMITER, RateLimited,
SIGNED_ADDRESS_LIMITER, SIGNER_LIMITER, UNAUTH_LIMITER,
};
use crate::config::{RoomBudget, RoomRateLimitsConfig};
use crate::server::AppState;
use crate::trust_tasks::VerifiedAdmission;
pub const SIGNER_BURST: u32 = 60;
pub const SIGNER_PERIOD: Duration = Duration::from_millis(500);
pub const SIGNED_ADDRESS_BURST: u32 = 120;
pub const SIGNED_ADDRESS_PERIOD: Duration = Duration::from_millis(250);
#[derive(Clone)]
pub struct TrustTaskLimits {
address: Arc<DefaultKeyedRateLimiter<IpAddr>>,
signed_address: Arc<DefaultKeyedRateLimiter<IpAddr>>,
signer: Arc<DefaultKeyedRateLimiter<String>>,
rooms: Arc<OnceLock<Arc<RoomLimiters>>>,
key_extractor: TrustedProxyKeyExtractor,
}
impl TrustTaskLimits {
pub fn new(
address: Arc<DefaultKeyedRateLimiter<IpAddr>>,
key_extractor: TrustedProxyKeyExtractor,
) -> Self {
Self::with_quotas(
address,
key_extractor,
quota(SIGNED_ADDRESS_PERIOD, SIGNED_ADDRESS_BURST),
quota(SIGNER_PERIOD, SIGNER_BURST),
)
}
fn with_quotas(
address: Arc<DefaultKeyedRateLimiter<IpAddr>>,
key_extractor: TrustedProxyKeyExtractor,
signed_address: Quota,
signer: Quota,
) -> Self {
Self {
address,
signed_address: Arc::new(RateLimiter::keyed(signed_address)),
signer: Arc::new(RateLimiter::keyed(signer)),
rooms: Arc::new(OnceLock::new()),
key_extractor,
}
}
async fn rooms(&self, state: &AppState) -> Arc<RoomLimiters> {
if let Some(rooms) = self.rooms.get() {
return rooms.clone();
}
let config = state.config.read().await.rooms.rate_limits.clone();
self.rooms
.get_or_init(|| Arc::new(RoomLimiters::new(&config)))
.clone()
}
}
pub const ROOM_FAILURE_COST: u32 = 25;
const ROOM_TRACKED_KEYS: usize = 10_000;
pub(crate) struct RoomLimiters {
address: DefaultKeyedRateLimiter<IpAddr>,
presenter: DefaultKeyedRateLimiter<[u8; 16]>,
transfer: DefaultKeyedRateLimiter<[u8; 16]>,
secret: [u8; 32],
}
impl RoomLimiters {
fn new(config: &RoomRateLimitsConfig) -> Self {
Self::with_budgets(config.address(), config.presenter(), config.transfer())
}
fn with_budgets(address: RoomBudget, presenter: RoomBudget, transfer: RoomBudget) -> Self {
let q = |b: RoomBudget| quota(b.interval, b.burst);
let mut secret = [0u8; 32];
rand::fill(&mut secret);
Self {
address: RateLimiter::keyed(q(address)),
presenter: RateLimiter::keyed(q(presenter)),
transfer: RateLimiter::keyed(q(transfer)),
secret,
}
}
fn member_key(&self, room_id: &str, presenter: &str) -> [u8; 16] {
use sha2::{Digest, Sha256};
let mut h = Sha256::new();
h.update(self.secret);
h.update(room_id.as_bytes());
h.update([0u8]);
h.update(presenter.as_bytes());
let digest = h.finalize();
let mut key = [0u8; 16];
key.copy_from_slice(&digest[..16]);
key
}
fn charge_address(&self, address: &IpAddr) -> Result<(), RateLimited> {
if self.address.len() > ROOM_TRACKED_KEYS {
self.address.retain_recent();
}
charge(&self.address, address).map_err(|wait| {
RateLimited::new(
ROOM_ADDRESS_LIMITER,
wait,
"too many room documents from this address; wait for the \
Retry-After period before retrying",
)
})
}
fn penalize(&self, address: &IpAddr) {
for _ in 0..ROOM_FAILURE_COST {
if self.address.check_key(address).is_err() {
break;
}
}
}
fn charge_member(
&self,
room_id: &str,
presenter: &str,
traffic: RoomTraffic,
) -> Result<(), RateLimited> {
let (limiter, name, what) = match traffic {
RoomTraffic::Operation => (&self.presenter, ROOM_PRESENTER_LIMITER, "room operations"),
RoomTraffic::Transfer => (&self.transfer, ROOM_TRANSFER_LIMITER, "file chunks"),
};
if limiter.len() > ROOM_TRACKED_KEYS {
limiter.retain_recent();
}
charge(limiter, &self.member_key(room_id, presenter)).map_err(|wait| {
RateLimited::new(
name,
wait,
format!(
"too many {what} from this presenter in this room; wait for the \
Retry-After period before retrying"
),
)
})
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum RoomTraffic {
Operation,
Transfer,
}
pub(crate) struct RoomAdmission {
limiters: Arc<RoomLimiters>,
address: IpAddr,
authorized: AtomicBool,
refusal: Mutex<Option<RateLimited>>,
}
tokio::task_local! {
static ROOM_ADMISSION: Arc<RoomAdmission>;
}
pub(crate) fn admit_room(room_id: &str, presenter: &str, traffic: RoomTraffic) -> bool {
ROOM_ADMISSION
.try_with(|room| {
room.authorized.store(true, Ordering::SeqCst);
match room.limiters.charge_member(room_id, presenter, traffic) {
Ok(()) => true,
Err(limited) => {
if let Ok(mut slot) = room.refusal.lock() {
*slot = Some(limited);
}
false
}
}
})
.unwrap_or(true)
}
fn quota(period: Duration, burst: u32) -> Quota {
Quota::with_period(period)
.expect("the period is a non-zero constant")
.allow_burst(NonZeroU32::new(burst).expect("the burst is a non-zero constant"))
}
fn charge<K>(limiter: &DefaultKeyedRateLimiter<K>, key: &K) -> Result<(), u64>
where
K: std::hash::Hash + Eq + Clone,
{
limiter.check_key(key).map_err(|not_until| {
not_until
.wait_time_from(DefaultClock::default().now())
.as_secs_f64()
.ceil() as u64
})
}
fn unauth_refusal(wait: u64) -> RateLimited {
RateLimited::new(
UNAUTH_LIMITER,
wait,
"too many unauthenticated requests from this address; wait for the \
Retry-After period before retrying",
)
}
#[derive(Clone, Copy, Debug)]
pub struct ClientAddress(pub IpAddr);
pub async fn client_address(
axum::extract::State(limits): axum::extract::State<TrustTaskLimits>,
mut request: Request,
next: Next,
) -> Response {
match limits.key_extractor.extract(&request) {
Ok(ip) => {
request.extensions_mut().insert(ClientAddress(ip));
request.extensions_mut().insert(limits);
next.run(request).await
}
Err(e) => super::rate_limit::governor_error_response(e),
}
}
struct Claimed {
signer: String,
principal: String,
}
pub struct Admission {
limits: TrustTaskLimits,
address: IpAddr,
claimed: Option<Claimed>,
verified: AtomicBool,
address_charged: AtomicBool,
refusal: Mutex<Option<RateLimited>>,
room: Option<Arc<RoomAdmission>>,
}
impl Admission {
pub async fn begin(
state: &AppState,
limits: &TrustTaskLimits,
address: IpAddr,
body: &[u8],
) -> Result<Self, RateLimited> {
let claim = read_claim(body);
if claim
.type_uri
.as_deref()
.is_some_and(crate::rooms::handlers::serves)
{
let limiters = limits.rooms(state).await;
limiters.charge_address(&address)?;
return Ok(Self {
limits: limits.clone(),
address,
claimed: None,
verified: AtomicBool::new(false),
address_charged: AtomicBool::new(false),
refusal: Mutex::new(None),
room: Some(Arc::new(RoomAdmission {
limiters,
address,
authorized: AtomicBool::new(false),
refusal: Mutex::new(None),
})),
});
}
let claimed = match claim.signer() {
Some(did) => known_principal(state, &did).await.map(|principal| Claimed {
signer: did,
principal,
}),
None => None,
};
let admission = Self {
limits: limits.clone(),
address,
claimed,
verified: AtomicBool::new(false),
address_charged: AtomicBool::new(false),
refusal: Mutex::new(None),
room: None,
};
if admission.claimed.is_some() {
charge(&limits.signed_address, &address).map_err(|wait| {
RateLimited::new(
SIGNED_ADDRESS_LIMITER,
wait,
"too many signed documents from this address; wait for the \
Retry-After period before retrying",
)
})?;
} else {
admission.charge_address()?;
}
Ok(admission)
}
fn charge_address(&self) -> Result<(), RateLimited> {
self.address_charged.store(true, Ordering::SeqCst);
charge(&self.limits.address, &self.address).map_err(unauth_refusal)
}
fn refuse(&self, limited: RateLimited) -> bool {
if let Ok(mut slot) = self.refusal.lock() {
*slot = Some(limited);
}
false
}
pub fn verified_principal(&self) -> Option<&str> {
self.claimed
.as_ref()
.filter(|_| self.verified.load(Ordering::SeqCst))
.map(|claimed| claimed.principal.as_str())
}
pub fn finish(self) -> Result<(), RateLimited> {
if let Some(room) = &self.room {
if let Some(limited) = room.refusal.lock().ok().and_then(|mut slot| slot.take()) {
return Err(limited);
}
if !room.authorized.load(Ordering::SeqCst) {
room.limiters.penalize(&room.address);
}
return Ok(());
}
if let Some(limited) = self.refusal.lock().ok().and_then(|mut slot| slot.take()) {
return Err(limited);
}
if self.claimed.is_some()
&& !self.verified.load(Ordering::SeqCst)
&& !self.address_charged.load(Ordering::SeqCst)
{
self.charge_address()?;
}
Ok(())
}
}
impl Admission {
pub async fn scope<F: std::future::Future>(&self, fut: F) -> F::Output {
match &self.room {
Some(room) => ROOM_ADMISSION.scope(room.clone(), fut).await,
None => fut.await,
}
}
}
impl VerifiedAdmission for Admission {
fn room_precharged(&self) -> bool {
self.room.is_some()
}
fn admit_verified(&self, signer: &str) -> bool {
if self.room.is_some() {
return true;
}
match &self.claimed {
None => true,
Some(claimed) if claimed.signer == signer => {
self.verified.store(true, Ordering::SeqCst);
match charge(&self.limits.signer, &claimed.principal) {
Ok(()) => true,
Err(wait) => self.refuse(RateLimited::new(
SIGNER_LIMITER,
wait,
"too many documents from this signer; wait for the \
Retry-After period before retrying",
)),
}
}
Some(_) => match self.charge_address() {
Ok(()) => true,
Err(limited) => self.refuse(limited),
},
}
}
}
#[cfg(test)]
fn claimed_signer(body: &[u8]) -> Option<String> {
read_claim(body).signer()
}
#[derive(Default, serde::Deserialize)]
struct Claim {
#[serde(rename = "type")]
type_uri: Option<String>,
issuer: Option<String>,
proof: Option<serde::de::IgnoredAny>,
}
impl Claim {
fn signer(self) -> Option<String> {
self.proof?;
self.issuer.filter(|did| did.starts_with("did:"))
}
}
fn read_claim(body: &[u8]) -> Claim {
serde_json::from_slice(body).unwrap_or_default()
}
async fn known_principal(state: &AppState, did: &str) -> Option<String> {
let now = crate::auth::session::now_epoch();
let live = |entry: Option<crate::acl::VtcAclEntry>| entry.is_some_and(|e| !e.is_expired(now));
match crate::acl::get_acl_entry(&state.acl_ks, did).await {
Ok(Some(entry)) => return live(Some(entry)).then(|| did.to_string()),
Ok(None) => {}
Err(_) => return None,
}
let delegation = crate::acl::console_key::resolve_delegated_admin(&state.console_keys_ks, did)
.await
.ok()??;
let admin = crate::acl::get_acl_entry(&state.acl_ks, &delegation.admin_did)
.await
.ok()?;
live(admin).then_some(delegation.admin_did)
}
#[cfg(test)]
mod tests {
use super::*;
fn limits(signed_address_burst: u32, signer_burst: u32) -> TrustTaskLimits {
let slow = Duration::from_secs(3600);
TrustTaskLimits::with_quotas(
Arc::new(RateLimiter::keyed(quota(slow, 2))),
TrustedProxyKeyExtractor::new(Vec::new()),
quota(slow, signed_address_burst),
quota(slow, signer_burst),
)
}
fn admission(limits: &TrustTaskLimits, claimed: Option<&str>) -> Admission {
Admission {
limits: limits.clone(),
address: IpAddr::from([192, 0, 2, 1]),
claimed: claimed.map(|did| Claimed {
signer: did.to_string(),
principal: did.to_string(),
}),
verified: AtomicBool::new(false),
address_charged: AtomicBool::new(false),
refusal: Mutex::new(None),
room: None,
}
}
#[test]
fn the_claimed_signer_is_the_issuer_of_a_signed_document() {
let signed = br#"{"issuer":"did:key:z6Mk1","proof":{"type":"DataIntegrityProof"}}"#;
assert_eq!(claimed_signer(signed).as_deref(), Some("did:key:z6Mk1"));
assert_eq!(claimed_signer(br#"{"issuer":"did:key:z6Mk1"}"#), None);
assert_eq!(claimed_signer(br#"{"issuer":"nope","proof":{}}"#), None);
assert_eq!(claimed_signer(b"not json"), None);
}
#[test]
fn a_verified_known_signer_is_charged_its_own_bucket_not_the_address() {
let limits = limits(10, 3);
for _ in 0..3 {
let a = admission(&limits, Some("did:key:zAdmin"));
assert!(a.admit_verified("did:key:zAdmin"));
a.finish().expect("within the signer's burst");
}
assert!(charge(&limits.address, &IpAddr::from([192, 0, 2, 1])).is_ok());
let a = admission(&limits, Some("did:key:zAdmin"));
assert!(!a.admit_verified("did:key:zAdmin"));
assert_eq!(a.finish().unwrap_err().limiter(), SIGNER_LIMITER);
}
#[test]
fn a_claim_that_never_verifies_is_charged_the_address() {
let limits = limits(10, 10);
for _ in 0..2 {
admission(&limits, Some("did:key:zAdmin"))
.finish()
.expect("within the address's burst");
}
let refused = admission(&limits, Some("did:key:zAdmin"))
.finish()
.unwrap_err();
assert_eq!(refused.limiter(), UNAUTH_LIMITER);
}
#[test]
fn the_verified_principal_is_known_only_once_the_proof_verifies() {
let limits = limits(10, 10);
let mut a = admission(&limits, Some("did:key:zConsoleKey"));
a.claimed.as_mut().expect("claimed").principal = "did:key:zAdmin".to_string();
assert_eq!(
a.verified_principal(),
None,
"a claim is not a verification"
);
assert!(a.admit_verified("did:key:zConsoleKey"));
assert_eq!(a.verified_principal(), Some("did:key:zAdmin"));
let other = admission(&limits, Some("did:key:zAdmin"));
assert!(other.admit_verified("did:key:zSomeoneElse"));
assert_eq!(other.verified_principal(), None);
assert_eq!(admission(&limits, None).verified_principal(), None);
}
#[test]
fn verifying_as_another_signer_is_charged_the_address() {
let limits = limits(10, 10);
let a = admission(&limits, Some("did:key:zAdmin"));
assert!(a.admit_verified("did:key:zSomeoneElse"));
assert!(a.finish().is_ok());
assert!(
charge(&limits.signer, &"did:key:zAdmin".to_string()).is_ok(),
"the claimed signer's bucket is not charged"
);
assert!(charge(&limits.address, &IpAddr::from([192, 0, 2, 1])).is_ok());
assert!(
charge(&limits.address, &IpAddr::from([192, 0, 2, 1])).is_err(),
"the mismatched document spent one of the address's two"
);
}
fn rooms(address: u32, presenter: u32, transfer: u32) -> Arc<RoomLimiters> {
let slow = Duration::from_secs(3600);
let b = |burst| RoomBudget {
burst,
interval: slow,
};
Arc::new(RoomLimiters::with_budgets(
b(address),
b(presenter),
b(transfer),
))
}
fn room_admission(limiters: &Arc<RoomLimiters>) -> Arc<RoomAdmission> {
Arc::new(RoomAdmission {
limiters: limiters.clone(),
address: IpAddr::from([192, 0, 2, 9]),
authorized: AtomicBool::new(false),
refusal: Mutex::new(None),
})
}
#[test]
fn a_room_document_off_the_https_door_is_admitted() {
assert!(admit_room(
"did:key:zRoom",
"did:key:zMember",
RoomTraffic::Operation
));
}
#[tokio::test]
async fn presenters_and_rooms_are_charged_apart() {
let limiters = rooms(100, 2, 1);
let room = room_admission(&limiters);
ROOM_ADMISSION
.scope(room.clone(), async {
for _ in 0..2 {
assert!(admit_room("r1", "alice", RoomTraffic::Operation));
}
assert!(!admit_room("r1", "alice", RoomTraffic::Operation));
assert!(admit_room("r1", "bob", RoomTraffic::Operation));
assert!(admit_room("r2", "alice", RoomTraffic::Operation));
assert!(admit_room("r1", "alice", RoomTraffic::Transfer));
assert!(!admit_room("r1", "alice", RoomTraffic::Transfer));
})
.await;
let refused = room.refusal.lock().unwrap().take().expect("recorded");
assert_eq!(
refused.limiter(),
super::super::rate_limit::ROOM_TRANSFER_LIMITER
);
assert!(room.authorized.load(Ordering::SeqCst));
}
#[test]
fn a_room_document_never_authorized_costs_the_address_extra() {
let limiters = rooms(ROOM_FAILURE_COST + 2, 10, 10);
let address = IpAddr::from([192, 0, 2, 9]);
limiters.charge_address(&address).expect("first");
limiters.penalize(&address);
limiters.charge_address(&address).expect("the last token");
assert_eq!(
limiters.charge_address(&address).unwrap_err().limiter(),
super::super::rate_limit::ROOM_ADDRESS_LIMITER
);
}
#[test]
fn the_member_key_names_nobody_and_separates_room_from_presenter() {
let limiters = rooms(1, 1, 1);
let k = limiters.member_key("ab", "c");
assert_ne!(k, limiters.member_key("a", "bc"));
assert_ne!(k, rooms(1, 1, 1).member_key("ab", "c"), "keyed per process");
}
}