Skip to main content

agent_uri_dht/
simulation.rs

1//! Simulated DHT implementation for evaluation.
2
3use std::collections::{HashMap, HashSet};
4use std::sync::RwLock;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::time::{Duration, Instant, SystemTime};
7
8use agent_uri::AgentUri;
9use agent_uri_attestation::{SigningKey, Verifier};
10use async_trait::async_trait;
11
12use crate::expiration::Expirations;
13use crate::{
14    Cursor, Dht, DhtError, DhtKey, DhtStats, Endpoint, MatchMode, MigrationResult, Mutation,
15    MutationKind, MutationProof, NodeId, Page, PeerAddr, Query, Quorum, ReadOptions, Registration,
16    SimulationConfig, WriteOptions, WriteReceipt,
17};
18
19/// Hands each simulator instance a distinct identity without pulling in a
20/// random-number dependency. Distinctness is what callers rely on; the values
21/// themselves carry no meaning.
22static NEXT_NODE_ID: AtomicU64 = AtomicU64::new(1);
23
24/// The indices that together hold the simulator's state.
25///
26/// They live behind one lock rather than one apiece. Every path that mutates
27/// state touches all of them, so a lock each never let two writers run
28/// concurrently; it only created an order to acquire them in, and therefore an
29/// order to get wrong. One lock has no order to get wrong.
30///
31/// The indices are consistent with each other only between operations, which
32/// is why they are grouped: a registration exists in all of them or in none,
33/// and nothing may observe the state in between.
34struct Indices {
35    /// Primary index: `DhtKey` -> Registrations
36    records: HashMap<DhtKey, Vec<Registration>>,
37
38    /// Secondary index: canonical `AgentUri` -> exact and ancestor keys.
39    keys_by_uri: HashMap<String, Vec<DhtKey>>,
40
41    /// Enforces that a stable trust-root/agent-ID pair has one immutable path.
42    path_by_identity: HashMap<String, String>,
43
44    /// When each registration lapses. Eviction consults this instead of
45    /// walking `records`, so it costs what is due rather than what is stored.
46    expirations: Expirations,
47}
48
49impl Indices {
50    fn new() -> Self {
51        Self {
52            records: HashMap::new(),
53            keys_by_uri: HashMap::new(),
54            path_by_identity: HashMap::new(),
55            expirations: Expirations::default(),
56        }
57    }
58
59    fn clear(&mut self) {
60        self.records.clear();
61        self.keys_by_uri.clear();
62        self.path_by_identity.clear();
63        self.expirations.clear();
64    }
65
66    /// Removes one registration from every index.
67    ///
68    /// Every removal goes through here. The indices only agree with each other
69    /// between operations, so a removal that updated three of the four would
70    /// leave the fourth speaking for a record that no longer exists: a name
71    /// still taken, a path still bound to an identity, or a deadline still
72    /// scheduled against nothing.
73    fn forget(&mut self, uri_str: &str) {
74        self.expirations.forget(uri_str);
75        if let Ok(uri) = AgentUri::parse(uri_str) {
76            self.path_by_identity
77                .remove(&SimulatedDht::identity_key(&uri));
78        }
79        let Some(keys) = self.keys_by_uri.remove(uri_str) else {
80            return;
81        };
82        for key in keys {
83            if let Some(registrations) = self.records.get_mut(&key) {
84                registrations.retain(|r| r.agent_uri().canonical() != *uri_str);
85                if registrations.is_empty() {
86                    self.records.remove(&key);
87                }
88            }
89        }
90    }
91
92    /// Removes every registration due at or before `now`, returning how many.
93    ///
94    /// Proportional to what is due, not to what is stored, which is what lets
95    /// writes carry it: see [`Expirations`] (issue #55).
96    fn evict_due(&mut self, now: SystemTime) -> usize {
97        let lapsed = self.expirations.take_due(now);
98        for uri_str in &lapsed {
99            self.forget(uri_str);
100        }
101        lapsed.len()
102    }
103
104    fn estimate_memory_usage(&self) -> usize {
105        // Rough estimate:
106        // - Each DhtKey: 32 bytes
107        // - Each Registration: ~500 bytes (URI + endpoints + attestation)
108        // - Each queued deadline: 16 bytes plus the URI it names
109        // - HashMap overhead: ~64 bytes per entry
110
111        let key_bytes = self.records.len() * (32 + 64);
112        let registration_bytes = self.keys_by_uri.len() * 500;
113        let uri_index_bytes = self
114            .keys_by_uri
115            .values()
116            .map(|keys| 100 + keys.len() * 32 + 64)
117            .sum::<usize>();
118        // Counted because it is one entry per expiry ever assigned, not one
119        // per record: a simulation that refreshes often carries superseded
120        // deadlines until they pass, and an estimate that ignored them would
121        // read lowest exactly when it mattered most.
122        let expiry_bytes = self.expirations.queued() * (16 + 100 + 64);
123
124        key_bytes + registration_bytes + uri_index_bytes + expiry_bytes
125    }
126}
127
128/// Simulated DHT for evaluation.
129///
130/// Single-process, in-memory implementation that faithfully models
131/// DHT behavior without network overhead. Suitable for evaluation
132/// and testing.
133///
134/// # Thread Safety
135///
136/// One `RwLock` covers all of the state, so lookups run concurrently with each
137/// other and writes are exclusive. There is deliberately not a lock per index:
138/// every write already touches all of them, so splitting them bought no
139/// concurrency and cost an acquisition order to get wrong.
140///
141/// # Examples
142///
143/// ```
144/// use agent_uri::{AgentUri, CapabilityPath, TrustRoot};
145/// use agent_uri_attestation::SigningKey;
146/// use agent_uri_dht::{
147///     Dht, Endpoint, MutationProof, Query, ReadOptions, Registration, SimulatedDht,
148///     SimulationConfig, WriteOptions,
149/// };
150/// use futures::executor::block_on;
151///
152/// // Explicitly unverified because this example isolates indexing behavior.
153/// let dht = SimulatedDht::new(
154///     SimulationConfig::default().with_verify_attestations(false)
155/// );
156///
157/// let uri = AgentUri::parse("agent://anthropic.com/assistant/chat/llm_01h455vb4pex5vsknk084sn02q").unwrap();
158/// let agent_key = SigningKey::generate();
159/// let registration = Registration::new(
160///     uri.clone(),
161///     agent_key.verifying_key(),
162///     vec![Endpoint::https("agent.anthropic.com")],
163/// );
164///
165/// let proof = MutationProof::sign_registration(&agent_key, &registration);
166///
167/// block_on(dht.register(registration, &proof, WriteOptions::default())).unwrap();
168///
169/// let query = Query::exact(
170///     TrustRoot::parse("anthropic.com").unwrap(),
171///     CapabilityPath::parse("assistant/chat").unwrap(),
172/// );
173/// let page = block_on(dht.lookup(&query, &ReadOptions::default())).unwrap();
174///
175/// assert_eq!(page.len(), 1);
176/// ```
177pub struct SimulatedDht {
178    /// Every index, under one lock. See [`Indices`] for why it is one.
179    indices: RwLock<Indices>,
180
181    /// Trusted roots used for registration-time attestation verification.
182    verifier: Verifier,
183
184    /// This instance's identity in the (notional) overlay.
185    node_id: NodeId,
186
187    /// Configuration
188    config: SimulationConfig,
189}
190
191impl SimulatedDht {
192    fn next_node_id() -> NodeId {
193        NodeId::from_bytes(NEXT_NODE_ID.fetch_add(1, Ordering::Relaxed).to_be_bytes())
194    }
195
196    /// The simulator holds one authoritative copy, so every quorum resolves to
197    /// one against it and every accepted write acknowledges exactly that.
198    fn receipt(&self, quorum: Quorum) -> WriteReceipt {
199        let required = quorum.resolve(self.config.replication_factor);
200        WriteReceipt::new(self.config.replication_factor.get(), required)
201    }
202
203    fn identity_key(uri: &AgentUri) -> String {
204        format!("{}/{}", uri.trust_root(), uri.agent_id())
205    }
206
207    fn ancestor_keys(uri: &AgentUri) -> Vec<DhtKey> {
208        (1..=uri.capability_path().depth())
209            .filter_map(|depth| {
210                DhtKey::derive_at_depth(uri.trust_root(), uri.capability_path(), depth)
211            })
212            .collect()
213    }
214
215    /// Creates a new simulated DHT with the given configuration.
216    #[must_use]
217    pub fn new(config: SimulationConfig) -> Self {
218        Self {
219            indices: RwLock::new(Indices::new()),
220            verifier: Verifier::new(),
221            node_id: Self::next_node_id(),
222            config,
223        }
224    }
225
226    /// Creates a simulated DHT with trusted roots for registration verification.
227    #[must_use]
228    pub fn with_verifier(config: SimulationConfig, verifier: Verifier) -> Self {
229        Self {
230            indices: RwLock::new(Indices::new()),
231            verifier,
232            node_id: Self::next_node_id(),
233            config,
234        }
235    }
236
237    /// Creates a new simulated DHT with default configuration.
238    #[must_use]
239    pub fn with_defaults() -> Self {
240        Self::new(SimulationConfig::default())
241    }
242
243    /// Returns the configuration.
244    #[must_use]
245    pub const fn config(&self) -> &SimulationConfig {
246        &self.config
247    }
248
249    /// Registers multiple agents in batch.
250    ///
251    /// More efficient than individual registrations for bulk setup.
252    ///
253    /// # Returns
254    ///
255    /// The number of successfully registered agents.
256    ///
257    /// # Errors
258    ///
259    /// Returns `DhtError` if a critical error occurs (though individual
260    /// registration failures are silently counted).
261    pub fn register_batch(
262        &self,
263        registrations: Vec<(Registration, MutationProof)>,
264    ) -> Result<usize, DhtError> {
265        let mut count = 0;
266        for (registration, proof) in registrations {
267            if self.do_register(&registration, &proof).is_ok() {
268                count += 1;
269            }
270        }
271        Ok(count)
272    }
273
274    /// Returns statistics about the DHT state.
275    ///
276    /// # Panics
277    ///
278    /// Panics if the internal lock is poisoned.
279    #[must_use]
280    pub fn stats(&self) -> DhtStats {
281        let indices = self.indices.read().expect("lock poisoned");
282
283        let total_registrations = indices.keys_by_uri.len();
284        let unique_keys = indices.records.len();
285        // Derived on demand rather than kept as a standing index. `stats` is a
286        // diagnostic that already walks both indices, so the parse cost lands
287        // here instead of on every write.
288        let unique_trust_roots = indices
289            .keys_by_uri
290            .keys()
291            .filter_map(|uri| AgentUri::parse(uri).ok())
292            .map(|uri| uri.trust_root().as_str().to_string())
293            .collect::<HashSet<_>>()
294            .len();
295
296        let max_registrations_per_key = indices.records.values().map(Vec::len).max().unwrap_or(0);
297
298        // Use f64::from for u32 to avoid precision loss; saturate for stats safety
299        let indexed_copies = indices.records.values().map(Vec::len).sum::<usize>();
300        let total_u32 = u32::try_from(indexed_copies).unwrap_or(u32::MAX);
301        let keys_u32 = u32::try_from(unique_keys).unwrap_or(u32::MAX);
302        let avg_registrations_per_key = if keys_u32 > 0 {
303            f64::from(total_u32) / f64::from(keys_u32)
304        } else {
305            0.0
306        };
307
308        // Estimate memory usage
309        let memory_bytes = indices.estimate_memory_usage();
310
311        DhtStats {
312            total_registrations,
313            unique_keys,
314            unique_trust_roots,
315            max_registrations_per_key,
316            avg_registrations_per_key,
317            path_depth_histogram: Vec::new(),
318            memory_bytes,
319        }
320    }
321
322    /// Simulates agent migration with timing.
323    ///
324    /// Takes the agent's signing key because it plays the part of the agent:
325    /// it reads the record's current sequence number and mints the proof the
326    /// migration needs, exactly as a migrating agent would.
327    ///
328    /// # Errors
329    ///
330    /// Returns `DhtError::NotFound` if the agent is not registered, or
331    /// `DhtError::Unauthorized` if `agent_key` is not the key the record names.
332    ///
333    /// # Panics
334    ///
335    /// Panics if the internal lock is poisoned.
336    pub fn simulate_migration(
337        &self,
338        agent_uri: &AgentUri,
339        new_endpoint: Endpoint,
340        agent_key: &SigningKey,
341    ) -> Result<MigrationResult, DhtError> {
342        let uri_str = agent_uri.canonical();
343
344        // Read the record the agent is about to move: its endpoints for the
345        // before/after report, its version to sign against. The guard is
346        // dropped before `apply_mutation` runs, because that takes the write
347        // lock and `RwLock` is not reentrant.
348        let current = {
349            let indices = self.indices.read().expect("lock poisoned");
350            let keys = indices
351                .keys_by_uri
352                .get(&uri_str)
353                .ok_or_else(|| DhtError::not_found(&uri_str))?;
354            let key = keys.last().ok_or_else(|| DhtError::not_found(&uri_str))?;
355
356            let registrations = indices
357                .records
358                .get(key)
359                .ok_or_else(|| DhtError::not_found(&uri_str))?;
360
361            registrations
362                .iter()
363                .find(|r| r.agent_uri().canonical() == uri_str)
364                .cloned()
365                .ok_or_else(|| DhtError::not_found(&uri_str))?
366        };
367        let old_endpoints = current.endpoints().to_vec();
368
369        let endpoints = std::slice::from_ref(&new_endpoint);
370        // Migrating leaves the expiry where it is, so the proof names the one
371        // the record already carries.
372        let proof = MutationProof::sign_next(
373            agent_key,
374            &current,
375            &Mutation::UpdateEndpoint {
376                endpoints,
377                expires_at: current.expires_at(),
378            },
379        );
380
381        // Time the update
382        let start = Instant::now();
383        self.apply_mutation(agent_uri, &proof, |record| Mutation::UpdateEndpoint {
384            endpoints,
385            expires_at: record.expires_at(),
386        })?;
387        let update_latency = start.elapsed();
388
389        Ok(MigrationResult::success(
390            uri_str,
391            old_endpoints,
392            vec![new_endpoint],
393            update_latency,
394        ))
395    }
396
397    /// Clears all registrations.
398    ///
399    /// # Panics
400    ///
401    /// Panics if the internal lock is poisoned.
402    pub fn clear(&self) {
403        self.indices.write().expect("lock poisoned").clear();
404    }
405
406    /// Removes expired registrations.
407    ///
408    /// Returns the number of registrations removed.
409    ///
410    /// A simulation under load does not need this call: writes evict what has
411    /// lapsed as they go, when [`SimulationConfig::auto_expire`] is set. It is
412    /// here for a store that has gone quiet with garbage still in it, and for
413    /// one that has turned automatic expiry off. It removes what is due either
414    /// way, because the configuration decides whether expiry happens unasked,
415    /// not whether it happens when asked.
416    ///
417    /// The cost is the number of records actually due rather than the number
418    /// stored: deadlines are kept in order, so a store where nothing has
419    /// lapsed costs one comparison (issue #55).
420    ///
421    /// # Panics
422    ///
423    /// Panics if the internal lock is poisoned.
424    pub fn expire_stale(&self) -> usize {
425        self.indices
426            .write()
427            .expect("lock poisoned")
428            .evict_due(SystemTime::now())
429    }
430}
431
432impl SimulatedDht {
433    /// Applies a registration to the indices. Shared by the async trait method
434    /// and by the synchronous evaluation helpers.
435    fn do_register(
436        &self,
437        registration: &Registration,
438        proof: &MutationProof,
439    ) -> Result<(), DhtError> {
440        if registration.endpoints().is_empty() {
441            return Err(DhtError::NoEndpoints);
442        }
443
444        let uri_str = registration.agent_uri().canonical();
445        let keys = Self::ancestor_keys(registration.agent_uri());
446        let exact_key = DhtKey::derive(
447            registration.agent_uri().trust_root(),
448            registration.agent_uri().capability_path(),
449        );
450        let identity = Self::identity_key(registration.agent_uri());
451
452        // Possession first. The attestation says the trust root vouches for a
453        // key; the proof says the party registering holds it. Checking the
454        // token alone would accept a token lifted from any public record.
455        proof.authorizes_registration(registration)?;
456
457        if self.config.verify_attestations {
458            let token = registration.attestation().ok_or_else(|| {
459                DhtError::invalid_attestation(&uri_str, "registration has no attestation token")
460            })?;
461            let claims = self
462                .verifier
463                .verify_for_capability(
464                    token,
465                    registration.agent_uri(),
466                    registration.agent_uri().capability_path(),
467                )
468                .map_err(|error| DhtError::invalid_attestation(&uri_str, error.to_string()))?;
469
470            // The token has to be about this record's key, or it is somebody
471            // else's token: an attestation is world-readable and says nothing
472            // about who presented it.
473            let attested = claims
474                .agent_verifying_key()
475                .map_err(|error| DhtError::invalid_attestation(&uri_str, error.to_string()))?;
476            if attested != *registration.agent_key() {
477                return Err(DhtError::AgentKeyMismatch { agent_uri: uri_str });
478            }
479        }
480
481        // Check and insert under one acquisition of the write lock. The
482        // identity binding, capacity checks, and ancestor writes are one
483        // atomic update.
484        {
485            let mut indices = self.indices.write().expect("lock poisoned");
486
487            // Before the checks, not after: a lapsed record is garbage, and
488            // garbage does not hold a name. Left in place it would keep the
489            // URI taken, the identity bound to its path, and a slot at its key
490            // occupied until somebody remembered to call `expire_stale`, which
491            // is the accumulation issue #55 reports.
492            if self.config.auto_expire {
493                indices.evict_due(SystemTime::now());
494            }
495
496            if indices.keys_by_uri.contains_key(&uri_str) {
497                return Err(DhtError::already_registered(&uri_str));
498            }
499            if let Some(existing_path) = indices.path_by_identity.get(&identity)
500                && existing_path != registration.agent_uri().capability_path().as_str()
501            {
502                return Err(DhtError::IdentityCapabilityConflict {
503                    identity,
504                    existing_path: existing_path.clone(),
505                    requested_path: registration
506                        .agent_uri()
507                        .capability_path()
508                        .as_str()
509                        .to_string(),
510                });
511            }
512            // Capacity is charged to the key this agent registers at, and to
513            // no other. The ancestor keys hold a copy too, but the agent did
514            // not choose them: it chose a path, and the ancestors follow. An
515            // ancestor bucket therefore holds its whole subtree, and charging
516            // a registration against it would let whoever first fills a
517            // popular top-level prefix lock every path beneath it, at the cost
518            // of registering under that prefix (issue #54).
519            if let Some(registrations) = indices.records.get(&exact_key)
520                && registrations.len() >= self.config.max_registrations_per_capability
521            {
522                return Err(DhtError::key_capacity_exceeded(
523                    format!("{exact_key}"),
524                    self.config.max_registrations_per_capability,
525                ));
526            }
527
528            // Primary index: write once at every ancestor key. A GET for any
529            // ancestor therefore returns that path and all descendants without
530            // attempting to enumerate SHA-256 children.
531            for key in &keys {
532                indices
533                    .records
534                    .entry(*key)
535                    .or_default()
536                    .push(registration.clone());
537            }
538
539            indices.expirations.set(&uri_str, registration.expires_at());
540            indices.keys_by_uri.insert(uri_str, keys);
541            indices.path_by_identity.insert(
542                identity,
543                registration
544                    .agent_uri()
545                    .capability_path()
546                    .as_str()
547                    .to_string(),
548            );
549        }
550
551        Ok(())
552    }
553
554    /// Rewrites every ancestor-key copy of a registration in place.
555    ///
556    /// `mutate` runs once per copy and must be idempotent across copies: the
557    /// copies are the same logical record, so they must end up identical.
558    fn rewrite_copies(
559        records: &mut HashMap<DhtKey, Vec<Registration>>,
560        keys: &[DhtKey],
561        uri_str: &str,
562        sequence: u64,
563        mut mutate: impl FnMut(&mut Registration),
564    ) -> Result<(), DhtError> {
565        for key in keys {
566            let registrations = records
567                .get_mut(key)
568                .ok_or_else(|| DhtError::not_found(uri_str))?;
569            let registration = registrations
570                .iter_mut()
571                .find(|r| r.agent_uri().canonical() == uri_str)
572                .ok_or_else(|| DhtError::not_found(uri_str))?;
573            mutate(registration);
574            registration.set_sequence(sequence);
575        }
576        Ok(())
577    }
578
579    /// Authorizes a write against the record it targets, then applies it.
580    ///
581    /// This is the only path by which a stored record changes, so it is the
582    /// only place authorization has to be enforced. Checking and applying
583    /// happen under one acquisition of the write lock: split apart, two
584    /// writers could clear the same sequence check and both proceed, and the
585    /// second would overwrite the first while its proof said otherwise.
586    ///
587    /// `describe` builds the mutation from the record as the store holds it,
588    /// rather than taking one ready-made. A migration's signed `expires_at` is
589    /// the record's own, and a caller that supplied it separately could
590    /// describe the same write differently from the record it targets. Reading
591    /// it off the record under the same lock that applies the write means it
592    /// cannot.
593    fn apply_mutation<'m>(
594        &self,
595        agent_uri: &AgentUri,
596        proof: &MutationProof,
597        describe: impl FnOnce(&Registration) -> Mutation<'m>,
598    ) -> Result<(), DhtError> {
599        let uri_str = agent_uri.canonical();
600
601        let mut guard = self.indices.write().expect("lock poisoned");
602        let indices = &mut *guard;
603
604        let keys = indices
605            .keys_by_uri
606            .get(&uri_str)
607            .cloned()
608            .ok_or_else(|| DhtError::not_found(&uri_str))?;
609
610        // Every ancestor key holds the same logical record, so the first copy
611        // found answers for all of them.
612        let current = keys
613            .iter()
614            .filter_map(|key| indices.records.get(key))
615            .flatten()
616            .find(|r| r.agent_uri().canonical() == uri_str)
617            .cloned()
618            .ok_or_else(|| DhtError::not_found(&uri_str))?;
619
620        if current.is_expired() && self.config.auto_expire {
621            return Err(DhtError::expired(&uri_str));
622        }
623
624        let mutation = describe(&current);
625        proof.authorizes(&current, &mutation)?;
626
627        match mutation {
628            // The expiry is applied as well as the endpoints, so that the
629            // stored record is exactly the record the proof describes and not
630            // merely one the proof permitted.
631            Mutation::UpdateEndpoint {
632                endpoints,
633                expires_at,
634            } => {
635                Self::rewrite_copies(
636                    &mut indices.records,
637                    &keys,
638                    &uri_str,
639                    proof.sequence(),
640                    |registration| {
641                        registration.update_endpoints(endpoints.to_vec());
642                        registration.set_expires_at(expires_at);
643                    },
644                )?;
645                indices.expirations.set(&uri_str, expires_at);
646            }
647            // The signed instant, not `now + ttl`. A store that recomputed it
648            // would write an expiry the agent never signed, and on a wire
649            // protocol that field would then be open to anyone who relays the
650            // record (issue #81).
651            Mutation::Refresh { expires_at, .. } => {
652                Self::rewrite_copies(
653                    &mut indices.records,
654                    &keys,
655                    &uri_str,
656                    proof.sequence(),
657                    |registration| registration.set_expires_at(expires_at),
658                )?;
659                indices.expirations.set(&uri_str, expires_at);
660            }
661            Mutation::Deregister => indices.forget(&uri_str),
662            // Registration creates a record rather than changing one, so it
663            // never reaches here; `do_register` authorizes it directly.
664            Mutation::Register { .. } => {
665                return Err(DhtError::Unauthorized {
666                    agent_uri: uri_str,
667                    operation: MutationKind::Register,
668                });
669            }
670        }
671
672        // After the write rather than before it. A write against a record that
673        // has lapsed is reported as expired above, and sweeping first would
674        // have taken that record out and reported it missing instead. Sweeping
675        // here is what keeps eviction paced by writes rather than by however
676        // long it has been since anyone called `expire_stale` (issue #55).
677        if self.config.auto_expire {
678            indices.evict_due(SystemTime::now());
679        }
680
681        Ok(())
682    }
683
684    fn do_lookup(
685        &self,
686        query: &Query,
687        options: &ReadOptions,
688    ) -> Result<Page<Registration>, DhtError> {
689        let key = DhtKey::derive(query.trust_root(), query.capability_path());
690
691        let mut matches: Vec<Registration> = {
692            let indices = self.indices.read().expect("lock poisoned");
693            indices
694                .records
695                .get(&key)
696                .map(|registrations| {
697                    registrations
698                        .iter()
699                        .filter(|r| match query.match_mode() {
700                            // Ancestor-key materialization means the key already
701                            // holds the whole subtree, so a prefix query keeps
702                            // everything and an exact query filters back down.
703                            MatchMode::Prefix => true,
704                            MatchMode::Exact => {
705                                r.agent_uri().capability_path() == query.capability_path()
706                            }
707                        })
708                        .filter(|r| !r.is_expired() || !self.config.auto_expire)
709                        .cloned()
710                        .collect()
711                })
712                .unwrap_or_default()
713        };
714
715        // Paging is only coherent over a stable order, and insertion order is
716        // not stable across deregistrations. Canonical URI is total and cheap.
717        matches.sort_by_key(|r| r.agent_uri().canonical());
718
719        let offset = options.cursor.map_or(0, |c| c.position() as usize);
720        if offset >= matches.len() {
721            return Ok(Page::complete(Vec::new()));
722        }
723
724        let page_size = self.config.page_size.get();
725        let end = offset.saturating_add(page_size).min(matches.len());
726        let items = matches[offset..end].to_vec();
727
728        if end < matches.len() {
729            let next = u32::try_from(end).map_err(|_| DhtError::QuorumFailed {
730                achieved: 0,
731                required: 1,
732            })?;
733            Ok(Page::partial(items, Cursor::at(next)))
734        } else {
735            Ok(Page::complete(items))
736        }
737    }
738}
739
740#[async_trait]
741impl Dht for SimulatedDht {
742    async fn register(
743        &self,
744        registration: Registration,
745        proof: &MutationProof,
746        options: WriteOptions,
747    ) -> Result<WriteReceipt, DhtError> {
748        self.do_register(&registration, proof)?;
749        Ok(self.receipt(options.quorum))
750    }
751
752    async fn update_endpoint(
753        &self,
754        agent_uri: &AgentUri,
755        new_endpoints: Vec<Endpoint>,
756        proof: &MutationProof,
757        options: WriteOptions,
758    ) -> Result<WriteReceipt, DhtError> {
759        // Checked before authorization: an empty list is a malformed request,
760        // not a rejected one, and saying so costs no signature verification.
761        if new_endpoints.is_empty() {
762            return Err(DhtError::NoEndpoints);
763        }
764        self.apply_mutation(agent_uri, proof, |record| Mutation::UpdateEndpoint {
765            endpoints: &new_endpoints,
766            expires_at: record.expires_at(),
767        })?;
768        Ok(self.receipt(options.quorum))
769    }
770
771    async fn refresh(
772        &self,
773        agent_uri: &AgentUri,
774        ttl: Duration,
775        expires_at: SystemTime,
776        proof: &MutationProof,
777        options: WriteOptions,
778    ) -> Result<WriteReceipt, DhtError> {
779        self.apply_mutation(agent_uri, proof, |_| Mutation::Refresh { ttl, expires_at })?;
780        Ok(self.receipt(options.quorum))
781    }
782
783    async fn deregister(
784        &self,
785        agent_uri: &AgentUri,
786        proof: &MutationProof,
787        _options: WriteOptions,
788    ) -> Result<(), DhtError> {
789        self.apply_mutation(agent_uri, proof, |_| Mutation::Deregister)
790    }
791
792    async fn lookup(
793        &self,
794        query: &Query,
795        options: &ReadOptions,
796    ) -> Result<Page<Registration>, DhtError> {
797        self.do_lookup(query, options)
798    }
799
800    fn local_id(&self) -> NodeId {
801        self.node_id.clone()
802    }
803
804    async fn bootstrap(&self, peers: &[PeerAddr]) -> Result<(), DhtError> {
805        // There is no overlay to join. Rejecting an empty peer list anyway
806        // keeps the failure shape identical to a real backend, so a caller that
807        // forgets to configure peers finds out here rather than in production.
808        if peers.is_empty() {
809            return Err(DhtError::NoPeers);
810        }
811        Ok(())
812    }
813}
814
815#[cfg(test)]
816mod tests {
817    use super::*;
818    use crate::RecordVersion;
819    // Comparing at millisecond resolution, which is what the signed payload and
820    // the serde representation both round to.
821    use crate::registration::system_time_to_millis;
822    use agent_uri::{CapabilityPath, TrustRoot};
823    use agent_uri_attestation::{AcceptAll, Issuer};
824    use futures::executor::block_on;
825    use std::num::NonZeroUsize;
826    use std::time::{Duration, SystemTime};
827
828    fn unverified_dht() -> SimulatedDht {
829        SimulatedDht::new(SimulationConfig::default().with_verify_attestations(false))
830    }
831
832    fn test_uri(suffix: &str) -> AgentUri {
833        AgentUri::parse(&format!(
834            "agent://anthropic.com/assistant/chat/llm_01h455vb4pex5vsknk084sn0{suffix}"
835        ))
836        .unwrap()
837    }
838
839    /// A URI at an arbitrary capability path, for tests about how capacity is
840    /// charged across a subtree rather than at one path.
841    fn uri_at(path: &str, suffix: &str) -> AgentUri {
842        AgentUri::parse(&format!(
843            "agent://anthropic.com/{path}/llm_01h455vb4pex5vsknk084sn0{suffix}"
844        ))
845        .unwrap()
846    }
847
848    /// A simulator that is full the moment anything registers, so that any
849    /// key a registration is charged against rejects the next one.
850    fn dht_with_capacity_for_one() -> SimulatedDht {
851        SimulatedDht::new(
852            SimulationConfig::new()
853                .with_max_registrations_per_capability(1)
854                .with_verify_attestations(false),
855        )
856    }
857
858    fn test_endpoint() -> Endpoint {
859        Endpoint::https("agent.anthropic.com:443")
860    }
861
862    fn registration(uri: &AgentUri, key: &SigningKey) -> Registration {
863        Registration::new(uri.clone(), key.verifying_key(), vec![test_endpoint()])
864    }
865
866    // The trait is async so that a distributed backend can implement it. The
867    // simulator never actually yields, so these drive it to completion without
868    // a runtime and keep the tests reading like the operations they describe.
869    /// Registers with a proof from `key`. Tests that pass a key other than the
870    /// record's own are checking that registration is a write like any other.
871    fn register(
872        dht: &SimulatedDht,
873        registration: Registration,
874        key: &SigningKey,
875    ) -> Result<WriteReceipt, DhtError> {
876        let proof = MutationProof::sign_registration(key, &registration);
877        block_on(dht.register(registration, &proof, WriteOptions::default()))
878    }
879
880    /// Returns the record the DHT currently holds, which is what a caller
881    /// preparing a write would have looked up.
882    fn record(dht: &SimulatedDht, uri: &AgentUri) -> Registration {
883        stored(dht, uri).expect("record is registered")
884    }
885
886    fn stored(dht: &SimulatedDht, uri: &AgentUri) -> Option<Registration> {
887        lookup(
888            dht,
889            &Query::exact(uri.trust_root().clone(), uri.capability_path().clone()),
890        )
891        .into_iter()
892        .find(|r| r.agent_uri() == uri)
893    }
894
895    /// The version a caller would sign against, or an invented one when there
896    /// is nothing to look up. Tests that write to an unregistered URI need a
897    /// proof to hand over, and a real caller working from a stale cache would
898    /// be in exactly that position.
899    fn version_at(dht: &SimulatedDht, uri: &AgentUri, sequence: u64) -> RecordVersion {
900        RecordVersion::new(
901            stored(dht, uri).map_or_else(SystemTime::now, |r| r.registered_at()),
902            sequence,
903        )
904    }
905
906    fn deregister(
907        dht: &SimulatedDht,
908        uri: &AgentUri,
909        key: &SigningKey,
910        sequence: u64,
911    ) -> Result<(), DhtError> {
912        let proof = MutationProof::sign(
913            key,
914            uri,
915            version_at(dht, uri, sequence),
916            &Mutation::Deregister,
917        );
918        block_on(dht.deregister(uri, &proof, WriteOptions::default()))
919    }
920
921    /// The expiry a migration proof has to name: the one the record already
922    /// carries, because a migration does not move it.
923    fn expiry_of(dht: &SimulatedDht, uri: &AgentUri) -> SystemTime {
924        stored(dht, uri).map_or_else(SystemTime::now, |r| r.expires_at())
925    }
926
927    fn update_endpoint(
928        dht: &SimulatedDht,
929        uri: &AgentUri,
930        key: &SigningKey,
931        sequence: u64,
932        endpoints: Vec<Endpoint>,
933    ) -> Result<WriteReceipt, DhtError> {
934        let proof = MutationProof::sign(
935            key,
936            uri,
937            version_at(dht, uri, sequence),
938            &Mutation::UpdateEndpoint {
939                endpoints: &endpoints,
940                expires_at: expiry_of(dht, uri),
941            },
942        );
943        block_on(dht.update_endpoint(uri, endpoints, &proof, WriteOptions::default()))
944    }
945
946    fn refresh(
947        dht: &SimulatedDht,
948        uri: &AgentUri,
949        key: &SigningKey,
950        sequence: u64,
951        ttl: Duration,
952    ) -> Result<WriteReceipt, DhtError> {
953        // Computed once and used for both the proof and the call, which is
954        // what every caller has to do: the store applies the instant the agent
955        // signed rather than deriving one of its own.
956        let expires_at = SystemTime::now() + ttl;
957        let proof = MutationProof::sign(
958            key,
959            uri,
960            version_at(dht, uri, sequence),
961            &Mutation::Refresh { ttl, expires_at },
962        );
963        block_on(dht.refresh(uri, ttl, expires_at, &proof, WriteOptions::default()))
964    }
965
966    fn lookup(dht: &SimulatedDht, query: &Query) -> Vec<Registration> {
967        block_on(dht.lookup(query, &ReadOptions::default()))
968            .unwrap()
969            .into_items()
970    }
971
972    fn exact(trust_root: &str, path: &str) -> Query {
973        Query::exact(
974            TrustRoot::parse(trust_root).unwrap(),
975            CapabilityPath::parse(path).unwrap(),
976        )
977    }
978
979    fn prefix(trust_root: &str, path: &str) -> Query {
980        Query::prefix(
981            TrustRoot::parse(trust_root).unwrap(),
982            CapabilityPath::parse(path).unwrap(),
983        )
984    }
985
986    #[test]
987    fn register_and_lookup_exact() {
988        let dht = unverified_dht();
989        let key = SigningKey::generate();
990        let uri = test_uri("2q");
991
992        register(&dht, registration(&uri, &key), &key).unwrap();
993
994        let results = lookup(&dht, &exact("anthropic.com", "assistant/chat"));
995
996        assert_eq!(results.len(), 1);
997        assert_eq!(results[0].agent_uri(), &uri);
998    }
999
1000    #[test]
1001    fn register_empty_endpoints_fails() {
1002        let dht = unverified_dht();
1003        let key = SigningKey::generate();
1004        let registration = Registration::new(test_uri("2q"), key.verifying_key(), vec![]);
1005
1006        let result = register(&dht, registration, &key);
1007
1008        assert!(matches!(result, Err(DhtError::NoEndpoints)));
1009    }
1010
1011    #[test]
1012    fn double_registration_fails() {
1013        let dht = unverified_dht();
1014        let key = SigningKey::generate();
1015        let uri = test_uri("2q");
1016
1017        register(&dht, registration(&uri, &key), &key).unwrap();
1018
1019        let result = register(&dht, registration(&uri, &key), &key);
1020
1021        assert!(matches!(result, Err(DhtError::AlreadyRegistered { .. })));
1022    }
1023
1024    #[test]
1025    fn deregister_removes_agent() {
1026        let dht = unverified_dht();
1027        let key = SigningKey::generate();
1028        let uri = test_uri("2q");
1029
1030        register(&dht, registration(&uri, &key), &key).unwrap();
1031        deregister(&dht, &uri, &key, 1).unwrap();
1032
1033        let results = lookup(&dht, &exact("anthropic.com", "assistant/chat"));
1034
1035        assert!(results.is_empty());
1036    }
1037
1038    #[test]
1039    fn lookup_prefix_finds_descendants() {
1040        let dht = unverified_dht();
1041        let key = SigningKey::generate();
1042
1043        let uri1 =
1044            AgentUri::parse("agent://anthropic.com/assistant/chat/llm_01h455vb4pex5vsknk084sn02q")
1045                .unwrap();
1046        let uri2 =
1047            AgentUri::parse("agent://anthropic.com/assistant/code/llm_01h455vb4pex5vsknk084sn02r")
1048                .unwrap();
1049
1050        register(&dht, registration(&uri1, &key), &key).unwrap();
1051        register(&dht, registration(&uri2, &key), &key).unwrap();
1052
1053        let results = lookup(&dht, &prefix("anthropic.com", "assistant"));
1054
1055        assert_eq!(results.len(), 2);
1056    }
1057
1058    #[test]
1059    fn lookup_is_scoped_to_one_trust_root() {
1060        let dht = unverified_dht();
1061        let key = SigningKey::generate();
1062
1063        let uri1 =
1064            AgentUri::parse("agent://anthropic.com/assistant/chat/llm_01h455vb4pex5vsknk084sn02q")
1065                .unwrap();
1066        let uri2 =
1067            AgentUri::parse("agent://openai.com/assistant/chat/llm_01h455vb4pex5vsknk084sn02r")
1068                .unwrap();
1069
1070        register(&dht, registration(&uri1, &key), &key).unwrap();
1071        register(&dht, registration(&uri2, &key), &key).unwrap();
1072
1073        // Identical capability paths under two authorities derive two distinct
1074        // keys, so neither lookup can observe the other's registration. This is
1075        // the cross-trust-root isolation that removing the global scan defends.
1076        let anthropic = lookup(&dht, &exact("anthropic.com", "assistant/chat"));
1077        let openai = lookup(&dht, &exact("openai.com", "assistant/chat"));
1078
1079        assert_eq!(anthropic.len(), 1);
1080        assert_eq!(
1081            anthropic[0].agent_uri().trust_root().as_str(),
1082            "anthropic.com"
1083        );
1084        assert_eq!(openai.len(), 1);
1085        assert_eq!(openai[0].agent_uri().trust_root().as_str(), "openai.com");
1086    }
1087
1088    #[test]
1089    fn update_endpoint_changes_endpoints() {
1090        let dht = unverified_dht();
1091        let key = SigningKey::generate();
1092        let uri = test_uri("2q");
1093
1094        register(&dht, registration(&uri, &key), &key).unwrap();
1095
1096        let new_endpoint = Endpoint::grpc("agent.anthropic.com:50051");
1097        update_endpoint(&dht, &uri, &key, 1, vec![new_endpoint.clone()]).unwrap();
1098
1099        let results = lookup(&dht, &exact("anthropic.com", "assistant/chat"));
1100
1101        assert_eq!(results.len(), 1);
1102        assert_eq!(results[0].endpoints(), &[new_endpoint]);
1103    }
1104
1105    #[test]
1106    fn stats_reports_correct_counts() {
1107        let dht = unverified_dht();
1108        let key = SigningKey::generate();
1109
1110        let uri1 = test_uri("2q");
1111        let uri2 =
1112            AgentUri::parse("agent://anthropic.com/assistant/code/llm_01h455vb4pex5vsknk084sn02r")
1113                .unwrap();
1114
1115        register(&dht, registration(&uri1, &key), &key).unwrap();
1116        register(&dht, registration(&uri2, &key), &key).unwrap();
1117
1118        let stats = dht.stats();
1119        assert_eq!(stats.total_registrations(), 2);
1120        assert_eq!(stats.unique_trust_roots(), 1);
1121    }
1122
1123    /// Capacity does bind somewhere: two agents at the same capability path
1124    /// are charged against the same key, and the second is refused.
1125    #[test]
1126    fn a_second_agent_at_a_full_capability_path_is_refused() {
1127        let dht = dht_with_capacity_for_one();
1128        let key = SigningKey::generate();
1129
1130        register(&dht, registration(&test_uri("2q"), &key), &key).unwrap();
1131
1132        let result = register(&dht, registration(&test_uri("2r"), &key), &key);
1133
1134        assert!(matches!(result, Err(DhtError::KeyCapacityExceeded { .. })));
1135    }
1136
1137    /// A sibling path shares an ancestor key and nothing else. Charging it for
1138    /// that ancestor would mean the first agent to register under a trust root
1139    /// decides how many others may, which is the subtree lockout of issue #54.
1140    #[test]
1141    fn a_full_ancestor_does_not_refuse_a_sibling_path() {
1142        let dht = dht_with_capacity_for_one();
1143        let key = SigningKey::generate();
1144
1145        let occupant = uri_at("assistant/chat", "2q");
1146        let sibling = uri_at("assistant/code", "2r");
1147        register(&dht, registration(&occupant, &key), &key).unwrap();
1148
1149        register(&dht, registration(&sibling, &key), &key)
1150            .expect("a path of its own has capacity of its own");
1151    }
1152
1153    /// The issue's other half: a deeper path whose ancestor is the crowded key
1154    /// itself. Its own key is empty, and that is the key it is charged for.
1155    #[test]
1156    fn a_full_ancestor_does_not_refuse_a_path_beneath_it() {
1157        let dht = dht_with_capacity_for_one();
1158        let key = SigningKey::generate();
1159
1160        register(
1161            &dht,
1162            registration(&uri_at("assistant/chat", "2q"), &key),
1163            &key,
1164        )
1165        .unwrap();
1166
1167        register(
1168            &dht,
1169            registration(&uri_at("assistant/chat/vision", "2r"), &key),
1170            &key,
1171        )
1172        .expect("an empty key is not full because its ancestor is");
1173    }
1174
1175    /// The consequence, stated outright: an ancestor bucket holds its whole
1176    /// subtree and is not bounded by the configured capacity. That is the
1177    /// direct model's premise rather than an oversight, and a store that
1178    /// cannot honor it is required to shard instead.
1179    #[test]
1180    fn an_ancestor_key_holds_more_than_one_path_may() {
1181        let dht = dht_with_capacity_for_one();
1182        let key = SigningKey::generate();
1183
1184        register(
1185            &dht,
1186            registration(&uri_at("assistant/chat", "2q"), &key),
1187            &key,
1188        )
1189        .unwrap();
1190        register(
1191            &dht,
1192            registration(&uri_at("assistant/code", "2r"), &key),
1193            &key,
1194        )
1195        .unwrap();
1196
1197        let under_assistant = lookup(&dht, &prefix("anthropic.com", "assistant"));
1198
1199        assert_eq!(under_assistant.len(), 2);
1200        assert!(dht.stats().max_registrations_per_key() > 1);
1201    }
1202
1203    #[test]
1204    fn clear_removes_all() {
1205        let dht = unverified_dht();
1206        let key = SigningKey::generate();
1207
1208        register(&dht, registration(&test_uri("2q"), &key), &key).unwrap();
1209
1210        dht.clear();
1211
1212        let stats = dht.stats();
1213        assert_eq!(stats.total_registrations(), 0);
1214    }
1215
1216    #[test]
1217    fn register_batch_counts_successes() {
1218        let dht = unverified_dht();
1219        let key = SigningKey::generate();
1220
1221        let registrations = vec![
1222            registration(&test_uri("2q"), &key),
1223            registration(&test_uri("2r"), &key),
1224            // Will fail - no endpoints
1225            Registration::new(test_uri("2s"), key.verifying_key(), vec![]),
1226        ]
1227        .into_iter()
1228        .map(|r| {
1229            let proof = MutationProof::sign_registration(&key, &r);
1230            (r, proof)
1231        })
1232        .collect();
1233
1234        let count = dht.register_batch(registrations).unwrap();
1235
1236        assert_eq!(count, 2);
1237    }
1238
1239    #[test]
1240    fn secure_default_rejects_missing_attestation() {
1241        let dht = SimulatedDht::with_defaults();
1242        let key = SigningKey::generate();
1243
1244        let result = register(&dht, registration(&test_uri("2q"), &key), &key);
1245
1246        assert!(matches!(result, Err(DhtError::InvalidAttestation { .. })));
1247    }
1248
1249    #[test]
1250    fn trusted_attestation_is_verified_before_registration() {
1251        let signing_key = SigningKey::generate();
1252        let issuer = Issuer::new(
1253            "anthropic.com",
1254            signing_key.clone(),
1255            Duration::from_hours(1),
1256        );
1257        let mut verifier = Verifier::new().with_revocation(AcceptAll);
1258        verifier.add_trusted_root("anthropic.com", signing_key.verifying_key());
1259        let dht = SimulatedDht::with_verifier(SimulationConfig::default(), verifier);
1260        let agent_key = SigningKey::generate();
1261        let uri = test_uri("2q");
1262        let token = issuer
1263            .issue(
1264                &uri,
1265                &agent_key.verifying_key(),
1266                vec!["assistant/chat".to_string()],
1267            )
1268            .unwrap();
1269
1270        register(
1271            &dht,
1272            registration(&uri, &agent_key).with_attestation(token),
1273            &agent_key,
1274        )
1275        .unwrap();
1276
1277        let results = lookup(
1278            &dht,
1279            &Query::exact(uri.trust_root().clone(), uri.capability_path().clone()),
1280        );
1281        assert_eq!(results.len(), 1);
1282    }
1283
1284    /// A trust root that vouches for `agent_key`, and a DHT that trusts it.
1285    fn attesting_dht(agent_key: &SigningKey, uri: &AgentUri) -> (SimulatedDht, String) {
1286        let root_key = SigningKey::generate();
1287        let issuer = Issuer::new("anthropic.com", root_key.clone(), Duration::from_hours(1));
1288        let mut verifier = Verifier::new().with_revocation(AcceptAll);
1289        verifier.add_trusted_root("anthropic.com", root_key.verifying_key());
1290        let token = issuer
1291            .issue(
1292                uri,
1293                &agent_key.verifying_key(),
1294                vec!["assistant/chat".to_string()],
1295            )
1296            .unwrap();
1297        (
1298            SimulatedDht::with_verifier(SimulationConfig::default(), verifier),
1299            token,
1300        )
1301    }
1302
1303    #[test]
1304    fn a_lifted_token_does_not_register_the_thief_at_its_own_endpoints() {
1305        // The replay this binding exists to stop. Registration records are
1306        // world-readable by design, so assume the attacker read one and kept
1307        // the token. Presenting it under their own key must fail, because the
1308        // trust root attested the victim's key, not theirs.
1309        let victim = SigningKey::generate();
1310        let attacker = SigningKey::generate();
1311        let uri = test_uri("2q");
1312        let (dht, lifted) = attesting_dht(&victim, &uri);
1313
1314        let hijack = Registration::new(
1315            uri.clone(),
1316            attacker.verifying_key(),
1317            vec![Endpoint::https("attacker.example")],
1318        )
1319        .with_attestation(lifted);
1320
1321        let result = register(&dht, hijack, &attacker);
1322
1323        assert!(matches!(result, Err(DhtError::AgentKeyMismatch { .. })));
1324    }
1325
1326    #[test]
1327    fn a_lifted_token_does_not_register_the_thief_under_the_victims_key_either() {
1328        // Keeping the victim's key in the record satisfies the token, so the
1329        // register proof is what has to stop this one: the attacker cannot
1330        // sign for a key they do not hold.
1331        let victim = SigningKey::generate();
1332        let attacker = SigningKey::generate();
1333        let uri = test_uri("2q");
1334        let (dht, lifted) = attesting_dht(&victim, &uri);
1335
1336        let hijack = Registration::new(
1337            uri.clone(),
1338            victim.verifying_key(),
1339            vec![Endpoint::https("attacker.example")],
1340        )
1341        .with_attestation(lifted);
1342
1343        let result = register(&dht, hijack, &attacker);
1344
1345        assert!(matches!(
1346            result,
1347            Err(DhtError::Unauthorized {
1348                operation: MutationKind::Register,
1349                ..
1350            })
1351        ));
1352    }
1353
1354    #[test]
1355    fn the_attested_agent_registers_normally() {
1356        let agent = SigningKey::generate();
1357        let uri = test_uri("2q");
1358        let (dht, token) = attesting_dht(&agent, &uri);
1359
1360        register(
1361            &dht,
1362            registration(&uri, &agent).with_attestation(token),
1363            &agent,
1364        )
1365        .unwrap();
1366
1367        assert_eq!(
1368            lookup(&dht, &exact("anthropic.com", "assistant/chat")).len(),
1369            1
1370        );
1371    }
1372
1373    #[test]
1374    fn a_registration_proof_does_not_cover_endpoints_it_did_not_sign() {
1375        // Signed by the right key, for the right URI, but over other
1376        // endpoints. Without binding them, a relay could forward a
1377        // registration pointing somewhere else.
1378        let dht = unverified_dht();
1379        let agent = SigningKey::generate();
1380        let uri = test_uri("2q");
1381        let honest = Registration::new(
1382            uri.clone(),
1383            agent.verifying_key(),
1384            vec![Endpoint::https("honest.example")],
1385        );
1386        let proof = MutationProof::sign_registration(&agent, &honest);
1387
1388        let substituted = Registration::new(
1389            uri,
1390            agent.verifying_key(),
1391            vec![Endpoint::https("attacker.example")],
1392        )
1393        .with_registered_at(honest.registered_at());
1394
1395        let result = block_on(dht.register(substituted, &proof, WriteOptions::default()));
1396
1397        assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1398    }
1399
1400    #[test]
1401    fn a_registration_proof_does_not_cover_a_different_expiry() {
1402        let dht = unverified_dht();
1403        let agent = SigningKey::generate();
1404        let uri = test_uri("2q");
1405        let honest = registration(&uri, &agent).with_ttl(Duration::from_mins(1));
1406        let proof = MutationProof::sign_registration(&agent, &honest);
1407
1408        let stretched = honest.clone().with_ttl(Duration::from_hours(24));
1409
1410        let result = block_on(dht.register(stretched, &proof, WriteOptions::default()));
1411
1412        assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1413    }
1414
1415    #[test]
1416    fn a_registration_proof_does_not_transfer_to_another_agent() {
1417        let dht = unverified_dht();
1418        let agent = SigningKey::generate();
1419        let mine = registration(&test_uri("2q"), &agent);
1420        let proof = MutationProof::sign_registration(&agent, &mine);
1421
1422        let theirs = registration(&test_uri("2r"), &agent).with_registered_at(mine.registered_at());
1423
1424        let result = block_on(dht.register(theirs, &proof, WriteOptions::default()));
1425
1426        assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1427    }
1428
1429    #[test]
1430    fn refresh_extends_an_expiring_registration() {
1431        let dht = unverified_dht();
1432        let key = SigningKey::generate();
1433        let uri = test_uri("2q");
1434        register(
1435            &dht,
1436            registration(&uri, &key).with_ttl(Duration::from_millis(50)),
1437            &key,
1438        )
1439        .unwrap();
1440
1441        let before = lookup(&dht, &exact("anthropic.com", "assistant/chat"))[0].expires_at();
1442        refresh(&dht, &uri, &key, 1, Duration::from_hours(1)).unwrap();
1443        let after = lookup(&dht, &exact("anthropic.com", "assistant/chat"))[0].expires_at();
1444
1445        assert!(
1446            after > before,
1447            "refresh must push expiry out, got {before:?} -> {after:?}"
1448        );
1449    }
1450
1451    #[test]
1452    fn a_refresh_lands_on_the_instant_the_agent_signed() {
1453        // The store applies the expiry from the proof rather than deriving one
1454        // from the TTL. A store that recomputed `now + ttl` would write an
1455        // instant the agent never signed, which is invisible here but leaves
1456        // the field open to anyone relaying the record on a real overlay
1457        // (issue #81).
1458        let dht = unverified_dht();
1459        let key = SigningKey::generate();
1460        let uri = test_uri("2q");
1461        register(&dht, registration(&uri, &key), &key).unwrap();
1462
1463        let ttl = Duration::from_hours(1);
1464        // Deliberately not `now + ttl`: only a store honouring the signed
1465        // instant lands here.
1466        let chosen = SystemTime::now() + Duration::from_hours(9);
1467        let proof = MutationProof::sign(
1468            &key,
1469            &uri,
1470            version_at(&dht, &uri, 1),
1471            &Mutation::Refresh {
1472                ttl,
1473                expires_at: chosen,
1474            },
1475        );
1476        block_on(dht.refresh(&uri, ttl, chosen, &proof, WriteOptions::default())).unwrap();
1477
1478        let stored = lookup(&dht, &exact("anthropic.com", "assistant/chat"))[0].expires_at();
1479        assert_eq!(
1480            system_time_to_millis(stored),
1481            system_time_to_millis(chosen),
1482            "the store wrote an expiry the agent did not sign"
1483        );
1484    }
1485
1486    #[test]
1487    fn a_migration_leaves_the_expiry_where_it_was() {
1488        // A migration moves an agent, not the moment its record lapses. The
1489        // proof names the record's own expiry, so a store that wrote a
1490        // different one would be writing a record its own proof does not
1491        // describe.
1492        let dht = unverified_dht();
1493        let key = SigningKey::generate();
1494        let uri = test_uri("2q");
1495        register(&dht, registration(&uri, &key), &key).unwrap();
1496        let before = record(&dht, &uri).expires_at();
1497
1498        update_endpoint(
1499            &dht,
1500            &uri,
1501            &key,
1502            1,
1503            vec![Endpoint::https("eu-west-1.example")],
1504        )
1505        .unwrap();
1506
1507        assert_eq!(
1508            system_time_to_millis(record(&dht, &uri).expires_at()),
1509            system_time_to_millis(before)
1510        );
1511    }
1512
1513    #[test]
1514    fn refresh_updates_every_ancestor_copy() {
1515        // The record is materialized at each ancestor key. A refresh that
1516        // touched only one would leave the others to expire, silently pruning
1517        // the agent from prefix queries while exact lookup still found it.
1518        let dht = unverified_dht();
1519        let key = SigningKey::generate();
1520        let uri = test_uri("2q");
1521        register(
1522            &dht,
1523            registration(&uri, &key).with_ttl(Duration::from_millis(50)),
1524            &key,
1525        )
1526        .unwrap();
1527
1528        refresh(&dht, &uri, &key, 1, Duration::from_hours(1)).unwrap();
1529
1530        let from_ancestor = lookup(&dht, &prefix("anthropic.com", "assistant"));
1531        assert_eq!(from_ancestor.len(), 1);
1532        assert!(!from_ancestor[0].is_expired());
1533    }
1534
1535    /// A registration that lapsed before it was ever looked at. Every test
1536    /// below about eviction needs one, and none of them should have to wait
1537    /// for a TTL to run out to get it.
1538    fn lapsed(uri: &AgentUri, key: &SigningKey) -> Registration {
1539        registration(uri, key).with_expires_at(SystemTime::now() - Duration::from_secs(1))
1540    }
1541
1542    /// What the store is holding, expired records included. `lookup` filters
1543    /// those out, so it cannot tell eviction from concealment.
1544    fn held(dht: &SimulatedDht) -> usize {
1545        dht.stats().total_registrations()
1546    }
1547
1548    #[test]
1549    fn expire_stale_removes_a_lapsed_record_from_every_index() {
1550        let dht = unverified_dht();
1551        let key = SigningKey::generate();
1552        let uri = test_uri("2q");
1553        register(&dht, lapsed(&uri, &key), &key).unwrap();
1554
1555        assert_eq!(dht.expire_stale(), 1);
1556
1557        assert_eq!(held(&dht), 0, "the record is still indexed by URI");
1558        assert_eq!(dht.stats().unique_keys(), 0, "an ancestor key outlived it");
1559        assert!(lookup(&dht, &prefix("anthropic.com", "assistant")).is_empty());
1560    }
1561
1562    #[test]
1563    fn expire_stale_leaves_a_live_record_alone() {
1564        let dht = unverified_dht();
1565        let key = SigningKey::generate();
1566        let uri = test_uri("2q");
1567        register(&dht, registration(&uri, &key), &key).unwrap();
1568
1569        assert_eq!(dht.expire_stale(), 0);
1570        assert!(stored(&dht, &uri).is_some());
1571    }
1572
1573    #[test]
1574    fn a_write_evicts_what_has_lapsed_without_being_asked() {
1575        // The accumulation issue #55 reports: eviction used to happen only
1576        // when someone remembered to ask for it, so a simulation that ran for
1577        // a day and never called `expire_stale` held a day of garbage.
1578        let dht = unverified_dht();
1579        let key = SigningKey::generate();
1580        register(&dht, lapsed(&test_uri("2q"), &key), &key).unwrap();
1581
1582        register(&dht, registration(&test_uri("2r"), &key), &key).unwrap();
1583
1584        assert_eq!(held(&dht), 1, "the lapsed record is still stored");
1585    }
1586
1587    #[test]
1588    fn a_lapsed_registration_does_not_hold_its_name() {
1589        // Garbage that is merely filtered still occupies the namespace: the
1590        // URI reads as taken, and the agent coming back to it is told so.
1591        let dht = unverified_dht();
1592        let key = SigningKey::generate();
1593        let uri = test_uri("2q");
1594        register(&dht, lapsed(&uri, &key), &key).unwrap();
1595
1596        register(&dht, registration(&uri, &key), &key)
1597            .expect("a lapsed record does not keep its URI registered");
1598
1599        assert_eq!(held(&dht), 1);
1600        assert!(stored(&dht, &uri).is_some_and(|r| !r.is_expired()));
1601    }
1602
1603    #[test]
1604    fn a_lapsed_registration_does_not_hold_its_capacity() {
1605        // The same thing one level out: the slot a lapsed record occupies at
1606        // its key was being counted against agents that arrived after it.
1607        let dht = dht_with_capacity_for_one();
1608        let key = SigningKey::generate();
1609        register(&dht, lapsed(&test_uri("2q"), &key), &key).unwrap();
1610
1611        register(&dht, registration(&test_uri("2r"), &key), &key)
1612            .expect("a lapsed record does not keep its slot");
1613    }
1614
1615    #[test]
1616    fn an_expiry_moved_earlier_is_evicted_at_the_instant_it_names() {
1617        // Deadlines are queued in order, so the entry this write supersedes
1618        // sits ahead of it. A schedule that waited for the entry it already
1619        // held would keep this record for the hour it was originally given.
1620        let dht = unverified_dht();
1621        let key = SigningKey::generate();
1622        let uri = test_uri("2q");
1623        register(&dht, registration(&uri, &key), &key).unwrap();
1624
1625        let ttl = Duration::from_hours(1);
1626        let lapsed_at = SystemTime::now() - Duration::from_secs(1);
1627        let proof = MutationProof::sign(
1628            &key,
1629            &uri,
1630            version_at(&dht, &uri, 1),
1631            &Mutation::Refresh {
1632                ttl,
1633                expires_at: lapsed_at,
1634            },
1635        );
1636        block_on(dht.refresh(&uri, ttl, lapsed_at, &proof, WriteOptions::default())).unwrap();
1637
1638        assert_eq!(held(&dht), 0, "the record lapsed the moment it was written");
1639    }
1640
1641    #[test]
1642    fn a_renewed_record_outlives_the_deadline_it_replaced() {
1643        // The other direction, and the reason a queued deadline is checked
1644        // against the record before it is acted on: the deadline this refresh
1645        // replaces still arrives, and acting on it would evict a record that
1646        // was renewed well before it lapsed.
1647        let dht = unverified_dht();
1648        let key = SigningKey::generate();
1649        let uri = test_uri("2q");
1650        register(
1651            &dht,
1652            registration(&uri, &key).with_ttl(Duration::from_millis(40)),
1653            &key,
1654        )
1655        .unwrap();
1656
1657        refresh(&dht, &uri, &key, 1, Duration::from_hours(1)).unwrap();
1658        std::thread::sleep(Duration::from_millis(60));
1659
1660        assert_eq!(dht.expire_stale(), 0, "the renewed record was evicted");
1661        assert!(stored(&dht, &uri).is_some());
1662    }
1663
1664    #[test]
1665    fn expiry_waits_to_be_asked_when_auto_expire_is_off() {
1666        // An evaluation measuring how much of a namespace has gone stale needs
1667        // the stale part still there to measure. Turning expiry off stops it
1668        // happening unasked; it does not stop it happening.
1669        let dht = SimulatedDht::new(
1670            SimulationConfig::new()
1671                .with_verify_attestations(false)
1672                .with_auto_expire(false),
1673        );
1674        let key = SigningKey::generate();
1675        register(&dht, lapsed(&test_uri("2q"), &key), &key).unwrap();
1676
1677        register(&dht, registration(&test_uri("2r"), &key), &key).unwrap();
1678        assert_eq!(held(&dht), 2, "a write evicted what the caller kept");
1679
1680        assert_eq!(dht.expire_stale(), 1);
1681        assert_eq!(held(&dht), 1);
1682    }
1683
1684    #[test]
1685    fn every_stored_registration_has_a_deadline() {
1686        // The invariant the whole schedule rests on. A record that reaches the
1687        // store without one is never evicted; a deadline left behind by a
1688        // record that has gone is a name that cannot be registered again, once
1689        // some later write re-uses it.
1690        let dht = unverified_dht();
1691        let key = SigningKey::generate();
1692        for suffix in ["2q", "2r", "2s"] {
1693            register(&dht, registration(&test_uri(suffix), &key), &key).unwrap();
1694        }
1695        refresh(&dht, &test_uri("2q"), &key, 1, Duration::from_hours(2)).unwrap();
1696        deregister(&dht, &test_uri("2s"), &key, 1).unwrap();
1697        update_endpoint(
1698            &dht,
1699            &test_uri("2r"),
1700            &key,
1701            1,
1702            vec![Endpoint::https("eu-west-1.example")],
1703        )
1704        .unwrap();
1705
1706        let indices = dht.indices.read().expect("lock poisoned");
1707        assert_eq!(
1708            indices.expirations.tracked(),
1709            indices.keys_by_uri.len(),
1710            "the schedule and the store disagree about what is registered"
1711        );
1712    }
1713
1714    #[test]
1715    fn refresh_of_unknown_agent_is_not_found() {
1716        let dht = unverified_dht();
1717        let key = SigningKey::generate();
1718        let result = refresh(&dht, &test_uri("2q"), &key, 1, Duration::from_hours(1));
1719        assert!(matches!(result, Err(DhtError::NotFound { .. })));
1720    }
1721
1722    #[test]
1723    fn lookup_pages_and_the_pages_partition_the_results() {
1724        let config = SimulationConfig::new()
1725            .with_verify_attestations(false)
1726            .with_page_size(NonZeroUsize::new(2).unwrap());
1727        let dht = SimulatedDht::new(config);
1728        let key = SigningKey::generate();
1729
1730        for suffix in ["2q", "2r", "2s", "2t", "2v"] {
1731            register(&dht, registration(&test_uri(suffix), &key), &key).unwrap();
1732        }
1733
1734        let query = exact("anthropic.com", "assistant/chat");
1735        let mut options = ReadOptions::default();
1736        let mut seen: Vec<String> = Vec::new();
1737        let mut pages = 0;
1738        loop {
1739            let page = block_on(dht.lookup(&query, &options)).unwrap();
1740            pages += 1;
1741            assert!(page.len() <= 2, "page exceeded the configured page size");
1742            let next = page.next_cursor();
1743            seen.extend(page.into_items().iter().map(|r| r.agent_uri().canonical()));
1744            match next {
1745                Some(cursor) => options = options.with_cursor(cursor),
1746                None => break,
1747            }
1748        }
1749
1750        assert_eq!(pages, 3, "5 results at page size 2 should take 3 pages");
1751        assert_eq!(seen.len(), 5);
1752        let mut deduped = seen.clone();
1753        deduped.sort();
1754        deduped.dedup();
1755        assert_eq!(deduped.len(), 5, "pages must not repeat a registration");
1756    }
1757
1758    #[test]
1759    fn paging_order_is_stable_across_calls() {
1760        // Cursors index into the result order, so an unstable order would let a
1761        // second page skip or repeat records that the first already moved past.
1762        let config = SimulationConfig::new()
1763            .with_verify_attestations(false)
1764            .with_page_size(NonZeroUsize::new(2).unwrap());
1765        let dht = SimulatedDht::new(config);
1766        let key = SigningKey::generate();
1767        for suffix in ["2v", "2q", "2s", "2r"] {
1768            register(&dht, registration(&test_uri(suffix), &key), &key).unwrap();
1769        }
1770
1771        let query = exact("anthropic.com", "assistant/chat");
1772        let first = block_on(dht.lookup(&query, &ReadOptions::default())).unwrap();
1773        let again = block_on(dht.lookup(&query, &ReadOptions::default())).unwrap();
1774
1775        assert_eq!(
1776            first
1777                .items()
1778                .iter()
1779                .map(|r| r.agent_uri().canonical())
1780                .collect::<Vec<_>>(),
1781            again
1782                .items()
1783                .iter()
1784                .map(|r| r.agent_uri().canonical())
1785                .collect::<Vec<_>>()
1786        );
1787    }
1788
1789    #[test]
1790    fn cursor_past_the_end_yields_an_exhausted_page() {
1791        let dht = unverified_dht();
1792        let key = SigningKey::generate();
1793        register(&dht, registration(&test_uri("2q"), &key), &key).unwrap();
1794
1795        let page = block_on(dht.lookup(
1796            &exact("anthropic.com", "assistant/chat"),
1797            &ReadOptions::default().with_cursor(Cursor::at(99)),
1798        ))
1799        .unwrap();
1800
1801        assert!(page.is_empty());
1802        assert!(!page.has_more());
1803    }
1804
1805    #[test]
1806    fn write_receipt_reports_the_resolved_quorum() {
1807        let dht = unverified_dht();
1808        let key = SigningKey::generate();
1809        let receipt = register(&dht, registration(&test_uri("2q"), &key), &key).unwrap();
1810
1811        // One authoritative copy, so every quorum resolves to one against it.
1812        assert_eq!(receipt.acknowledged(), 1);
1813        assert_eq!(receipt.required(), 1);
1814        assert!(!receipt.exceeded_quorum());
1815    }
1816
1817    #[test]
1818    fn instances_have_distinct_identities() {
1819        let a = unverified_dht();
1820        let b = unverified_dht();
1821        assert_ne!(a.local_id(), b.local_id());
1822        assert!(!a.local_id().is_empty());
1823    }
1824
1825    #[test]
1826    fn bootstrap_without_peers_fails_the_way_a_real_backend_would() {
1827        let dht = unverified_dht();
1828        assert!(matches!(
1829            block_on(dht.bootstrap(&[])),
1830            Err(DhtError::NoPeers)
1831        ));
1832        assert!(block_on(dht.bootstrap(&[PeerAddr::new("/ip4/127.0.0.1/tcp/4001")])).is_ok());
1833    }
1834
1835    #[test]
1836    fn the_simulator_is_usable_as_a_trait_object() {
1837        // Dyn-compatibility is the reason this trait uses async_trait at all,
1838        // so it needs a test that actually fails if that regresses.
1839        let key = SigningKey::generate();
1840        let dht: Box<dyn Dht> = Box::new(unverified_dht());
1841        let record = registration(&test_uri("2q"), &key);
1842        let proof = MutationProof::sign_registration(&key, &record);
1843        block_on(dht.register(record, &proof, WriteOptions::default())).unwrap();
1844
1845        let page = block_on(dht.lookup(
1846            &exact("anthropic.com", "assistant/chat"),
1847            &ReadOptions::default(),
1848        ))
1849        .unwrap();
1850        assert_eq!(page.len(), 1);
1851    }
1852
1853    #[test]
1854    fn agent_id_cannot_be_reused_under_a_different_capability_path() {
1855        let dht = unverified_dht();
1856        let key = SigningKey::generate();
1857        let first =
1858            AgentUri::parse("agent://anthropic.com/assistant/chat/llm_01h455vb4pex5vsknk084sn02q")
1859                .unwrap();
1860        let reclassified =
1861            AgentUri::parse("agent://anthropic.com/assistant/code/llm_01h455vb4pex5vsknk084sn02q")
1862                .unwrap();
1863        register(&dht, registration(&first, &key), &key).unwrap();
1864
1865        let result = register(&dht, registration(&reclassified, &key), &key);
1866
1867        assert!(matches!(
1868            result,
1869            Err(DhtError::IdentityCapabilityConflict { .. })
1870        ));
1871
1872        // Ancestor materialization must not make the ancestor look like a
1873        // legitimate second identity for the same Agent ID.
1874        let dht = unverified_dht();
1875        let first =
1876            AgentUri::parse("agent://anthropic.com/assistant/chat/llm_01h455vb4pex5vsknk084sn02q")
1877                .unwrap();
1878        let ancestor_reuse =
1879            AgentUri::parse("agent://anthropic.com/assistant/llm_01h455vb4pex5vsknk084sn02q")
1880                .unwrap();
1881        register(&dht, registration(&first, &key), &key).unwrap();
1882
1883        let result = register(&dht, registration(&ancestor_reuse, &key), &key);
1884
1885        assert!(matches!(
1886            result,
1887            Err(DhtError::IdentityCapabilityConflict { .. })
1888        ));
1889    }
1890
1891    #[test]
1892    fn an_attacker_cannot_repoint_an_agents_endpoints() {
1893        // The hijack this authorization exists to stop. The attacker knows the
1894        // URI, which is public, and holds every key except the agent's.
1895        let dht = unverified_dht();
1896        let agent = SigningKey::generate();
1897        let attacker = SigningKey::generate();
1898        let uri = test_uri("2q");
1899        register(&dht, registration(&uri, &agent), &agent).unwrap();
1900
1901        let result = update_endpoint(
1902            &dht,
1903            &uri,
1904            &attacker,
1905            1,
1906            vec![Endpoint::https("attacker.example")],
1907        );
1908
1909        assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1910        assert_eq!(
1911            lookup(&dht, &exact("anthropic.com", "assistant/chat"))[0].endpoints(),
1912            &[test_endpoint()]
1913        );
1914    }
1915
1916    #[test]
1917    fn an_attacker_cannot_evict_an_agent() {
1918        let dht = unverified_dht();
1919        let agent = SigningKey::generate();
1920        let attacker = SigningKey::generate();
1921        let uri = test_uri("2q");
1922        register(&dht, registration(&uri, &agent), &agent).unwrap();
1923
1924        let result = deregister(&dht, &uri, &attacker, 1);
1925
1926        assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1927        assert_eq!(
1928            lookup(&dht, &exact("anthropic.com", "assistant/chat")).len(),
1929            1
1930        );
1931    }
1932
1933    #[test]
1934    fn an_attacker_cannot_refresh_an_agent() {
1935        // Refresh looks harmless, but keeping a record alive that its owner
1936        // meant to let expire is still a write nobody authorized.
1937        let dht = unverified_dht();
1938        let agent = SigningKey::generate();
1939        let attacker = SigningKey::generate();
1940        let uri = test_uri("2q");
1941        register(&dht, registration(&uri, &agent), &agent).unwrap();
1942
1943        let result = refresh(&dht, &uri, &attacker, 1, Duration::from_hours(9));
1944
1945        assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1946    }
1947
1948    #[test]
1949    fn a_proof_does_not_authorize_endpoints_it_did_not_sign() {
1950        // Signing the operation is not enough on its own: the proof has to
1951        // cover the arguments, or an intercepted migration could be re-aimed
1952        // while keeping the agent's own signature.
1953        let dht = unverified_dht();
1954        let agent = SigningKey::generate();
1955        let uri = test_uri("2q");
1956        register(&dht, registration(&uri, &agent), &agent).unwrap();
1957
1958        let honest = vec![Endpoint::https("honest.example")];
1959        let held = record(&dht, &uri);
1960        let proof = MutationProof::sign_next(
1961            &agent,
1962            &held,
1963            &Mutation::UpdateEndpoint {
1964                endpoints: &honest,
1965                expires_at: held.expires_at(),
1966            },
1967        );
1968
1969        let result = block_on(dht.update_endpoint(
1970            &uri,
1971            vec![Endpoint::https("attacker.example")],
1972            &proof,
1973            WriteOptions::default(),
1974        ));
1975
1976        assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1977    }
1978
1979    #[test]
1980    fn a_successful_mutation_advances_the_record_sequence() {
1981        let dht = unverified_dht();
1982        let key = SigningKey::generate();
1983        let uri = test_uri("2q");
1984        register(&dht, registration(&uri, &key), &key).unwrap();
1985
1986        assert_eq!(
1987            lookup(&dht, &exact("anthropic.com", "assistant/chat"))[0].sequence(),
1988            0
1989        );
1990
1991        refresh(&dht, &uri, &key, 4, Duration::from_hours(1)).unwrap();
1992
1993        // Sequences only have to increase, not increase by one: an agent that
1994        // lost a response has no way to learn which of its writes landed.
1995        assert_eq!(
1996            lookup(&dht, &exact("anthropic.com", "assistant/chat"))[0].sequence(),
1997            4
1998        );
1999    }
2000
2001    #[test]
2002    fn every_ancestor_copy_carries_the_new_sequence() {
2003        // A copy left behind at a lower sequence would accept a replay that the
2004        // exact key already rejected.
2005        let dht = unverified_dht();
2006        let key = SigningKey::generate();
2007        let uri = test_uri("2q");
2008        register(&dht, registration(&uri, &key), &key).unwrap();
2009
2010        refresh(&dht, &uri, &key, 3, Duration::from_hours(1)).unwrap();
2011
2012        assert_eq!(
2013            lookup(&dht, &prefix("anthropic.com", "assistant"))[0].sequence(),
2014            3
2015        );
2016    }
2017
2018    #[test]
2019    fn a_replayed_proof_is_refused_the_second_time() {
2020        let dht = unverified_dht();
2021        let key = SigningKey::generate();
2022        let uri = test_uri("2q");
2023        register(&dht, registration(&uri, &key), &key).unwrap();
2024
2025        let endpoints = vec![Endpoint::https("first.example")];
2026        let held = record(&dht, &uri);
2027        let proof = MutationProof::sign_next(
2028            &key,
2029            &held,
2030            &Mutation::UpdateEndpoint {
2031                endpoints: &endpoints,
2032                expires_at: held.expires_at(),
2033            },
2034        );
2035
2036        block_on(dht.update_endpoint(&uri, endpoints.clone(), &proof, WriteOptions::default()))
2037            .unwrap();
2038        let replay =
2039            block_on(dht.update_endpoint(&uri, endpoints, &proof, WriteOptions::default()));
2040
2041        assert!(matches!(
2042            replay,
2043            Err(DhtError::StaleSequence {
2044                presented: 1,
2045                current: 1,
2046                ..
2047            })
2048        ));
2049    }
2050
2051    #[test]
2052    fn a_captured_proof_does_not_reach_a_later_record_instance() {
2053        // Proofs travel in the clear, so assume the attacker kept one. A
2054        // re-registered record starts counting from zero again, so the sequence
2055        // number alone would leave the captured copy sitting above it.
2056        let dht = unverified_dht();
2057        let key = SigningKey::generate();
2058        let uri = test_uri("2q");
2059        register(&dht, registration(&uri, &key), &key).unwrap();
2060
2061        let captured = MutationProof::sign_next(&key, &record(&dht, &uri), &Mutation::Deregister);
2062        block_on(dht.deregister(&uri, &captured, WriteOptions::default())).unwrap();
2063
2064        // The agent comes back a second later, which is what makes this a
2065        // different record instance rather than a continuation of the old one.
2066        register(
2067            &dht,
2068            registration(&uri, &key).with_registered_at(SystemTime::now() + Duration::from_secs(1)),
2069            &key,
2070        )
2071        .unwrap();
2072        let replay = block_on(dht.deregister(&uri, &captured, WriteOptions::default()));
2073
2074        assert!(matches!(replay, Err(DhtError::Unauthorized { .. })));
2075        assert_eq!(
2076            lookup(&dht, &exact("anthropic.com", "assistant/chat")).len(),
2077            1
2078        );
2079    }
2080
2081    #[test]
2082    fn continuing_the_sequence_defeats_a_captured_proof_without_relying_on_the_clock() {
2083        // Instance identity is the registration time, which a coarse clock can
2084        // fail to distinguish. An agent that remembers how far it counted does
2085        // not depend on the clock: it re-registers above every proof it ever
2086        // signed.
2087        let dht = unverified_dht();
2088        let key = SigningKey::generate();
2089        let uri = test_uri("2q");
2090        let opened_at = SystemTime::now();
2091        register(
2092            &dht,
2093            registration(&uri, &key).with_registered_at(opened_at),
2094            &key,
2095        )
2096        .unwrap();
2097
2098        let captured = MutationProof::sign_next(&key, &record(&dht, &uri), &Mutation::Deregister);
2099        block_on(dht.deregister(&uri, &captured, WriteOptions::default())).unwrap();
2100
2101        // Same instant, so the instance binding cannot help here.
2102        register(
2103            &dht,
2104            registration(&uri, &key)
2105                .with_registered_at(opened_at)
2106                .with_sequence(captured.sequence()),
2107            &key,
2108        )
2109        .unwrap();
2110        let replay = block_on(dht.deregister(&uri, &captured, WriteOptions::default()));
2111
2112        assert!(matches!(replay, Err(DhtError::StaleSequence { .. })));
2113        assert_eq!(
2114            lookup(&dht, &exact("anthropic.com", "assistant/chat")).len(),
2115            1
2116        );
2117    }
2118
2119    #[test]
2120    fn a_write_at_the_records_own_sequence_does_not_advance_it() {
2121        let dht = unverified_dht();
2122        let key = SigningKey::generate();
2123        let uri = test_uri("2q");
2124        register(&dht, registration(&uri, &key), &key).unwrap();
2125
2126        // The record opens at 0, so 0 is already spent.
2127        let result = refresh(&dht, &uri, &key, 0, Duration::from_hours(1));
2128
2129        assert!(matches!(
2130            result,
2131            Err(DhtError::StaleSequence {
2132                presented: 0,
2133                current: 0,
2134                ..
2135            })
2136        ));
2137    }
2138
2139    #[test]
2140    fn migration_is_authorized_by_the_agents_own_key() {
2141        let dht = unverified_dht();
2142        let agent = SigningKey::generate();
2143        let attacker = SigningKey::generate();
2144        let uri = test_uri("2q");
2145        register(&dht, registration(&uri, &agent), &agent).unwrap();
2146
2147        let moved = Endpoint::https("eu-west-1.agent.anthropic.com");
2148        assert!(
2149            dht.simulate_migration(&uri, moved.clone(), &attacker)
2150                .is_err()
2151        );
2152
2153        let result = dht.simulate_migration(&uri, moved.clone(), &agent).unwrap();
2154
2155        assert_eq!(result.new_endpoints(), &[moved]);
2156    }
2157
2158    /// Registration and migration used to take the simulator's indices in
2159    /// opposite orders, back when each had its own lock: a registration held
2160    /// the record store and waited on the URI index while a migration held the
2161    /// URI index and waited on the record store, and neither ever moved
2162    /// (issue #47). Reading the lock sites is what found that; nothing
2163    /// ever ran the two against each other, which is how it outlived three
2164    /// refactors of this file.
2165    ///
2166    /// One lock has replaced the three, so the hang is now impossible by
2167    /// construction rather than avoided by convention, and this test cannot
2168    /// fail for the original reason. What it does is keep the pairing
2169    /// exercised: it asserts the answers a caller is owed when the two paths
2170    /// overlap, and it bounds the wait, because a lock-ordering regression
2171    /// reports as a stuck run rather than a failing one unless something is
2172    /// holding a stopwatch.
2173    #[test]
2174    fn concurrent_registration_and_migration_do_not_deadlock() {
2175        use std::sync::Arc;
2176        use std::sync::mpsc;
2177        use std::thread;
2178
2179        const WORKERS: usize = 4;
2180        const ROUNDS: usize = 60;
2181        /// Crockford base32, which is what an agent ID's characters are drawn
2182        /// from. `test_uri` appends two of them.
2183        const ALPHABET: &[u8] = b"0123456789abcdefghjkmnpqrstvwxyz";
2184
2185        fn agent(index: usize) -> AgentUri {
2186            test_uri(&format!(
2187                "{}{}",
2188                ALPHABET[index / ALPHABET.len()] as char,
2189                ALPHABET[index % ALPHABET.len()] as char,
2190            ))
2191        }
2192
2193        let dht = Arc::new(unverified_dht());
2194
2195        // Two disjoint sets, so the threads contend on the indices rather than
2196        // on each other's records. A migrator moves an agent that is already
2197        // registered; a churner adds and removes its own. Sharing an agent
2198        // would test sequence handling instead, and would make a lost race
2199        // look like the deadlock this is watching for.
2200        let migrators: Vec<(AgentUri, SigningKey)> = (0..WORKERS)
2201            .map(|index| {
2202                let (uri, key) = (agent(index), SigningKey::generate());
2203                register(&dht, registration(&uri, &key), &key).expect("a fresh agent registers");
2204                (uri, key)
2205            })
2206            .collect();
2207        let churners: Vec<(AgentUri, SigningKey)> = (WORKERS..2 * WORKERS)
2208            .map(|index| (agent(index), SigningKey::generate()))
2209            .collect();
2210
2211        let mut handles = Vec::new();
2212        for (uri, key) in migrators {
2213            let dht = Arc::clone(&dht);
2214            handles.push(thread::spawn(move || {
2215                for round in 0..ROUNDS {
2216                    let endpoint = Endpoint::https(format!("eu-west-{round}.anthropic.com"));
2217                    dht.simulate_migration(&uri, endpoint, &key)
2218                        .expect("a registered agent migrates");
2219                }
2220            }));
2221        }
2222        for (uri, key) in churners {
2223            let dht = Arc::clone(&dht);
2224            handles.push(thread::spawn(move || {
2225                for _ in 0..ROUNDS {
2226                    register(&dht, registration(&uri, &key), &key)
2227                        .expect("a deregistered agent registers again");
2228                    deregister(&dht, &uri, &key, 1).expect("its own key removes it");
2229                }
2230            }));
2231        }
2232
2233        // Joining on this thread rather than the test's, so that a hang is a
2234        // failure with a name on it instead of a run that never returns. A
2235        // panicking worker still surfaces as a panic: its join result is
2236        // carried back and unwrapped below.
2237        let (done, finished) = mpsc::channel();
2238        thread::spawn(move || {
2239            let outcomes: Vec<_> = handles.into_iter().map(thread::JoinHandle::join).collect();
2240            let _ = done.send(outcomes);
2241        });
2242
2243        let outcomes = finished
2244            .recv_timeout(Duration::from_mins(1))
2245            .expect("registration and migration deadlocked against each other");
2246        for outcome in outcomes {
2247            outcome.expect("no worker panicked");
2248        }
2249
2250        // Every migrator's agent is still registered and every churner's last
2251        // act was to leave, so the indices agree on who is present.
2252        let registered = lookup(&dht, &prefix("anthropic.com", "assistant"));
2253        assert_eq!(registered.len(), WORKERS);
2254        assert_eq!(dht.stats().total_registrations, WORKERS);
2255    }
2256}