This repository has no description
0

Configure Feed

Select the types of activity you want to include in your feed.

core / knot2 / crates / knot-ssh / src / roster.rs
19 kB 522 lines
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}