use std::cell::Cell;
thread_local! {
static DEPTH: Cell<u32> = const { Cell::new(0) };
static WORK: Cell<u64> = const { Cell::new(0) };
static WORK_LIMIT: Cell<u64> = const { Cell::new(u64::MAX) };
static STEPS: Cell<u64> = const { Cell::new(0) };
static STEP_LIMIT: Cell<u64> = const { Cell::new(u64::MAX) };
static SPENT: Cell<bool> = const { Cell::new(false) };
}
const MAX_DEPTH: u32 = 48;
pub(super) const MAX_PARSE_WORK_CEILING: u64 = 20_000;
const WORK_BASE: u64 = 2_048;
const WORK_PER_BYTE: u64 = 8;
const STEP_BYTES: usize = 16;
pub(super) const STEPS_PER_BYTE: u64 = 8;
const STEP_BASE: u64 = 4_096;
pub(super) fn reset(input_len: usize) {
let len = input_len as u64;
DEPTH.with(|d| d.set(0));
WORK.with(|w| w.set(0));
WORK_LIMIT.with(|l| l.set((WORK_BASE + WORK_PER_BYTE * len).min(MAX_PARSE_WORK_CEILING)));
STEPS.with(|s| s.set(0));
STEP_LIMIT.with(|l| l.set(STEP_BASE + STEPS_PER_BYTE * len));
SPENT.with(|s| s.set(false));
}
pub(super) fn spent() -> bool {
SPENT.with(Cell::get)
}
pub(super) fn charge(attempts: u64, bytes: usize) -> bool {
let over = STEPS.with(|s| {
let n = s.get().saturating_add(attempts + (bytes / STEP_BYTES) as u64);
s.set(n);
n > STEP_LIMIT.with(Cell::get)
});
if over {
SPENT.with(|s| s.set(true));
}
!over
}
#[cfg(test)]
pub(super) fn work() -> u64 {
WORK.with(Cell::get)
}
#[cfg(test)]
pub(super) fn steps() -> u64 {
STEPS.with(Cell::get)
}
pub(super) struct DepthGuard;
impl DepthGuard {
pub(super) fn enter() -> Option<Self> {
let over = WORK.with(|w| {
let n = w.get().saturating_add(1);
w.set(n);
n > WORK_LIMIT.with(Cell::get)
});
if over {
SPENT.with(|s| s.set(true));
return None;
}
DEPTH.with(|d| {
if d.get() >= MAX_DEPTH {
None
} else {
d.set(d.get() + 1);
Some(DepthGuard)
}
})
}
}
impl Drop for DepthGuard {
fn drop(&mut self) {
DEPTH.with(|d| d.set(d.get().saturating_sub(1)));
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_charge_counts_its_attempts_plus_one_step_per_step_bytes() {
reset(0);
assert!(charge(1, 0));
assert!(charge(1, STEP_BYTES * 3));
assert_eq!(steps(), 1 + 4);
}
#[test]
fn running_out_of_steps_is_sticky_until_the_next_reset() {
reset(0);
assert!(!spent());
assert!(!charge(0, STEP_BYTES * (STEP_BASE as usize + 1)));
assert!(spent());
assert!(!charge(1, 0), "a spent budget must keep refusing");
reset(0);
assert!(!spent());
assert!(charge(1, 0));
}
#[test]
fn the_step_allowance_grows_with_the_input() {
reset(1_000);
let allowance = (STEP_BASE + STEPS_PER_BYTE * 1_000) as usize;
assert!(charge(0, STEP_BYTES * allowance));
assert!(!charge(1, 0));
}
#[test]
fn running_out_of_work_is_sticky_too() {
reset(0);
let guards: Vec<_> = (0..MAX_DEPTH).map_while(|_| DepthGuard::enter()).collect();
assert_eq!(guards.len(), MAX_DEPTH as usize, "the depth cap is not the budget");
assert!(!spent(), "hitting the depth cap alone must not mark the budget spent");
drop(guards);
while DepthGuard::enter().is_some() {}
assert!(spent());
}
}