This repository has no description
1use std::collections::{HashMap, HashSet};
2use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
3use std::sync::{Arc, Mutex};
4use std::time::Duration;
5
6use futures::StreamExt;
7use knot_atproto::Atproto;
8use knot_index::{Index, IndexGeneration, Resolved};
9use knot_runtime::{Clock, HttpTransport, UnixMicros};
10use knot_types::{AccountDid, OfferedKey};
11
12const FRESH_TTL: Duration = Duration::from_secs(60);
13const DEGRADED_TTL: Duration = Duration::from_secs(5);
14const BACKOFF_SHIFT_LIMIT: u32 = 4;
15const MISS_REVALIDATE_BUDGET: Duration = Duration::from_secs(30);
16const RESOLVE_FANOUT: usize = 16;
17
18fn degraded_ttl(consecutive_failures: u32) -> Duration {
19 let secs = DEGRADED_TTL
20 .as_secs()
21 .saturating_mul(1u64 << consecutive_failures.min(BACKOFF_SHIFT_LIMIT))
22 .min(FRESH_TTL.as_secs());
23 Duration::from_secs(secs)
24}
25
26struct Freshness {
27 due: UnixMicros,
28 generation: IndexGeneration,
29}
30
31#[derive(Debug, PartialEq, Eq)]
32enum Staleness {
33 Fresh,
34 Revalidate,
35 Cold,
36}
37
38pub(crate) struct KeyRoster {
39 by_did: Mutex<HashMap<AccountDid, HashSet<OfferedKey>>>,
40 recognized: Mutex<HashSet<OfferedKey>>,
41 freshness: Mutex<Option<Freshness>>,
42 failures: AtomicU32,
43 refresh: tokio::sync::Mutex<()>,
44 refresh_in_flight: AtomicBool,
45}
46
47impl KeyRoster {
48 pub(crate) fn new() -> Self {
49 Self {
50 by_did: Mutex::new(HashMap::new()),
51 recognized: Mutex::new(HashSet::new()),
52 freshness: Mutex::new(None),
53 failures: AtomicU32::new(0),
54 refresh: tokio::sync::Mutex::new(()),
55 refresh_in_flight: AtomicBool::new(false),
56 }
57 }
58
59 pub(crate) fn recognizes(&self, key: &OfferedKey) -> bool {
60 self.recognized
61 .lock()
62 .unwrap_or_else(|poisoned| poisoned.into_inner())
63 .contains(key)
64 }
65
66 pub(crate) fn did_for(&self, key: &OfferedKey) -> Option<AccountDid> {
67 self.by_did
68 .lock()
69 .unwrap_or_else(|poisoned| poisoned.into_inner())
70 .iter()
71 .find(|(_, keys)| keys.contains(key))
72 .map(|(did, _)| did.clone())
73 }
74
75 fn is_fresh(&self, now: UnixMicros, generation: IndexGeneration) -> bool {
76 self.freshness
77 .lock()
78 .unwrap_or_else(|poisoned| poisoned.into_inner())
79 .as_ref()
80 .is_some_and(|fresh| now.get() < fresh.due.get() && fresh.generation == generation)
81 }
82
83 pub(crate) fn prime<H: HttpTransport, C: Clock>(
84 self: &Arc<Self>,
85 index: &Arc<Index>,
86 atproto: &Arc<Atproto<H, C>>,
87 ) {
88 self.spawn_refresh(index, atproto);
89 }
90
91 pub(crate) fn ensure_fresh<H: HttpTransport, C: Clock>(
92 self: &Arc<Self>,
93 index: &Arc<Index>,
94 atproto: &Arc<Atproto<H, C>>,
95 ) {
96 match self.staleness(atproto.now(), index.generation()) {
97 Staleness::Fresh => {}
98 Staleness::Revalidate | Staleness::Cold => self.spawn_refresh(index, atproto),
99 }
100 }
101
102 pub(crate) async fn recognizes_fresh<H: HttpTransport, C: Clock>(
103 self: &Arc<Self>,
104 key: &OfferedKey,
105 index: &Arc<Index>,
106 atproto: &Arc<Atproto<H, C>>,
107 ) -> bool {
108 if self.recognizes(key) {
109 self.ensure_fresh(index, atproto);
110 return true;
111 }
112 if self.is_fresh(atproto.now(), index.generation()) {
113 return false;
114 }
115 let _ = tokio::time::timeout(MISS_REVALIDATE_BUDGET, self.refresh(index, atproto)).await;
116 self.recognizes(key)
117 }
118
119 fn staleness(&self, now: UnixMicros, generation: IndexGeneration) -> Staleness {
120 match self
121 .freshness
122 .lock()
123 .unwrap_or_else(|poisoned| poisoned.into_inner())
124 .as_ref()
125 {
126 None => Staleness::Cold,
127 Some(fresh) if fresh.generation != generation => Staleness::Revalidate,
128 Some(fresh) if now.get() < fresh.due.get() => Staleness::Fresh,
129 Some(_) => Staleness::Revalidate,
130 }
131 }
132
133 fn spawn_refresh<H: HttpTransport, C: Clock>(
134 self: &Arc<Self>,
135 index: &Arc<Index>,
136 atproto: &Arc<Atproto<H, C>>,
137 ) {
138 if self
139 .refresh_in_flight
140 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
141 .is_err()
142 {
143 return;
144 }
145 let roster = Arc::clone(self);
146 let index = Arc::clone(index);
147 let atproto = Arc::clone(atproto);
148 tokio::spawn(async move {
149 let _in_flight = InFlightGuard(&roster.refresh_in_flight);
150 roster.refresh(&index, &atproto).await;
151 });
152 }
153
154 async fn refresh<H: HttpTransport, C: Clock>(&self, index: &Index, atproto: &Atproto<H, C>) {
155 let _single_flight = self.refresh.lock().await;
156 if self.is_fresh(atproto.now(), index.generation()) {
157 return;
158 }
159 let generation = index.generation();
160 let (dids, incomplete) = relevant_dids(index);
161 let resolved: Vec<(AccountDid, Option<Vec<OfferedKey>>)> = futures::stream::iter(dids)
162 .map(|did| async move {
163 let keys = atproto.resolve_pubkeys(&did).await.ok();
164 (did, keys)
165 })
166 .buffer_unordered(RESOLVE_FANOUT)
167 .collect()
168 .await;
169 let any_failed = resolved.iter().any(|(_, keys)| keys.is_none());
170 let relevant: HashSet<AccountDid> = resolved.iter().map(|(did, _)| did.clone()).collect();
171 {
172 let mut by_did = self
173 .by_did
174 .lock()
175 .unwrap_or_else(|poisoned| poisoned.into_inner());
176 by_did.retain(|did, _| relevant.contains(did));
177 resolved.into_iter().for_each(|(did, keys)| {
178 if let Some(keys) = keys {
179 by_did.insert(did, keys.into_iter().collect());
180 }
181 });
182 let union: HashSet<OfferedKey> = by_did.values().flatten().cloned().collect();
183 *self
184 .recognized
185 .lock()
186 .unwrap_or_else(|poisoned| poisoned.into_inner()) = union;
187 }
188 let ttl = if any_failed {
189 degraded_ttl(self.failures.fetch_add(1, Ordering::Relaxed))
190 } else {
191 self.failures.store(0, Ordering::Relaxed);
192 if incomplete { DEGRADED_TTL } else { FRESH_TTL }
193 };
194 let due = UnixMicros::new(atproto.now().get().saturating_add(ttl.as_micros() as u64));
195 *self
196 .freshness
197 .lock()
198 .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(Freshness { due, generation });
199 }
200}
201
202struct InFlightGuard<'a>(&'a AtomicBool);
203
204impl Drop for InFlightGuard<'_> {
205 fn drop(&mut self) {
206 self.0.store(false, Ordering::Release);
207 }
208}
209
210fn relevant_dids(index: &Index) -> (Vec<AccountDid>, bool) {
211 let (mut dids, incomplete): (Vec<AccountDid>, bool) = index
212 .hosted_repos()
213 .iter()
214 .map(|repo| {
215 let (owner, owner_warming) = match index.owner_of(repo) {
216 Resolved::Ready(Some(owner)) => (Some(AccountDid::from(owner)), false),
217 Resolved::Ready(None) => (None, false),
218 Resolved::Warming => (None, true),
219 };
220 let (collaborators, collaborators_warming) = match index.collaborators_of(repo) {
221 Resolved::Ready(collaborators) => (collaborators, false),
222 Resolved::Warming => (Vec::new(), true),
223 };
224 (
225 owner.into_iter().chain(collaborators).collect::<Vec<_>>(),
226 owner_warming || collaborators_warming,
227 )
228 })
229 .fold(
230 (Vec::new(), false),
231 |(mut acc, warming), (dids, repo_warming)| {
232 acc.extend(dids);
233 (acc, warming || repo_warming)
234 },
235 );
236 dids.sort();
237 dids.dedup();
238 (dids, incomplete)
239}
240
241#[cfg(test)]
242mod tests {
243 use super::*;
244 use std::sync::Arc;
245 use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
246
247 use knot_atproto::Atproto;
248 use knot_cob::{CobHome, CobStore};
249 use knot_cobs::{Registration, RegistryChange};
250 use knot_git::{Layout, Repo};
251 use knot_runtime::{
252 FakeHttp, HttpRequest, HttpResponse, K256Signer, NetworkError, SeededEntropy, Signer,
253 };
254 use knot_types::{KnotId, OwnerDid, RepoDid, RepoName, RepoRkey, UnixSeconds, crypto};
255 use russh::keys::{Algorithm, PrivateKey};
256 use url::Url;
257
258 struct SharedClock(Arc<AtomicU64>);
259 impl Clock for SharedClock {
260 fn now_unix_micros(&self) -> UnixMicros {
261 UnixMicros::new(self.0.load(Ordering::SeqCst))
262 }
263 }
264
265 fn line_and_offered() -> (String, OfferedKey) {
266 let key = PrivateKey::random(&mut crate::EntropyRng, Algorithm::Ed25519).unwrap();
267 let public = key.public_key();
268 (
269 public.to_openssh().unwrap(),
270 OfferedKey::from_bytes(public.to_bytes().unwrap()),
271 )
272 }
273
274 type Responder = Box<dyn Fn(&HttpRequest) -> Result<HttpResponse, NetworkError> + Send + Sync>;
275
276 struct Harness {
277 index: Arc<Index>,
278 atproto: Arc<Atproto<FakeHttp<Responder>, SharedClock>>,
279 published: Arc<std::sync::Mutex<Vec<String>>>,
280 list_calls: Arc<AtomicUsize>,
281 _dir: tempfile::TempDir,
282 }
283
284 fn harness(initial: Vec<String>) -> Harness {
285 let dir = tempfile::tempdir().unwrap();
286 let meta_path = dir.path().join("meta");
287 Repo::create(&meta_path).unwrap();
288 let layout = Layout::new(dir.path().join("repos"));
289 let repo_did = RepoDid::new("did:plc:squid").unwrap();
290 layout.create(&repo_did).unwrap();
291 let cob_signer = K256Signer::generate(&SeededEntropy::new(2));
292 {
293 let meta = Repo::open(&meta_path).unwrap();
294 CobStore::new(&meta)
295 .create(
296 &CobHome::from(&KnotId::new("did:web:nel.pet").unwrap()),
297 &RegistryChange::Register(Registration {
298 owner: OwnerDid::new("did:plc:nel").unwrap(),
299 rkey: RepoRkey::new("anemone").unwrap(),
300 name: RepoName::new("anemone").unwrap(),
301 repo: repo_did.clone(),
302 created_at: UnixSeconds::new(1),
303 }),
304 &cob_signer,
305 UnixSeconds::new(1),
306 )
307 .unwrap();
308 }
309 let index = Arc::new(Index::new(meta_path, layout.clone()));
310 index.rebuild().unwrap();
311
312 let published = Arc::new(std::sync::Mutex::new(initial));
313 let list_calls = Arc::new(AtomicUsize::new(0));
314 let multikey = crypto::multikey(
315 0xe7,
316 K256Signer::generate(&SeededEntropy::new(7))
317 .public_key()
318 .as_bytes(),
319 );
320 let clock = Arc::new(AtomicU64::new(1_000_000_000));
321
322 let responder: Responder = {
323 let published = Arc::clone(&published);
324 let list_calls = Arc::clone(&list_calls);
325 Box::new(move |request: &HttpRequest| {
326 let host = request.url.host_str().unwrap_or_default().to_string();
327 let body = if host == "pds.oyster.cafe" {
328 list_calls.fetch_add(1, Ordering::SeqCst);
329 let records: Vec<_> = published
330 .lock()
331 .unwrap()
332 .iter()
333 .map(|line| {
334 serde_json::json!({
335 "uri": "at://did:plc:nel/sh.tangled.publicKey/1",
336 "value": {
337 "$type": "sh.tangled.publicKey",
338 "key": line,
339 "name": "laptop",
340 "createdAt": "2026-06-08T00:00:00Z"
341 }
342 })
343 })
344 .collect();
345 serde_json::to_vec(&serde_json::json!({ "records": records })).unwrap()
346 } else if host == "plc.directory" {
347 serde_json::to_vec(&serde_json::json!({
348 "id": "did:plc:nel",
349 "alsoKnownAs": ["at://nel.pet"],
350 "verificationMethod": [{
351 "id": "did:plc:nel#atproto",
352 "type": "Multikey",
353 "controller": "did:plc:nel",
354 "publicKeyMultibase": multikey
355 }],
356 "service": [{
357 "id": "#atproto_pds",
358 "type": "AtprotoPersonalDataServer",
359 "serviceEndpoint": "https://pds.oyster.cafe"
360 }]
361 }))
362 .unwrap()
363 } else {
364 return Ok(HttpResponse {
365 status: http::StatusCode::NOT_FOUND,
366 headers: http::HeaderMap::new(),
367 body: bytes::Bytes::new(),
368 });
369 };
370 Ok(HttpResponse {
371 status: http::StatusCode::OK,
372 headers: http::HeaderMap::new(),
373 body: bytes::Bytes::from(body),
374 })
375 })
376 };
377
378 let atproto = Arc::new(Atproto::new(
379 FakeHttp::new(responder),
380 SharedClock(clock),
381 KnotId::new("did:web:nel.pet").unwrap(),
382 knot_atproto::PlcDirectory::new(Url::parse("https://plc.directory/").unwrap()).unwrap(),
383 ));
384
385 Harness {
386 index,
387 atproto,
388 published,
389 list_calls,
390 _dir: dir,
391 }
392 }
393
394 async fn wait_recognized(roster: &KeyRoster, key: &OfferedKey) {
395 for _ in 0..1000 {
396 if roster.recognizes(key) {
397 return;
398 }
399 tokio::task::yield_now().await;
400 }
401 }
402
403 #[tokio::test]
404 async fn an_acl_write_makes_a_freshly_published_key_recognized_without_waiting_for_the_ttl() {
405 let (line1, offered1) = line_and_offered();
406 let (line2, offered2) = line_and_offered();
407
408 let Harness {
409 index,
410 atproto,
411 published,
412 list_calls,
413 _dir,
414 } = harness(vec![line1]);
415
416 let roster = Arc::new(KeyRoster::new());
417 roster.ensure_fresh(&index, &atproto);
418 wait_recognized(&roster, &offered1).await;
419 assert!(roster.recognizes(&offered1));
420 assert_eq!(list_calls.load(Ordering::SeqCst), 1);
421
422 published.lock().unwrap().push(line2.clone());
423
424 roster.ensure_fresh(&index, &atproto);
425 assert!(
426 !roster.recognizes(&offered2),
427 "stable index and unexpired TTL still serves cached roster, no re-resolution"
428 );
429 assert_eq!(list_calls.load(Ordering::SeqCst), 1);
430
431 index.refresh_members().unwrap();
432 roster.ensure_fresh(&index, &atproto);
433 wait_recognized(&roster, &offered2).await;
434 assert!(
435 roster.recognizes(&offered2),
436 "ACL write bumps generation, so roster revalidates off the auth path"
437 );
438 assert_eq!(
439 list_calls.load(Ordering::SeqCst),
440 2,
441 "exactly one async re-resolution off the auth path"
442 );
443 }
444
445 #[tokio::test]
446 async fn a_miss_against_a_stale_roster_blocks_bounded_to_revalidate_before_rejecting() {
447 let (line1, offered1) = line_and_offered();
448 let (line2, offered2) = line_and_offered();
449 let Harness {
450 index,
451 atproto,
452 published,
453 list_calls,
454 _dir,
455 } = harness(vec![line1]);
456 let roster = Arc::new(KeyRoster::new());
457
458 assert!(
459 roster.recognizes_fresh(&offered1, &index, &atproto).await,
460 "the first handshake blocks on the primed resolve and recognizes the published key"
461 );
462 assert_eq!(list_calls.load(Ordering::SeqCst), 1);
463
464 published.lock().unwrap().push(line2.clone());
465 index.refresh_members().unwrap();
466
467 assert!(
468 roster.recognizes_fresh(&offered2, &index, &atproto).await,
469 "a generation-bumped miss blocks to revalidate and picks up the new key on the first attempt"
470 );
471 assert_eq!(
472 list_calls.load(Ordering::SeqCst),
473 2,
474 "the miss triggers exactly one bounded re-resolution"
475 );
476 }
477
478 #[test]
479 fn staleness_classifies_cold_fresh_and_revalidate() {
480 let roster = KeyRoster::new();
481 assert_eq!(
482 roster.staleness(UnixMicros::new(0), IndexGeneration::new(0)),
483 Staleness::Cold,
484 "with no roster yet the first auth is cold and must revalidate before it can answer a miss"
485 );
486 *roster
487 .freshness
488 .lock()
489 .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(Freshness {
490 due: UnixMicros::new(1_000),
491 generation: IndexGeneration::new(0),
492 });
493 assert_eq!(
494 roster.staleness(UnixMicros::new(500), IndexGeneration::new(0)),
495 Staleness::Fresh
496 );
497 assert_eq!(
498 roster.staleness(UnixMicros::new(500), IndexGeneration::new(1)),
499 Staleness::Revalidate,
500 "an ACL write moves the generation, so the cached roster is stale"
501 );
502 assert_eq!(
503 roster.staleness(UnixMicros::new(2_000), IndexGeneration::new(0)),
504 Staleness::Revalidate,
505 "an expired ttl at the same generation is stale too"
506 );
507 }
508
509 #[test]
510 fn degraded_ttl_backs_off_from_the_short_retry_to_the_fresh_ceiling() {
511 assert_eq!(degraded_ttl(0), Duration::from_secs(5));
512 assert_eq!(degraded_ttl(1), Duration::from_secs(10));
513 assert_eq!(degraded_ttl(2), Duration::from_secs(20));
514 assert_eq!(degraded_ttl(3), Duration::from_secs(40));
515 assert_eq!(degraded_ttl(4), Duration::from_secs(60));
516 assert_eq!(
517 degraded_ttl(50),
518 Duration::from_secs(60),
519 "a persistently unresolvable did clamps the retry to the fresh ttl instead of storming"
520 );
521 }
522}