Skip to main content

rusty_time_core/
server.rs

1//! Server-side policy: client tracking, rate limiting, and interleaved mode.
2//!
3//! All decisions, no I/O. The daemon owns sockets and NTS; this module owns
4//! *what the answer should be*, which is what makes it testable without a
5//! network and portable to wasm.
6//!
7//! The client key is generic on purpose: the core must not know what an IP
8//! address is (mission plan §4 — the core knows bytes and timestamps, never a
9//! product type). The daemon instantiates it with the peer address.
10
11use crate::ntp::NtpTimestamp;
12use std::collections::HashMap;
13use std::hash::{Hash, Hasher};
14
15/// Rate-limit policy, mirroring chrony's `ratelimit interval/burst/leak`.
16#[derive(Clone, Copy, Debug, PartialEq)]
17pub struct RateLimitConfig {
18    /// log2 of the mean seconds between responses to one client. 3 = one
19    /// response per 8 s, chrony's default.
20    pub interval_log2: i8,
21    /// How many responses a client may take back to back after being quiet.
22    pub burst: u32,
23    /// Emit a Kiss-o'-Death to one dropped request in 2^leak_shift. Never 0:
24    /// answering *every* dropped request turns the limiter into the
25    /// amplifier it exists to prevent.
26    pub leak_shift: u8,
27    /// Ceiling on responses per second **across all clients**. 0 disables it.
28    ///
29    /// Per-client limiting alone is defeated by table churn: once the client
30    /// population exceeds the table, every request arrives from an address we
31    /// have forgotten, gets a fresh bucket, and is answered. TIMECORP S12b
32    /// (100k addresses into a 16k table) showed the reply ratio climbing back
33    /// to 100% for exactly that reason. A global bucket is the backstop that
34    /// bounds total output no matter how the address space is shuffled.
35    ///
36    /// Set high enough not to hinder a busy legitimate server; it exists to
37    /// bound the worst case, not to shape normal traffic.
38    pub global_rate_hz: f64,
39    /// Global burst allowance, in responses.
40    pub global_burst: f64,
41}
42
43impl Default for RateLimitConfig {
44    fn default() -> Self {
45        RateLimitConfig {
46            interval_log2: 3,
47            burst: 8,
48            leak_shift: 2,
49            global_rate_hz: 20_000.0,
50            global_burst: 40_000.0,
51        }
52    }
53}
54
55/// What the server should do with one request.
56#[derive(Clone, Copy, Debug, PartialEq, Eq)]
57pub enum Disposition {
58    /// Answer normally.
59    Respond,
60    /// Over the limit: answer with a Kiss-o'-Death RATE so a well-behaved
61    /// client backs off.
62    KissOfDeath,
63    /// Over the limit and not this client's turn for a KoD: say nothing.
64    Drop,
65}
66
67/// Mark a server response's timestamps so the two are distinguishable by their
68/// lowest bit: **receive has bit 0 set, transmit has it clear** (RFC 9769 and
69/// chrony's `ntp_core.c`).
70///
71/// Two things depend on this. A server can then recognise an interleaved
72/// request statelessly — an origin field with bit 0 set is echoing a *receive*
73/// timestamp — and receive can never accidentally equal transmit, which would
74/// make the mode ambiguous. The cost is the bottom bit of a 232-picosecond
75/// unit, far below any clock's resolution.
76pub fn mark_server_timestamps(receive: &mut NtpTimestamp, transmit: &mut NtpTimestamp) {
77    receive.0 |= 1;
78    transmit.0 &= !1;
79}
80
81/// Which timestamps a response should carry.
82#[derive(Clone, Copy, Debug, PartialEq, Eq)]
83pub enum ResponseMode {
84    /// The ordinary exchange: receive and transmit from *this* exchange.
85    Basic,
86    /// Interleaved: the reply carries **this** exchange's receive timestamp
87    /// (so the client can interleave again next time) together with the
88    /// *actual, post-send* transmit timestamp of the earlier response the
89    /// client named in its origin field. That transmit value is the whole
90    /// point — a basic reply must write its transmit field before the packet
91    /// leaves, so it can only ever be an estimate.
92    ///
93    /// Two field rules, both of which fail silently rather than loudly when
94    /// broken (RFC 9769; chrony `ntp_core.c` lines 1241/1251/1290, 2744-2754):
95    ///
96    /// * **origin** echoes the request's *receive* field, not its transmit —
97    ///   that is how the client recognises an interleaved reply.
98    /// * **receive** is the current exchange's, not the previous one's.
99    ///   Reporting the previous receive pairs the client's T1/T4 with a
100    ///   mismatched T2/T3 and shows up as a delay inflated by exactly one
101    ///   poll interval — measured against chrony as 4.009 s on a 4 s poll.
102    Interleaved {
103        /// The true transmit timestamp of the response the client's origin
104        /// field identifies.
105        prev_transmit: NtpTimestamp,
106    },
107}
108
109/// Per-client state: enough for rate limiting, interleaved mode and the MRU
110/// report, and no more — this is multiplied by every client that has ever
111/// spoken to us.
112#[derive(Clone, Copy, Debug)]
113pub struct ClientRecord {
114    pub last_seen: f64,
115    /// Token bucket level, in responses.
116    tokens: f64,
117    pub requests: u64,
118    pub responses: u64,
119    pub dropped: u64,
120    /// Dropped requests since the last Kiss-o'-Death, for deterministic leak.
121    drops_since_kod: u32,
122    /// T2 of the last request we accepted.
123    pub last_receive: Option<NtpTimestamp>,
124    /// The true transmit timestamp of our last response, once the driver
125    /// reports it. `None` until then, which is why interleaved mode cannot
126    /// answer the very first request.
127    pub last_transmit: Option<NtpTimestamp>,
128    /// The receive timestamp we put in our last response. A client asks for
129    /// interleaved mode by echoing exactly this.
130    pub last_receive_sent: Option<NtpTimestamp>,
131    /// Whether the last answered request was actually served interleaved.
132    ///
133    /// Distinct from "we *could* serve it interleaved": every client we have
134    /// answered once is capable, but only a client that asks is using it, and
135    /// an operator reading a client log needs the second fact.
136    pub interleaved_now: bool,
137}
138
139impl ClientRecord {
140    fn new(now: f64, burst: u32) -> Self {
141        ClientRecord {
142            last_seen: now,
143            tokens: burst as f64,
144            requests: 0,
145            responses: 0,
146            dropped: 0,
147            drops_since_kod: 0,
148            last_receive: None,
149            last_transmit: None,
150            last_receive_sent: None,
151            interleaved_now: false,
152        }
153    }
154}
155
156/// Aggregate counters for the `status.serverstats` op.
157#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
158pub struct ServerStats {
159    pub requests: u64,
160    pub responses: u64,
161    pub dropped_rate_limit: u64,
162    pub kiss_of_death: u64,
163    pub interleaved_responses: u64,
164    /// Requests refused before any per-client work (bad mode, too short).
165    pub refused: u64,
166    /// Clients evicted from the table because it was full.
167    pub evicted: u64,
168}
169
170/// End-of-list sentinel for the intrusive recency links.
171const NIL: u32 = u32::MAX;
172
173/// Odd 64-bit multiplier for the key mixer. Any odd constant with a good bit
174/// spread works; this one is xxHash's prime 1.
175const MIX: u64 = 0x9e37_79b1_85eb_ca87;
176
177/// A fast, **seeded** hasher for short keys.
178///
179/// The default `HashMap` hasher is SipHash-1-3, chosen for resistance to
180/// collision floods. That resistance is not optional here — the key is a
181/// client's source address, which an attacker picks — but SipHash's cost is
182/// out of proportion to a 4-to-17-byte key: with everything else in the client
183/// table fixed, callgrind still attributed ~32% of the server's per-request
184/// instructions to hashing one address once.
185///
186/// So this keeps the property and drops the price. The seed is drawn from the
187/// OS once per process via `RandomState`, exactly as SipHash's keys are, so an
188/// attacker cannot compute which addresses collide without first learning a
189/// secret they never see. What is given up is SipHash's *proof* against an
190/// adversary who somehow does learn the seed; what is kept is the practical
191/// defence, on a table that is additionally bounded to a fixed capacity with
192/// LRU eviction, so even a successful collision attack cannot grow a chain
193/// without bound.
194#[derive(Clone, Copy)]
195pub struct ClientHashBuilder {
196    seed: u64,
197}
198
199impl Default for ClientHashBuilder {
200    fn default() -> Self {
201        // `RUSTY_TIME_HASH_SEED` pins the seed. It exists for measurement: a
202        // random seed changes which keys collide, which moves the probe count,
203        // which makes an instruction-count harness reproducible only to about
204        // 0.002% instead of exactly. That is far below any effect worth acting
205        // on, but an instrument that is exact is worth more than one that is
206        // nearly exact, and the pin costs one environment read per table.
207        //
208        // It is emphatically NOT for production: a known seed is a known
209        // collision set, which is the property this hasher is seeded to deny.
210        if let Ok(pinned) = std::env::var("RUSTY_TIME_HASH_SEED")
211            && let Ok(seed) = pinned.parse::<u64>()
212        {
213            return ClientHashBuilder { seed };
214        }
215        // One OS-random draw per table, reusing std's entropy source rather
216        // than adding a dependency for it.
217        use std::hash::BuildHasher as _;
218        let seed = std::collections::hash_map::RandomState::new().hash_one(0xA5A5_5A5Au64);
219        ClientHashBuilder { seed }
220    }
221}
222
223impl std::hash::BuildHasher for ClientHashBuilder {
224    type Hasher = ClientHasher;
225    fn build_hasher(&self) -> ClientHasher {
226        ClientHasher { state: self.seed }
227    }
228}
229
230pub struct ClientHasher {
231    state: u64,
232}
233
234impl ClientHasher {
235    #[inline]
236    fn mix(&mut self, value: u64) {
237        self.state = (self.state ^ value).wrapping_mul(MIX);
238    }
239}
240
241impl Hasher for ClientHasher {
242    #[inline]
243    fn write(&mut self, bytes: &[u8]) {
244        // `as_chunks` rather than `chunks_exact(8)`: the width is a constant, so
245        // this yields `&[u8; 8]` directly and the fallible conversion in the
246        // loop disappears.
247        let (chunks, rest) = bytes.as_chunks::<8>();
248        for chunk in chunks {
249            self.mix(u64::from_le_bytes(*chunk));
250        }
251        if !rest.is_empty() {
252            let mut buf = [0u8; 8];
253            buf[..rest.len()].copy_from_slice(rest);
254            // Fold the length in so trailing zero bytes cannot alias a shorter
255            // key against a longer one.
256            self.mix(u64::from_le_bytes(buf) ^ (rest.len() as u64) << 56);
257        }
258    }
259
260    #[inline]
261    fn write_u8(&mut self, value: u8) {
262        self.mix(value as u64);
263    }
264
265    #[inline]
266    fn write_u32(&mut self, value: u32) {
267        self.mix(value as u64);
268    }
269
270    #[inline]
271    fn write_u64(&mut self, value: u64) {
272        self.mix(value);
273    }
274
275    #[inline]
276    fn write_usize(&mut self, value: usize) {
277        self.mix(value as u64);
278    }
279
280    #[inline]
281    fn finish(&self) -> u64 {
282        // splitmix64's finalizer: full avalanche in a handful of instructions,
283        // which is what stops near-adjacent addresses landing in near-adjacent
284        // buckets.
285        let mut z = self.state;
286        z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
287        z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
288        z ^ (z >> 31)
289    }
290}
291
292/// One client's storage: its key, its record, and its place in the recency
293/// list. Slots are stable — an index handed out stays valid until the client
294/// is evicted — which is what makes the list links safe as plain integers.
295struct Slot<K> {
296    key: K,
297    record: ClientRecord,
298    /// Toward the most-recently-used end.
299    prev: u32,
300    /// Toward the least-recently-used end.
301    next: u32,
302    /// Bumped every time this slot is handed to a new client, so a handle kept
303    /// across an eviction is detected rather than silently addressing whoever
304    /// took the slot over.
305    generation: u32,
306}
307
308/// A resolved position in the table.
309///
310/// One request touches the same client four times — admit, choose a response
311/// mode, record what was sent, then record the true transmit time after the
312/// packet leaves. Looking the key up each time cost four hashes of the same
313/// address, which callgrind put at ~51% of what the server does per request
314/// once the recency index was fixed. A handle turns the three follow-ups into
315/// array indexing.
316///
317/// It is deliberately not a raw index: the generation makes a handle that
318/// outlived its client detectably stale instead of quietly wrong.
319#[derive(Clone, Copy, Debug, PartialEq, Eq)]
320pub struct ClientHandle {
321    slot: u32,
322    generation: u32,
323}
324
325impl ClientHandle {
326    /// A handle that resolves to nothing — what a request rejected by the
327    /// global limiter gets, having never reached a per-client record.
328    pub const INVALID: ClientHandle = ClientHandle {
329        slot: NIL,
330        generation: 0,
331    };
332}
333
334/// Bounded most-recently-used client table.
335///
336/// The bound is the point: an unbounded map is a memory-exhaustion lever for
337/// anyone willing to spoof source addresses. When full, the least recently
338/// seen client is evicted — losing its interleaved state, which costs it one
339/// exchange, not correctness.
340///
341/// **Eviction is indexed, not scanned.** The obvious implementation — walk the
342/// map for the oldest `last_seen` — is O(capacity) per admission, and a public
343/// server facing more clients than the table holds evicts on nearly every
344/// packet. TIMECORP S12b (100k clients into a 16k table) went from "runs" to
345/// "does not finish" on exactly that, which is what the scenario is for. The
346/// `order` index makes it O(log n) *and* deterministic, where a `HashMap` scan
347/// depends on iteration order that differs between instances.
348pub struct ClientTable<K: Eq + Hash + Ord + Clone> {
349    /// Key to slot. The only hashed structure, and the reason a request costs
350    /// one hash instead of several.
351    index: HashMap<K, u32, ClientHashBuilder>,
352    /// Records in stable storage. Slots are handed out from `free` and never
353    /// move, which is what lets the ordering below be pointers rather than
354    /// comparisons.
355    slots: Vec<Slot<K>>,
356    free: Vec<u32>,
357    /// Ends of the intrusive most-recently-used list threaded through `slots`.
358    mru: u32,
359    lru: u32,
360    /// Global token bucket: level and the time it was last refilled.
361    global_tokens: f64,
362    global_last: Option<f64>,
363    /// Global drops since the last global Kiss-o'-Death.
364    global_drops_since_kod: u32,
365    capacity: usize,
366    config: RateLimitConfig,
367    /// Tokens a client earns per second — `2^-interval_log2`, precomputed.
368    ///
369    /// It derives only from `config`, which is fixed at construction, so
370    /// recomputing it per request was a `powi` call on the hot path for a
371    /// constant. Callgrind put the server's per-request cost at 3053 Ir; this
372    /// is one of the cheaper pieces of that, and it is free to remove.
373    refill_per_s: f64,
374    pub stats: ServerStats,
375}
376
377impl<K: Eq + Hash + Ord + Clone> ClientTable<K> {
378    pub fn new(capacity: usize, config: RateLimitConfig) -> Self {
379        ClientTable {
380            index: HashMap::with_capacity_and_hasher(capacity.max(1), ClientHashBuilder::default()),
381            slots: Vec::with_capacity(capacity.max(1)),
382            free: Vec::new(),
383            mru: NIL,
384            lru: NIL,
385            global_tokens: config.global_burst,
386            global_last: None,
387            global_drops_since_kod: 0,
388            capacity: capacity.max(1),
389            config,
390            refill_per_s: 2f64.powi(-(config.interval_log2 as i32)),
391            stats: ServerStats::default(),
392        }
393    }
394
395    /// Detach a slot from the recency list.
396    fn unlink(&mut self, i: u32) {
397        let (prev, next) = {
398            let slot = &self.slots[i as usize];
399            (slot.prev, slot.next)
400        };
401        if prev == NIL {
402            self.mru = next;
403        } else {
404            self.slots[prev as usize].next = next;
405        }
406        if next == NIL {
407            self.lru = prev;
408        } else {
409            self.slots[next as usize].prev = prev;
410        }
411    }
412
413    /// Put a detached slot at the most-recent end.
414    fn link_front(&mut self, i: u32) {
415        let old = self.mru;
416        {
417            let slot = &mut self.slots[i as usize];
418            slot.prev = NIL;
419            slot.next = old;
420        }
421        if old == NIL {
422            self.lru = i;
423        } else {
424            self.slots[old as usize].prev = i;
425        }
426        self.mru = i;
427    }
428
429    /// Move a client to the most-recent end of the eviction order.
430    ///
431    /// Six pointer writes, no hashing, no comparisons, no allocation. The
432    /// previous form kept a `BTreeSet<(seq, K)>` and did a remove plus an
433    /// insert — two O(log n) tree walks and two key clones — on *every*
434    /// request. Callgrind attributed 23% of the whole server hot path to that
435    /// tree's `search_tree`, for an index that is only ever read when the
436    /// table is full and something has to be evicted.
437    fn touch(&mut self, i: u32) {
438        if self.mru == i {
439            return; // already most-recent; the common case for a chatty client
440        }
441        self.unlink(i);
442        self.link_front(i);
443    }
444
445    /// The real per-client footprint: the record, the slot links that order
446    /// it, and the index entry that finds it.
447    ///
448    /// Reported by the type rather than estimated by the caller, because the
449    /// corpus quotes this number and an estimate drifts silently when the
450    /// structure changes. It did: the figure used to be `size_of::<ClientRecord>()`
451    /// alone, which stopped being the whole story the moment records moved
452    /// into slots.
453    pub fn bytes_per_client() -> usize {
454        core::mem::size_of::<Slot<K>>()
455            + core::mem::size_of::<K>()          // the index's own copy of the key
456            + core::mem::size_of::<u32>()        // the slot number it maps to
457            + 1 // hashbrown's control byte
458    }
459
460    pub fn len(&self) -> usize {
461        self.index.len()
462    }
463
464    pub fn is_empty(&self) -> bool {
465        self.index.is_empty()
466    }
467
468    pub fn get(&self, key: &K) -> Option<&ClientRecord> {
469        let i = *self.index.get(key)?;
470        Some(&self.slots[i as usize].record)
471    }
472
473    /// The MRU report: most recently seen first, at most `limit` entries.
474    ///
475    /// Walks the recency list, which is **already** in this order — `admit`
476    /// moves a client to the front and sets `last_seen` in the same breath, so
477    /// list order and descending `last_seen` are the same thing.
478    ///
479    /// The previous form cloned every record in the table into a `Vec`, sorted
480    /// all of them, and threw away all but `limit`. At the daemon's capacity
481    /// of 16384 that is ~1.5 MiB copied and an O(n log n) sort to answer a
482    /// ten-row status query. This is O(limit) and allocates once, for exactly
483    /// the rows returned.
484    pub fn most_recent(&self, limit: usize) -> Vec<(K, ClientRecord)> {
485        let mut out = Vec::with_capacity(limit.min(self.index.len()));
486        let mut at = self.mru;
487        while at != NIL && out.len() < limit {
488            let slot = &self.slots[at as usize];
489            out.push((slot.key.clone(), slot.record));
490            at = slot.next;
491        }
492        out
493    }
494
495    /// Drop the least recently seen client. O(1): it is the list's tail.
496    fn evict_one(&mut self) {
497        let victim = self.lru;
498        if victim == NIL {
499            return;
500        }
501        self.unlink(victim);
502        let key = self.slots[victim as usize].key.clone();
503        self.index.remove(&key);
504        self.free.push(victim);
505        self.stats.evicted += 1;
506    }
507
508    /// Take a slot for a new client, reusing an evicted one where possible.
509    fn alloc_slot(&mut self, key: K, record: ClientRecord) -> u32 {
510        let i = match self.free.pop() {
511            Some(i) => {
512                let slot = &mut self.slots[i as usize];
513                slot.key = key;
514                slot.record = record;
515                slot.prev = NIL;
516                slot.next = NIL;
517                // New occupant, new generation: any handle still naming this
518                // slot from its previous client now fails to resolve.
519                slot.generation = slot.generation.wrapping_add(1);
520                i
521            }
522            None => {
523                self.slots.push(Slot {
524                    key,
525                    record,
526                    prev: NIL,
527                    next: NIL,
528                    generation: 0,
529                });
530                (self.slots.len() - 1) as u32
531            }
532        };
533        self.link_front(i);
534        i
535    }
536
537    /// Turn a handle back into a slot, or `None` if it has gone stale.
538    fn resolve(&self, handle: ClientHandle) -> Option<usize> {
539        let i = handle.slot as usize;
540        let slot = self.slots.get(i)?;
541        (slot.generation == handle.generation).then_some(i)
542    }
543
544    fn handle_for(&self, slot: u32) -> ClientHandle {
545        ClientHandle {
546            slot,
547            generation: self.slots[slot as usize].generation,
548        }
549    }
550
551    /// Admit one request: refill the client's bucket, decide its fate, and
552    /// record it. `now` is monotonic seconds.
553    pub fn admit(&mut self, key: &K, now: f64) -> Disposition {
554        self.admit_handle(key, now).0
555    }
556
557    /// `admit`, also returning the handle that addresses this client, so the
558    /// rest of the request never has to hash the key again.
559    pub fn admit_handle(&mut self, key: &K, now: f64) -> (Disposition, ClientHandle) {
560        self.stats.requests += 1;
561
562        // Global ceiling first. It is checked before the per-client bucket so
563        // that a churned-address flood — which defeats per-client limiting by
564        // never reusing an address — still cannot make us answer without
565        // bound.
566        if self.config.global_rate_hz > 0.0 {
567            let last = self.global_last.unwrap_or(now);
568            let elapsed = (now - last).max(0.0);
569            self.global_tokens = (self.global_tokens + elapsed * self.config.global_rate_hz)
570                .min(self.config.global_burst);
571            self.global_last = Some(now);
572
573            if self.global_tokens < 1.0 {
574                self.stats.dropped_rate_limit += 1;
575                self.global_drops_since_kod += 1;
576                let period = 1u32 << self.config.leak_shift.min(16);
577                if self.global_drops_since_kod >= period {
578                    self.global_drops_since_kod = 0;
579                    self.stats.kiss_of_death += 1;
580                    return (Disposition::KissOfDeath, ClientHandle::INVALID);
581                }
582                return (Disposition::Drop, ClientHandle::INVALID);
583            }
584        }
585
586        // One hash for the common case (a client we already know), two for a
587        // client we have never seen. The previous form hashed three times per
588        // request: `contains_key`, then `touch`'s `get_mut`, then a final
589        // `get_mut` to reach the record.
590        let slot = match self.index.get(key) {
591            Some(&i) => {
592                self.touch(i);
593                i
594            }
595            None => {
596                if self.index.len() >= self.capacity {
597                    self.evict_one();
598                }
599                let i = self.alloc_slot(key.clone(), ClientRecord::new(now, self.config.burst));
600                self.index.insert(key.clone(), i);
601                i
602            }
603        };
604
605        let config = self.config;
606        let rate = self.refill_per_s;
607        let handle = self.handle_for(slot);
608        let record = &mut self.slots[slot as usize].record;
609
610        // Refill: one token per 2^interval seconds since we last saw them.
611        let elapsed = (now - record.last_seen).max(0.0);
612        record.tokens = (record.tokens + elapsed * rate).min(config.burst as f64);
613        record.last_seen = now;
614        record.requests += 1;
615
616        if record.tokens >= 1.0 {
617            record.tokens -= 1.0;
618            record.responses += 1;
619            self.stats.responses += 1;
620            // Spend a global token only when a response is actually produced.
621            self.global_tokens -= 1.0;
622            return (Disposition::Respond, handle);
623        }
624
625        record.dropped += 1;
626        record.drops_since_kod += 1;
627        self.stats.dropped_rate_limit += 1;
628
629        // Deterministic leak: every 2^leak_shift-th drop gets a KoD. A
630        // deterministic rule is testable, and the client cannot tell the
631        // difference from a probabilistic one.
632        let period = 1u32 << config.leak_shift.min(16);
633        if record.drops_since_kod >= period {
634            record.drops_since_kod = 0;
635            self.stats.kiss_of_death += 1;
636            (Disposition::KissOfDeath, handle)
637        } else {
638            (Disposition::Drop, handle)
639        }
640    }
641
642    /// Decide basic vs interleaved for an admitted request.
643    ///
644    /// The client signals interleaved mode by setting its origin timestamp to
645    /// the *receive* timestamp we sent last time, rather than the transmit
646    /// timestamp. That is unforgeable in the useful sense: only a client that
647    /// actually saw our last response knows it.
648    pub fn response_mode(&mut self, key: &K, request_origin: NtpTimestamp) -> ResponseMode {
649        match self.index.get(key) {
650            Some(&i) => {
651                let handle = self.handle_for(i);
652                self.response_mode_at(handle, request_origin)
653            }
654            None => ResponseMode::Basic,
655        }
656    }
657
658    /// `response_mode` addressed by handle — no hashing.
659    pub fn response_mode_at(
660        &mut self,
661        handle: ClientHandle,
662        request_origin: NtpTimestamp,
663    ) -> ResponseMode {
664        // One lookup, not two. The original read the record with `get` and then
665        // re-found the same entry with `get_mut` to store `interleaved_now` —
666        // a second hash of the same key on the per-request path, where hashing
667        // was measured at ~38% of all instructions.
668        let Some(i) = self.resolve(handle) else {
669            return ResponseMode::Basic;
670        };
671        let record = &mut self.slots[i].record;
672        let (Some(sent_receive), Some(prev_transmit)) =
673            (record.last_receive_sent, record.last_transmit)
674        else {
675            return ResponseMode::Basic;
676        };
677        // The client names a specific earlier response by echoing the receive
678        // timestamp we reported for it. We keep one slot, so only the most
679        // recent qualifies; anything older falls back to basic rather than
680        // answering with a transmit timestamp from the wrong exchange.
681        let interleaved = request_origin == sent_receive && !request_origin.is_zero();
682        record.interleaved_now = interleaved;
683        if interleaved {
684            self.stats.interleaved_responses += 1;
685            ResponseMode::Interleaved { prev_transmit }
686        } else {
687            ResponseMode::Basic
688        }
689    }
690
691    /// Record what we received and what we told the client, after answering.
692    pub fn note_response(&mut self, key: &K, receive: NtpTimestamp, receive_sent: NtpTimestamp) {
693        if let Some(&i) = self.index.get(key) {
694            let handle = self.handle_for(i);
695            self.note_response_at(handle, receive, receive_sent);
696        }
697    }
698
699    /// `note_response` addressed by handle — no hashing.
700    pub fn note_response_at(
701        &mut self,
702        handle: ClientHandle,
703        receive: NtpTimestamp,
704        receive_sent: NtpTimestamp,
705    ) {
706        if let Some(i) = self.resolve(handle) {
707            let record = &mut self.slots[i].record;
708            record.last_receive = Some(receive);
709            record.last_receive_sent = Some(receive_sent);
710        }
711    }
712
713    /// Record the true transmit timestamp of the response just sent. Called
714    /// after `send`, which is the whole point of interleaved mode — this is a
715    /// timestamp the basic exchange cannot report because the packet has not
716    /// left yet when its own transmit field is written.
717    pub fn note_transmit(&mut self, key: &K, transmit: NtpTimestamp) {
718        if let Some(&i) = self.index.get(key) {
719            self.slots[i as usize].record.last_transmit = Some(transmit);
720        }
721    }
722
723    /// `note_transmit` addressed by handle — no hashing.
724    ///
725    /// This is the one called after `send`, so a handle taken before the write
726    /// is used after it. The generation check is what makes that safe: if the
727    /// client was evicted in between, the update is dropped rather than landing
728    /// on whoever inherited the slot.
729    pub fn note_transmit_at(&mut self, handle: ClientHandle, transmit: NtpTimestamp) {
730        if let Some(i) = self.resolve(handle) {
731            self.slots[i].record.last_transmit = Some(transmit);
732        }
733    }
734
735    pub fn note_refused(&mut self) {
736        self.stats.refused += 1;
737    }
738}
739
740#[cfg(test)]
741mod tests {
742    use super::*;
743
744    fn table() -> ClientTable<u32> {
745        ClientTable::new(
746            4,
747            RateLimitConfig {
748                interval_log2: 3, // one per 8 s
749                burst: 2,
750                leak_shift: 2, // KoD every 4th drop
751                // Global ceiling off: these tests are about per-client policy.
752                global_rate_hz: 0.0,
753                global_burst: 0.0,
754            },
755        )
756    }
757
758    #[test]
759    fn address_churn_cannot_defeat_the_limiter() {
760        // Every request arrives from an address never seen before, so the
761        // per-client bucket is always full and the table churns. Without a
762        // global ceiling this answers 100% of the flood — which is what
763        // TIMECORP S12b measured before the ceiling existed.
764        let mut t = ClientTable::<u32>::new(
765            1_024,
766            RateLimitConfig {
767                interval_log2: 3,
768                burst: 8,
769                leak_shift: 4,
770                global_rate_hz: 100.0,
771                global_burst: 100.0,
772            },
773        );
774        let mut answered = 0u64;
775        // 20_000 requests from 20_000 distinct addresses inside one second.
776        for client in 0..20_000u32 {
777            if t.admit(&client, client as f64 * 5e-5) == Disposition::Respond {
778                answered += 1;
779            }
780        }
781        // Ceiling is 100/s with a 100 burst, over ~1 s: ~200 at the very most.
782        assert!(
783            answered <= 250,
784            "address churn produced {answered} answers against a 100/s ceiling"
785        );
786        assert!(answered > 0, "the ceiling must not block everything");
787    }
788
789    #[test]
790    fn the_global_ceiling_refills_over_time() {
791        let mut t = ClientTable::<u32>::new(
792            16,
793            RateLimitConfig {
794                interval_log2: -10, // per-client effectively unlimited
795                burst: 1_000_000,
796                leak_shift: 8,
797                global_rate_hz: 10.0,
798                global_burst: 10.0,
799            },
800        );
801        let mut answered = 0;
802        for i in 0..100u32 {
803            if t.admit(&(i % 4), 0.0) == Disposition::Respond {
804                answered += 1;
805            }
806        }
807        assert!(answered <= 11, "burst exceeded: {answered}");
808        // Ten seconds later the bucket has refilled.
809        assert_eq!(t.admit(&0, 10.0), Disposition::Respond);
810    }
811
812    #[test]
813    fn burst_is_allowed_then_the_limiter_bites() {
814        let mut t = table();
815        assert_eq!(t.admit(&1, 0.0), Disposition::Respond);
816        assert_eq!(t.admit(&1, 0.0), Disposition::Respond);
817        // Burst of 2 spent; further immediate requests are not answered.
818        assert!(matches!(
819            t.admit(&1, 0.0),
820            Disposition::Drop | Disposition::KissOfDeath
821        ));
822        assert_eq!(t.stats.responses, 2);
823    }
824
825    #[test]
826    fn tokens_refill_over_time() {
827        let mut t = table();
828        let _ = t.admit(&1, 0.0);
829        let _ = t.admit(&1, 0.0);
830        assert_ne!(t.admit(&1, 0.0), Disposition::Respond);
831        // 8 s later exactly one token is back.
832        assert_eq!(t.admit(&1, 8.0), Disposition::Respond);
833        assert_ne!(t.admit(&1, 8.0), Disposition::Respond);
834    }
835
836    #[test]
837    fn kiss_of_death_leaks_at_the_configured_rate_not_every_drop() {
838        let mut t = table();
839        let _ = t.admit(&1, 0.0);
840        let _ = t.admit(&1, 0.0);
841        // 12 further requests in the same instant: all over the limit.
842        let mut kods = 0;
843        for _ in 0..12 {
844            if t.admit(&1, 0.0) == Disposition::KissOfDeath {
845                kods += 1;
846            }
847        }
848        // leak_shift 2 => one KoD per 4 drops => 3 of 12.
849        assert_eq!(kods, 3, "KoD leak rate wrong");
850        // The point of leaking: we must answer far less than we are asked, or
851        // the limiter is itself an amplifier.
852        assert!(
853            (kods as u64) < t.stats.dropped_rate_limit,
854            "KoD count must stay below the drop count"
855        );
856    }
857
858    #[test]
859    fn one_client_cannot_starve_another() {
860        let mut t = table();
861        for _ in 0..50 {
862            let _ = t.admit(&1, 0.0);
863        }
864        // A quiet client still gets its full burst.
865        assert_eq!(t.admit(&2, 0.0), Disposition::Respond);
866        assert_eq!(t.admit(&2, 0.0), Disposition::Respond);
867    }
868
869    #[test]
870    fn table_is_bounded_and_evicts_the_stalest() {
871        let mut t = table(); // capacity 4
872        for client in 0..4u32 {
873            let _ = t.admit(&client, client as f64);
874        }
875        assert_eq!(t.len(), 4);
876        // A fifth client evicts client 0, the least recently seen.
877        let _ = t.admit(&99, 10.0);
878        assert_eq!(t.len(), 4, "table exceeded its bound");
879        assert!(t.get(&0).is_none(), "stalest client was not evicted");
880        assert!(t.get(&99).is_some());
881        assert_eq!(t.stats.evicted, 1);
882    }
883
884    #[test]
885    fn interleaved_requires_the_client_to_echo_our_receive_timestamp() {
886        let mut t = table();
887        let _ = t.admit(&1, 0.0);
888        let rx1 = NtpTimestamp::from_unix(1_756_224_000, 0);
889
890        // First exchange: nothing to interleave with yet.
891        assert_eq!(
892            t.response_mode(&1, NtpTimestamp(0x1111)),
893            ResponseMode::Basic
894        );
895        t.note_response(&1, rx1, rx1);
896        let tx1 = NtpTimestamp::from_unix(1_756_224_000, 500);
897        t.note_transmit(&1, tx1);
898
899        // Second exchange, client echoes our receive timestamp: interleaved,
900        // and it gets the true transmit of the response it named.
901        let _ = t.admit(&1, 8.0);
902        match t.response_mode(&1, rx1) {
903            ResponseMode::Interleaved { prev_transmit } => {
904                assert_eq!(prev_transmit, tx1);
905            }
906            other => panic!("expected interleaved, got {other:?}"),
907        }
908        assert_eq!(t.stats.interleaved_responses, 1);
909    }
910
911    #[test]
912    fn interleaved_flag_tracks_use_not_capability() {
913        let mut t = table();
914        let rx1 = NtpTimestamp::from_unix(1_756_224_000, 0);
915        let _ = t.admit(&1, 0.0);
916        t.note_response(&1, rx1, rx1);
917        t.note_transmit(&1, rx1);
918
919        // Capable now, but this client has never asked.
920        let _ = t.admit(&1, 8.0);
921        let _ = t.response_mode(&1, NtpTimestamp::ZERO);
922        assert!(
923            !t.get(&1).expect("record").interleaved_now,
924            "a client that never asked must not be reported as using interleaved"
925        );
926
927        // Now it asks.
928        let _ = t.admit(&1, 16.0);
929        let _ = t.response_mode(&1, rx1);
930        assert!(t.get(&1).expect("record").interleaved_now);
931
932        // And stops asking again.
933        let _ = t.admit(&1, 24.0);
934        let _ = t.response_mode(&1, NtpTimestamp::ZERO);
935        assert!(!t.get(&1).expect("record").interleaved_now);
936    }
937
938    #[test]
939    fn a_client_echoing_the_wrong_value_gets_basic_mode() {
940        let mut t = table();
941        let _ = t.admit(&1, 0.0);
942        let rx1 = NtpTimestamp::from_unix(1_756_224_000, 0);
943        t.note_response(&1, rx1, rx1);
944        t.note_transmit(&1, NtpTimestamp::from_unix(1_756_224_000, 500));
945
946        let _ = t.admit(&1, 8.0);
947        // Some other value — an off-path guess — must not unlock interleaved.
948        assert_eq!(
949            t.response_mode(&1, NtpTimestamp(0xDEAD_BEEF)),
950            ResponseMode::Basic
951        );
952        // Nor may a zero origin.
953        assert_eq!(t.response_mode(&1, NtpTimestamp::ZERO), ResponseMode::Basic);
954    }
955
956    #[test]
957    fn eviction_downgrades_to_basic_rather_than_lying() {
958        let mut t = table(); // capacity 4
959        let rx = NtpTimestamp::from_unix(1_756_224_000, 0);
960        let _ = t.admit(&1, 0.0);
961        t.note_response(&1, rx, rx);
962        t.note_transmit(&1, rx);
963
964        // Push client 1 out.
965        for c in 10..15u32 {
966            let _ = t.admit(&c, 100.0 + c as f64);
967        }
968        assert!(t.get(&1).is_none());
969        // Its next request must be answered in basic mode, not with another
970        // client's timestamps.
971        let _ = t.admit(&1, 200.0);
972        assert_eq!(t.response_mode(&1, rx), ResponseMode::Basic);
973    }
974
975    #[test]
976    fn heavy_client_churn_stays_tractable() {
977        // A public server sees far more addresses than its table holds, so
978        // eviction runs on nearly every packet. A scan-based eviction is
979        // O(capacity) each time and this test does not finish; the indexed one
980        // is O(log n). The assertion is a *count*, not a duration: every
981        // admission must do a bounded amount of index work, which shows up as
982        // the table never exceeding capacity while churning far past it.
983        let capacity = 4_096;
984        let mut t = ClientTable::<u32>::new(capacity, RateLimitConfig::default());
985        let churn = 200_000u32;
986        for client in 0..churn {
987            let _ = t.admit(&client, client as f64 * 0.001);
988        }
989        assert_eq!(t.len(), capacity, "table must sit exactly at capacity");
990        assert_eq!(
991            t.stats.evicted,
992            (churn as u64) - capacity as u64,
993            "every client past capacity must have cost exactly one eviction"
994        );
995        // True LRU: the survivors are the most recent `capacity` clients.
996        assert!(t.get(&(churn - 1)).is_some(), "newest client was evicted");
997        assert!(t.get(&0).is_none(), "oldest client survived");
998    }
999
1000    #[test]
1001    fn eviction_is_true_lru_not_insertion_order() {
1002        let mut t = ClientTable::<u32>::new(3, RateLimitConfig::default());
1003        let _ = t.admit(&1, 0.0);
1004        let _ = t.admit(&2, 1.0);
1005        let _ = t.admit(&3, 2.0);
1006        // Touch client 1 so it is no longer the stalest.
1007        let _ = t.admit(&1, 3.0);
1008        // Inserting a fourth must evict client 2, not client 1.
1009        let _ = t.admit(&4, 4.0);
1010        assert!(t.get(&1).is_some(), "recently used client was evicted");
1011        assert!(t.get(&2).is_none(), "stalest client should have gone");
1012        assert!(t.get(&3).is_some());
1013        assert!(t.get(&4).is_some());
1014    }
1015
1016    #[test]
1017    fn mru_report_is_ordered_and_bounded() {
1018        let mut t = ClientTable::<u32>::new(16, RateLimitConfig::default());
1019        for c in 0..10u32 {
1020            let _ = t.admit(&c, c as f64);
1021        }
1022        let mru = t.most_recent(3);
1023        assert_eq!(mru.len(), 3);
1024        assert_eq!(mru[0].0, 9, "most recent first");
1025        assert_eq!(mru[2].0, 7);
1026    }
1027}