use trust_tasks_rs::{ErrorPayload, ErrorResponse, Payload, RejectReason, TrustTask};
pub const DEFAULT_MAX_DOCUMENT_BYTES: usize = 64 * 1024;
pub const DID_REGISTER_MAX_DOCUMENT_BYTES: usize = 1024 * 1024;
pub const DECLARED: &[(&str, usize)] = &[(
<trust_tasks_rs::specs::did_management::did::register::v0_1::Payload as Payload>::TYPE_URI,
DID_REGISTER_MAX_DOCUMENT_BYTES,
)];
pub const DETAILS_MAX_BYTES: &str = "maxBytes";
pub fn largest_max_document_bytes(served: &[&str]) -> usize {
largest_in(DECLARED, served)
}
fn largest_in(declared: &[(&str, usize)], served: &[&str]) -> usize {
declared
.iter()
.filter(|(uri, _)| served.contains(uri))
.map(|(_, max)| *max)
.fold(DEFAULT_MAX_DOCUMENT_BYTES, usize::max)
}
pub fn max_document_bytes(type_uri: &str, served: &[&str]) -> usize {
max_in(type_uri, DECLARED, served)
}
fn max_in(type_uri: &str, declared: &[(&str, usize)], served: &[&str]) -> usize {
declared
.iter()
.find(|(uri, _)| *uri == type_uri && served.contains(uri))
.map_or(DEFAULT_MAX_DOCUMENT_BYTES, |(_, max)| *max)
}
const MAX_NAMED_TYPE_CHARS: usize = 256;
pub fn check(body: &[u8], served: &[&str]) -> Result<(), ErrorResponse> {
check_with(body, |type_uri| max_document_bytes(type_uri, served))
}
pub const LARGE_DOCUMENT_BUDGET_PER_WINDOW: u64 = 5;
pub const LARGE_DOCUMENT_BUDGET_WINDOW_SECS: u64 = 60;
pub const LARGE_DOCUMENT_GLOBAL_BUDGET_PER_WINDOW: u64 = 60;
const MAX_TRACKED_ADDRESSES: usize = 10_000;
#[derive(Debug, Clone, Copy)]
struct Bucket {
spent: u64,
window_start: u64,
cost: u64,
}
#[derive(Debug, Default)]
pub struct LargeDocumentBudget {
inner: std::sync::Mutex<BudgetState>,
}
#[derive(Debug, Default)]
struct BudgetState {
buckets: std::collections::HashMap<String, Bucket>,
global_spent: u64,
global_window_start: u64,
}
impl LargeDocumentBudget {
pub fn new() -> Self {
Self::default()
}
fn try_charge(&self, address: &str, now: u64) -> Result<(), u64> {
let mut state = self
.inner
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let state = &mut *state;
if now.saturating_sub(state.global_window_start) >= LARGE_DOCUMENT_BUDGET_WINDOW_SECS {
state.global_spent = 0;
state.global_window_start = now;
}
let global_retry =
(state.global_window_start + LARGE_DOCUMENT_BUDGET_WINDOW_SECS).saturating_sub(now);
if state.global_spent >= LARGE_DOCUMENT_GLOBAL_BUDGET_PER_WINDOW {
return Err(global_retry);
}
let buckets = &mut state.buckets;
if buckets.len() >= MAX_TRACKED_ADDRESSES && !buckets.contains_key(address) {
buckets.retain(|_, b| {
b.cost > 1 || now.saturating_sub(b.window_start) < LARGE_DOCUMENT_BUDGET_WINDOW_SECS
});
if buckets.len() >= MAX_TRACKED_ADDRESSES {
return Err(global_retry);
}
}
let entry = buckets.entry(address.to_string()).or_insert(Bucket {
spent: 0,
window_start: now,
cost: 1,
});
if now.saturating_sub(entry.window_start) >= LARGE_DOCUMENT_BUDGET_WINDOW_SECS {
entry.spent = 0;
entry.window_start = now;
}
if entry.spent.saturating_add(entry.cost) > LARGE_DOCUMENT_BUDGET_PER_WINDOW {
let retry_after_secs =
(entry.window_start + LARGE_DOCUMENT_BUDGET_WINDOW_SECS).saturating_sub(now);
return Err(retry_after_secs);
}
entry.spent += entry.cost;
state.global_spent += 1;
Ok(())
}
fn penalize(&self, address: &str) {
let mut state = self
.inner
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let entry = state.buckets.entry(address.to_string()).or_insert(Bucket {
spent: 0,
window_start: 0,
cost: 1,
});
entry.cost = entry.cost.saturating_mul(2).min(1 << 20);
}
#[cfg(test)]
fn cost(&self, address: &str) -> u64 {
self.inner
.lock()
.unwrap()
.buckets
.get(address)
.map(|b| b.cost)
.unwrap_or(1)
}
}
#[derive(Debug)]
pub struct LargeDocumentCharge {
address: String,
claimed_issuer: String,
}
pub async fn check_for_known_issuer(
body: &[u8],
served: &[&str],
acl: &crate::server::store::KeyspaceHandle,
budget: &LargeDocumentBudget,
address: &str,
now: u64,
) -> Result<Option<LargeDocumentCharge>, ErrorResponse> {
if body.len() <= DEFAULT_MAX_DOCUMENT_BYTES {
return Ok(None);
}
let claimed_issuer = match peek_issuer(body) {
Some(issuer)
if matches!(
crate::server::acl::get_acl_entry(acl, &issuer).await,
Ok(Some(_))
) =>
{
Some(issuer)
}
_ => None,
};
let Some(claimed_issuer) = claimed_issuer else {
check_with(body, |_| DEFAULT_MAX_DOCUMENT_BYTES)?;
return Ok(None);
};
if let Err(retry_after_secs) = budget.try_charge(address, now) {
return Err(unrouted_error(
RejectReason::Unavailable {
retry_after: Some(
chrono::Utc::now() + chrono::Duration::seconds(retry_after_secs as i64),
),
}
.into(),
));
}
check(body, served)?;
Ok(Some(LargeDocumentCharge {
address: address.to_string(),
claimed_issuer,
}))
}
pub fn settle_large_document_charge(
budget: &LargeDocumentBudget,
charge: Option<&LargeDocumentCharge>,
verified_issuer: Option<&str>,
) {
let Some(charge) = charge else {
return;
};
if verified_issuer != Some(charge.claimed_issuer.as_str()) {
budget.penalize(&charge.address);
}
}
fn peek_issuer(body: &[u8]) -> Option<String> {
#[derive(serde::Deserialize)]
struct IssuerOnly {
issuer: String,
}
serde_json::from_slice::<IssuerOnly>(body)
.ok()
.map(|t| t.issuer)
}
fn check_with(body: &[u8], limit: impl Fn(&str) -> usize) -> Result<(), ErrorResponse> {
if body.len() <= DEFAULT_MAX_DOCUMENT_BYTES {
return Ok(());
}
let type_uri = peek_type(body);
let max = type_uri
.as_deref()
.map_or(DEFAULT_MAX_DOCUMENT_BYTES, &limit);
if body.len() <= max {
return Ok(());
}
let named = match type_uri.as_deref() {
Some(t) if t.chars().count() <= MAX_NAMED_TYPE_CHARS => t,
Some(_) => "its type",
None => "a document of no readable type",
};
let payload: ErrorPayload = RejectReason::MalformedRequest {
reason: format!(
"the document is {} bytes; {named} accepts at most {max}",
body.len()
),
}
.into();
Err(unrouted_error(payload.with_details(
serde_json::json!({ DETAILS_MAX_BYTES: max }),
)))
}
fn peek_type(body: &[u8]) -> Option<String> {
#[derive(serde::Deserialize)]
struct TypeOnly {
#[serde(rename = "type")]
type_uri: String,
}
serde_json::from_slice::<TypeOnly>(body)
.ok()
.map(|t| t.type_uri)
}
fn unrouted_error(payload: ErrorPayload) -> ErrorResponse {
TrustTask {
id: format!("urn:uuid:{}", uuid::Uuid::new_v4()),
thread_id: None,
parent_thread_id: None,
ceremony: None,
type_uri: super::framework_error_type_uri(),
issuer: None,
recipient: None,
issued_at: Some(chrono::Utc::now()),
expires_at: None,
payload,
context: None,
proof: None,
extra: Default::default(),
}
}
#[cfg(test)]
mod tests {
use super::*;
const DID_REGISTER: &str =
<trust_tasks_rs::specs::did_management::did::register::v0_1::Payload as Payload>::TYPE_URI;
const ACL_SHOW: &str = <trust_tasks_rs::specs::acl::show::v0_1::Payload as Payload>::TYPE_URI;
fn register_doc(issuer: &str, bytes: usize) -> Vec<u8> {
serde_json::to_vec(&serde_json::json!({
"type": DID_REGISTER,
"issuer": issuer,
"payload": { "log": "x".repeat(bytes) },
}))
.unwrap()
}
async fn acl_with(did: &str) -> crate::server::store::KeyspaceHandle {
let dir = tempfile::tempdir().expect("tempdir");
let cfg = crate::server::config::StoreConfig {
data_dir: dir.path().to_path_buf(),
..Default::default()
};
std::mem::forget(dir);
let store = crate::server::store::Store::open(&cfg)
.await
.expect("open store");
let acl = store
.keyspace(crate::server::store::KS_ACL)
.expect("acl keyspace");
crate::server::acl::store_acl_entry(
&acl,
&crate::server::acl::AclEntry {
did: did.into(),
role: crate::server::acl::Role::Admin,
label: None,
created_at: 1_700_000_000,
max_total_size: None,
max_did_count: None,
domains: crate::server::domain::DomainScope::All,
},
)
.await
.expect("store entry");
acl
}
#[tokio::test]
async fn a_raised_limit_is_granted_only_to_a_known_issuer() {
let acl = acl_with("did:web:admin.example").await;
let served = [DID_REGISTER];
let budget = LargeDocumentBudget::new();
let big = register_doc("did:web:admin.example", 200 * 1024);
assert!(
check_for_known_issuer(&big, &served, &acl, &budget, "ip:1.2.3.4", 0)
.await
.is_ok()
);
let stranger = register_doc("did:web:stranger.example", 200 * 1024);
assert!(
check_for_known_issuer(&stranger, &served, &acl, &budget, "ip:5.6.7.8", 0)
.await
.is_err()
);
}
#[tokio::test]
async fn repeated_large_documents_from_one_address_are_refused_past_the_budget() {
let acl = acl_with("did:web:admin.example").await;
let served = [DID_REGISTER];
let budget = LargeDocumentBudget::new();
let doc = register_doc("did:web:admin.example", 200 * 1024);
for _ in 0..LARGE_DOCUMENT_BUDGET_PER_WINDOW {
assert!(
check_for_known_issuer(&doc, &served, &acl, &budget, "ip:9.9.9.9", 0)
.await
.is_ok()
);
}
let err = check_for_known_issuer(&doc, &served, &acl, &budget, "ip:9.9.9.9", 0)
.await
.unwrap_err();
assert_eq!(refusal(err)["code"], "unavailable");
assert!(
check_for_known_issuer(&doc, &served, &acl, &budget, "ip:1.1.1.1", 0)
.await
.is_ok()
);
}
#[tokio::test]
async fn a_verified_known_issuer_within_its_budget_passes() {
let acl = acl_with("did:web:admin.example").await;
let served = [DID_REGISTER];
let budget = LargeDocumentBudget::new();
let doc = register_doc("did:web:admin.example", 200 * 1024);
let charge = check_for_known_issuer(&doc, &served, &acl, &budget, "ip:1.2.3.4", 0)
.await
.expect("within budget")
.expect("a large document charges the budget");
settle_large_document_charge(&budget, Some(&charge), Some("did:web:admin.example"));
assert_eq!(
budget.cost("ip:1.2.3.4"),
1,
"a verified claim is not penalised"
);
}
#[tokio::test]
async fn an_unverified_or_mismatched_issuer_penalises_the_address() {
let budget = LargeDocumentBudget::new();
let charge = LargeDocumentCharge {
address: "ip:1.2.3.4".to_string(),
claimed_issuer: "did:web:admin.example".to_string(),
};
settle_large_document_charge(&budget, Some(&charge), None);
assert_eq!(budget.cost("ip:1.2.3.4"), 2);
settle_large_document_charge(&budget, Some(&charge), Some("did:web:someone-else.example"));
assert_eq!(budget.cost("ip:1.2.3.4"), 4);
settle_large_document_charge(&budget, None, None);
assert_eq!(budget.cost("ip:1.2.3.4"), 4);
}
#[test]
fn many_addresses_together_are_held_to_the_global_budget() {
let budget = LargeDocumentBudget::new();
for i in 0..LARGE_DOCUMENT_GLOBAL_BUDGET_PER_WINDOW {
assert!(budget.try_charge(&format!("vid:{i}"), 1_000).is_ok());
}
assert!(
budget.try_charge("vid:fresh", 1_000).is_err(),
"a fresh address does not escape the service-wide cap"
);
assert!(
budget
.try_charge("vid:fresh", 1_000 + LARGE_DOCUMENT_BUDGET_WINDOW_SECS)
.is_ok(),
"the next window has room again"
);
}
#[test]
fn a_full_map_keeps_its_penalties() {
let budget = LargeDocumentBudget::new();
budget.penalize("ip:abuser");
{
let mut state = budget.inner.lock().unwrap();
for i in 0..MAX_TRACKED_ADDRESSES {
state.buckets.insert(
format!("ip:{i}"),
Bucket {
spent: 0,
window_start: 0,
cost: 1,
},
);
}
}
assert!(budget.try_charge("ip:new", 10_000).is_ok());
assert_eq!(budget.cost("ip:abuser"), 2);
}
#[test]
fn a_full_map_of_live_addresses_refuses_a_new_one() {
let budget = LargeDocumentBudget::new();
{
let mut state = budget.inner.lock().unwrap();
for i in 0..MAX_TRACKED_ADDRESSES {
state.buckets.insert(
format!("ip:{i}"),
Bucket {
spent: 1,
window_start: 10_000,
cost: 1,
},
);
}
}
assert!(budget.try_charge("ip:new", 10_000).is_err());
assert!(
budget.try_charge("ip:0", 10_000).is_ok(),
"a tracked address still charges"
);
}
#[tokio::test]
async fn small_documents_are_never_charged() {
let acl = acl_with("did:web:admin.example").await;
let served = [DID_REGISTER];
let budget = LargeDocumentBudget::new();
let small = register_doc("did:web:admin.example", 10);
assert!(small.len() <= DEFAULT_MAX_DOCUMENT_BYTES);
for _ in 0..(LARGE_DOCUMENT_BUDGET_PER_WINDOW * 3) {
let charge = check_for_known_issuer(&small, &served, &acl, &budget, "ip:2.2.2.2", 0)
.await
.expect("small document admitted");
assert!(
charge.is_none(),
"a small document never charges the budget"
);
}
}
fn served_register(type_uri: &str) -> usize {
max_in(type_uri, DECLARED, &[DID_REGISTER])
}
fn document_of(type_uri: &str, len: usize) -> Vec<u8> {
let head = format!(r#"{{"type":"{type_uri}","payload":{{"pad":""#);
let tail = r#""}}"#;
let pad = len - head.len() - tail.len();
let body = format!("{head}{}{tail}", "x".repeat(pad));
assert_eq!(body.len(), len);
body.into_bytes()
}
fn refusal(err: ErrorResponse) -> serde_json::Value {
serde_json::to_value(&err.payload).unwrap()
}
#[test]
fn every_declared_type_is_raised_above_the_default() {
for (uri, max) in DECLARED {
assert!(
*max > DEFAULT_MAX_DOCUMENT_BYTES,
"{uri} declares {max}, which is not above the default; drop it"
);
}
let mut uris: Vec<&str> = DECLARED.iter().map(|(u, _)| *u).collect();
uris.sort_unstable();
uris.dedup();
assert_eq!(uris.len(), DECLARED.len(), "a type is declared twice");
}
#[test]
fn an_undeclared_type_takes_the_default() {
assert_eq!(served_register(ACL_SHOW), DEFAULT_MAX_DOCUMENT_BYTES);
assert_eq!(
served_register(DID_REGISTER),
DID_REGISTER_MAX_DOCUMENT_BYTES
);
}
#[test]
fn a_declared_type_that_is_not_served_takes_the_default() {
assert_eq!(
max_in(DID_REGISTER, DECLARED, &[]),
DEFAULT_MAX_DOCUMENT_BYTES
);
assert_eq!(largest_in(DECLARED, &[]), DEFAULT_MAX_DOCUMENT_BYTES);
assert_eq!(
largest_in(DECLARED, &[DID_REGISTER]),
DID_REGISTER_MAX_DOCUMENT_BYTES
);
assert_eq!(
max_document_bytes(DID_REGISTER, &[]),
DEFAULT_MAX_DOCUMENT_BYTES
);
assert_eq!(
max_document_bytes(DID_REGISTER, &[DID_REGISTER]),
DID_REGISTER_MAX_DOCUMENT_BYTES
);
assert!(largest_max_document_bytes(&[]) >= DEFAULT_MAX_DOCUMENT_BYTES);
}
#[test]
fn a_document_at_its_types_limit_is_admitted() {
assert!(check(&document_of(ACL_SHOW, DEFAULT_MAX_DOCUMENT_BYTES), &[]).is_ok());
assert!(
check_with(
&document_of(DID_REGISTER, DID_REGISTER_MAX_DOCUMENT_BYTES),
served_register
)
.is_ok()
);
}
#[test]
fn a_document_over_the_default_is_refused_for_a_type_that_declares_nothing() {
let err = check(&document_of(ACL_SHOW, DEFAULT_MAX_DOCUMENT_BYTES + 1), &[]).unwrap_err();
let payload = refusal(err);
assert_eq!(payload["code"], "malformedRequest");
assert_eq!(
payload["details"][DETAILS_MAX_BYTES],
DEFAULT_MAX_DOCUMENT_BYTES
);
}
#[test]
fn a_1mib_register_document_is_admitted_and_over_the_limit_is_refused() {
assert!(
check(
&document_of(DID_REGISTER, DID_REGISTER_MAX_DOCUMENT_BYTES),
&[DID_REGISTER]
)
.is_ok()
);
let err = check(
&document_of(DID_REGISTER, DID_REGISTER_MAX_DOCUMENT_BYTES + 1),
&[DID_REGISTER],
)
.unwrap_err();
let payload = refusal(err);
assert_eq!(payload["code"], "malformedRequest");
assert_eq!(
payload["details"][DETAILS_MAX_BYTES],
DID_REGISTER_MAX_DOCUMENT_BYTES
);
}
#[test]
fn an_unserved_type_claiming_a_large_limit_gets_only_the_default() {
let err = check(
&document_of(DID_REGISTER, DEFAULT_MAX_DOCUMENT_BYTES + 1),
&[], )
.unwrap_err();
let payload = refusal(err);
assert_eq!(
payload["details"][DETAILS_MAX_BYTES],
DEFAULT_MAX_DOCUMENT_BYTES
);
}
#[test]
fn the_type_is_the_top_level_one_wherever_it_sits() {
let pad = "x".repeat(DEFAULT_MAX_DOCUMENT_BYTES);
let late = format!(
r#"{{"payload":{{"type":"{ACL_SHOW}","pad":"{pad}"}},"type":"{DID_REGISTER}"}}"#
);
assert!(check_with(late.as_bytes(), served_register).is_ok());
let nested_only = format!(r#"{{"payload":{{"type":"{DID_REGISTER}","pad":"{pad}"}}}}"#);
assert_eq!(
refusal(check_with(nested_only.as_bytes(), served_register).unwrap_err())["details"]
[DETAILS_MAX_BYTES],
DEFAULT_MAX_DOCUMENT_BYTES
);
}
#[test]
fn a_late_type_in_the_body_is_still_read() {
let pad = "x".repeat(DEFAULT_MAX_DOCUMENT_BYTES);
let late = format!(r#"{{"pad":"{pad}","type":"{DID_REGISTER}"}}"#);
assert!(check_with(late.as_bytes(), served_register).is_ok());
}
#[test]
fn a_duplicated_type_takes_the_default() {
let pad = "x".repeat(DEFAULT_MAX_DOCUMENT_BYTES);
let body = format!(r#"{{"type":"{ACL_SHOW}","type":"{DID_REGISTER}","pad":"{pad}"}}"#);
assert_eq!(
refusal(check_with(body.as_bytes(), served_register).unwrap_err())["details"]
[DETAILS_MAX_BYTES],
DEFAULT_MAX_DOCUMENT_BYTES
);
}
#[test]
fn a_refusal_does_not_echo_an_oversized_type() {
let huge = "t".repeat(DEFAULT_MAX_DOCUMENT_BYTES + 1);
let body = format!(r#"{{"type":"{huge}"}}"#);
let payload = refusal(check(body.as_bytes(), &[]).unwrap_err());
let text = payload.to_string();
assert!(
text.len() < 1024,
"the refusal echoes the body: {} bytes",
text.len()
);
assert_eq!(
payload["details"],
serde_json::json!({ DETAILS_MAX_BYTES: DEFAULT_MAX_DOCUMENT_BYTES })
);
}
#[test]
fn an_oversized_body_with_no_readable_type_takes_the_default() {
let junk = vec![b'['; DEFAULT_MAX_DOCUMENT_BYTES + 1];
let payload = refusal(check(&junk, &[]).unwrap_err());
assert_eq!(payload["code"], "malformedRequest");
assert_eq!(
payload["details"][DETAILS_MAX_BYTES],
DEFAULT_MAX_DOCUMENT_BYTES
);
}
}