1use {
11 super::handle::AuditReceiver,
12 crate::{
13 chain::ChainState,
14 checkpoint::Checkpoint,
15 event::AuditEvent,
16 ids::{EpochMs, Seq},
17 ocsf::ProductInfo,
18 signing::FileEd25519Signer,
19 sink::{AuditError, AuditRecord, AuditSink},
20 spool::{PendingLosses, Spool},
21 stats::AuditStats,
22 },
23 krabka_units::prelude::{Time, TimeExt as _},
24 qubit_clock::Timer,
25 std::sync::Arc,
26 tracing::field::Empty,
27};
28
29pub struct AuditWriterParams {
31 pub sink: Arc<dyn AuditSink>,
32 pub product: ProductInfo,
33 pub signer: Option<Arc<FileEd25519Signer>>,
34 pub checkpoint_every_n: u64,
38 pub checkpoint_every: Time,
39 pub chain: ChainState,
41 pub spool: Option<Spool>,
43 pub stats: Arc<AuditStats>,
44 pub replay_every: Time,
46 pub timer: Arc<dyn Timer>,
51}
52
53pub struct AuditWriter {
58 rx: AuditReceiver,
59 pub(super) sink: Arc<dyn AuditSink>,
60 product: ProductInfo,
61 chain: ChainState,
62 signer: Option<Arc<FileEd25519Signer>>,
63 checkpoint_every_n: u64,
64 checkpoint_every: Time,
65 since_checkpoint: u64,
66 pub(super) spool: Option<Spool>,
67 pub(super) spooling: bool,
68 pub(super) stats: Arc<AuditStats>,
69 replay_every: Time,
70 timer: Arc<dyn Timer>,
71 pending_losses: Arc<PendingLosses>,
72}
73
74impl AuditWriter {
75 #[must_use]
76 pub fn new(rx: AuditReceiver, params: AuditWriterParams) -> Self {
77 let spooling = params.spool.as_ref().is_some_and(|s| !s.is_empty());
78 let pending_losses = rx.pending_losses();
79 if let Some(spool) = ¶ms.spool {
80 params.stats.set_depth(spool.count(), spool.size());
81 }
82 Self {
83 rx,
84 sink: params.sink,
85 product: params.product,
86 chain: params.chain,
87 signer: params.signer,
88 checkpoint_every_n: params.checkpoint_every_n,
89 checkpoint_every: params.checkpoint_every,
90 since_checkpoint: 0,
91 spool: params.spool,
92 spooling,
93 stats: params.stats,
94 replay_every: params.replay_every,
95 timer: params.timer,
96 pending_losses,
97 }
98 }
99
100 #[tracing::instrument(
104 level = "info",
105 skip_all,
106 fields(checkpoint_every_n = self.checkpoint_every_n, spooling = self.spooling)
107 )]
108 pub async fn run(mut self) {
109 let timer = Arc::clone(&self.timer);
125 let Some(mut ckpt) = arm(&*timer, self.checkpoint_every.to_std(), CHECKPOINT_TASK) else {
126 return self.finish().await;
127 };
128 let Some(mut replay) = arm(&*timer, self.replay_every.to_std(), REPLAY_TASK) else {
129 return self.finish().await;
130 };
131
132 loop {
133 tokio::select! {
134 maybe = self.rx.recv_message() => {
135 match maybe {
136 Some(message) => {
137 let durable = message.acknowledgement.is_some();
138 let result = self.write_chained(&message.event, durable).await;
139 let fatal = matches!(&result, Err(AuditError::Indeterminate(_)));
140 if result.is_err() {
141 self.stats.inc_dropped();
142 if !durable {
143 self.pending_losses.add(1);
144 if let Err(error) = self.pending_losses.persist() {
145 tracing::error!(%error, "failed to persist pending audit loss count");
146 }
147 }
148 }
149 if let Some(acknowledgement) = message.acknowledgement {
150 let _ = acknowledgement.send(result);
151 }
152 if fatal {
153 tracing::error!("audit writer stopped after indeterminate durable write");
154 return;
155 }
156 if self.checkpoint_every_n > 0
157 && self.since_checkpoint >= self.checkpoint_every_n
158 {
159 self.emit_checkpoint().await;
160 }
161 }
162 None => break,
163 }
164 }
165 outcome = &mut ckpt => {
166 if !fired(outcome, CHECKPOINT_TASK) {
167 break;
168 }
169 if self.since_checkpoint > 0 {
170 self.emit_checkpoint().await;
171 }
172 let Some(next) =
173 arm(&*timer, self.checkpoint_every.to_std(), CHECKPOINT_TASK)
174 else {
175 break;
176 };
177 ckpt = next;
178 }
179 outcome = &mut replay => {
180 if !fired(outcome, REPLAY_TASK) {
181 break;
182 }
183 if self.spooling
184 && let Err(error) = self.try_replay().await
185 {
186 tracing::error!(%error, "audit writer stopped after indeterminate replay");
187 return;
188 }
189 if !self.spooling {
190 let _ = self.write_pending_loss_marker(false).await;
191 }
192 let Some(next) = arm(&*timer, self.replay_every.to_std(), REPLAY_TASK) else {
193 break;
194 };
195 replay = next;
196 }
197 }
198 }
199 self.finish().await;
200 }
201
202 async fn finish(mut self) {
209 let _ = self.write_pending_loss_marker(true).await;
210 if self.since_checkpoint > 0 {
211 self.emit_checkpoint().await;
212 }
213 }
214
215 #[tracing::instrument(
216 level = "debug",
217 skip_all,
218 fields(class = ?event.class(), seq = Empty)
219 )]
220 async fn write_chained(&mut self, event: &AuditEvent, durable: bool) -> Result<(), AuditError> {
221 self.write_pending_loss_marker(durable).await?;
222 let mut record = AuditRecord::from_event(event, &self.product);
223 let seq = self.chain.next_seq();
224 let prev = self.chain.head();
225 tracing::Span::current().record("seq", seq);
226 record.push_chain_headers(seq, &prev);
227 self.write_or_spool(&record, durable).await?;
228 let _ = self.chain.extend(&record.value);
229 self.since_checkpoint += 1;
230 Ok(())
231 }
232
233 async fn write_pending_loss_marker(&mut self, durable: bool) -> Result<(), AuditError> {
235 if let Some(spool) = &mut self.spool {
236 let pending_losses = Arc::clone(&self.pending_losses);
237 let marker = pending_losses.persist_with(|batch| {
238 let mut marker = AuditRecord::records_lost(batch.count, batch.generation);
239 marker.push_chain_headers(self.chain.next_seq(), &self.chain.head());
240 spool.append_loss_marker(&marker)?;
241 Ok(marker)
242 })?;
243 if let Some(marker) = marker {
244 self.stats.inc_spooled();
245 self.stats.set_depth(spool.count(), spool.size());
246 let _ = self.chain.extend(&marker.value);
247 self.since_checkpoint += 1;
248 self.spooling = true;
249 }
250 return Ok(());
251 }
252
253 let Some(batch) = self.pending_losses.snapshot() else {
254 return Ok(());
255 };
256 let mut marker = AuditRecord::records_lost(batch.count, batch.generation);
257 let seq = self.chain.next_seq();
258 let prev = self.chain.head();
259 marker.push_chain_headers(seq, &prev);
260 self.write_or_spool(&marker, durable).await?;
261 let _ = self.chain.extend(&marker.value);
262 self.since_checkpoint += 1;
263 self.pending_losses.commit(batch);
264 Ok(())
265 }
266
267 #[tracing::instrument(
268 level = "debug",
269 skip_all,
270 fields(since_checkpoint = self.since_checkpoint, seq_high = Empty)
271 )]
272 async fn emit_checkpoint(&mut self) {
273 let Some(signer) = self.signer.clone() else {
274 self.since_checkpoint = 0;
275 return;
276 };
277 let seq_high = Seq(self.chain.next_seq().saturating_sub(1));
278 tracing::Span::current().record("seq_high", seq_high.0);
279 let head = self.chain.head();
280 let cp = Checkpoint::signed(signer.as_ref(), seq_high, &head, EpochMs(now_ms()));
281 let record = cp.to_record();
282 match self.write_or_spool(&record, false).await {
283 Ok(()) => self.since_checkpoint = 0,
284 Err(_) => self.stats.inc_dropped(),
285 }
286 }
287}
288
289const CHECKPOINT_TASK: &str = "audit checkpoint";
291
292const REPLAY_TASK: &str = "audit replay";
294
295krabka_macros::timer_hooks! {
296 arm("could not arm the audit timer; stopping the writer");
313
314 fired("the armed audit timer failed; stopping the writer");
320}
321
322krabka_macros::epoch_millis_fn!(
323#[cfg_attr(test, mutants::skip)]
326fn now_ms,
327i64::MAX
328);
329
330#[cfg(test)]
331mod tests {
332 use {
333 super::*,
334 crate::{
335 NodeId,
336 chain::{GENESIS_HEAD, chain_hash, to_hex},
337 event::AuditEventClass,
338 log::{
339 AuditLog,
340 test_support::{
341 AuditEventCount, CheckpointFrequency, CheckpointSetup,
342 FailableSink as TestFailableSink, ROOMY_CAP, SinkStats, await_until,
343 failed_sink_stats, finish_lifecycle_writer, finish_writer, header, life,
344 params, quiet_params, roomy_spool, spawn_writer, test_signer,
345 },
346 },
347 sink::{AuditSink, MemorySink},
348 spool::PendingLosses,
349 test_support::{TempSpool, TestSigner},
350 },
351 assert2::check,
352 qubit_clock::{
353 MonotonicClock, MonotonicInstant, StdMonotonicClock, TimeError, TimerFuture,
354 },
355 std::path::Path,
356 tokio::task::JoinHandle,
357 };
358
359 #[derive(
360 Debug,
361 Clone,
362 Copy,
363 PartialEq,
364 Eq,
365 derive_more::Display,
366 derive_more::From,
367 derive_more::Into,
368 )]
369 struct QueueCapacity(usize);
370
371 struct MemoryWriterSetup {
372 capacity: QueueCapacity,
373 checkpoints: CheckpointSetup,
374 }
375
376 impl Default for MemoryWriterSetup {
377 fn default() -> Self {
378 Self {
379 capacity: QueueCapacity(16),
380 checkpoints: CheckpointSetup {
381 frequency: CheckpointFrequency::Every(AuditEventCount(1_000_000)),
382 ..Default::default()
383 },
384 }
385 }
386 }
387
388 struct MemoryWriter {
389 log: Arc<AuditLog>,
390 sink: Arc<MemorySink>,
391 handle: JoinHandle<()>,
392 }
393
394 #[bon::builder]
395 fn memory_writer(
396 directory: &Path,
397 #[builder(into, default = MemoryWriterSetup::default().capacity)] capacity: QueueCapacity,
398 #[builder(default = MemoryWriterSetup::default().checkpoints)] checkpoints: CheckpointSetup,
399 ) -> MemoryWriter {
400 let setup = MemoryWriterSetup {
401 capacity,
402 checkpoints,
403 };
404 let MemoryWriterSetup {
405 capacity,
406 checkpoints,
407 } = setup;
408 let spool = Spool::open(directory, ROOMY_CAP).unwrap();
409 let (log, receiver) = AuditLog::new(capacity.0);
410 let sink = Arc::new(MemorySink::default());
411 let handle = spawn_writer(
412 receiver,
413 quiet_params()
414 .sink(Arc::<MemorySink>::clone(&sink))
415 .spool(spool)
416 .stats(checkpoints.stats)
417 .maybe_signer(checkpoints.signer)
418 .frequency(checkpoints.frequency)
419 .call(),
420 );
421 MemoryWriter { log, sink, handle }
422 }
423
424 struct SignedWriterFixture {
425 public_key: Vec<u8>,
426 directory: tempfile::TempDir,
427 writer: MemoryWriter,
428 }
429
430 #[derive(Clone, Copy, krabka_macros::FieldDefaults)]
431 struct SignedWriterSetup {
432 #[default(QueueCapacity(16))]
433 capacity: QueueCapacity,
434 #[default(CheckpointFrequency::Every(AuditEventCount(1_000_000)))]
435 frequency: CheckpointFrequency,
436 }
437
438 #[bon::builder]
439 fn signed_writer(
440 #[builder(into, default = QueueCapacity (16))] capacity: QueueCapacity,
441 #[builder(default = CheckpointFrequency::Every(AuditEventCount(1_000_000)))]
442 frequency: CheckpointFrequency,
443 ) -> SignedWriterFixture {
444 let setup = SignedWriterSetup {
445 capacity,
446 frequency,
447 };
448 let TestSigner { signer, public_key } = test_signer();
449 let directory = tempfile::tempdir().unwrap();
450 let writer = memory_writer()
451 .directory(directory.path())
452 .capacity(setup.capacity)
453 .checkpoints(CheckpointSetup {
454 signer: Some(signer),
455 frequency: setup.frequency,
456 ..Default::default()
457 })
458 .call();
459 SignedWriterFixture {
460 public_key,
461 directory,
462 writer,
463 }
464 }
465
466 struct PendingLossWriter {
467 log: Arc<AuditLog>,
468 receiver: AuditReceiver,
469 directory: tempfile::TempDir,
470 _losses: Arc<PendingLosses>,
472 params: AuditWriterParams,
473 }
474
475 fn pending_loss_writer<T: AuditSink + 'static>(
476 sink: &Arc<T>,
477 stats: impl FnOnce() -> Arc<AuditStats>,
478 signer: Arc<FileEd25519Signer>,
479 ) -> PendingLossWriter {
480 let (log, receiver) = AuditLog::new(16);
481 let losses = receiver.pending_losses();
482 losses.add(1);
483 let directory = tempfile::tempdir().unwrap();
484 let spool = Spool::open(directory.path(), ROOMY_CAP).unwrap();
485 let mut params = params(Arc::<T>::clone(sink), spool, stats());
486 params.signer = Some(signer);
487 params.checkpoint_every_n = 2;
488 PendingLossWriter {
489 log,
490 receiver,
491 directory,
492 _losses: losses,
493 params,
494 }
495 }
496
497 #[tokio::test]
498 async fn emitted_events_reach_the_sink_in_order() {
499 let dir = tempfile::tempdir().unwrap();
500 let MemoryWriter { log, sink, handle } = memory_writer().directory(dir.path()).call();
501
502 finish_lifecycle_writer(log, handle, &[NodeId(1), NodeId(2), NodeId(3)]).await;
503
504 let recs = sink.records();
505 check!((recs.len(), recs[0].class) == (3, AuditEventClass::ApplicationLifecycle));
506 let v0: serde_json::Value = serde_json::from_slice(&recs[0].value).unwrap();
508 check!(v0["device"]["uid"] == "1");
509 }
510
511 #[tokio::test]
512 async fn chained_records_carry_seq_and_prev_hash() {
513 let dir = tempfile::tempdir().unwrap();
514 let MemoryWriter {
516 log,
517 sink,
518 handle: h,
519 } = memory_writer().directory(dir.path()).call();
520 finish_lifecycle_writer(log, h, &[NodeId(1), NodeId(2)]).await;
521
522 let recs = sink.records();
523 check!(recs.len() == 2); let seq0 = header(&recs[0], "seq");
526 let seq1 = header(&recs[1], "seq");
527 check!(
528 (seq0, seq1, header(&recs[0], "prev_hash"))
529 == (
530 Some("0".to_owned()),
531 Some("1".to_owned()),
532 Some("0".repeat(64)),
533 )
534 );
535 let expect = to_hex(&chain_hash(&GENESIS_HEAD, 0, &recs[0].value));
537 check!(header(&recs[1], "prev_hash") == Some(expect));
538 }
539
540 #[tokio::test]
541 async fn checkpoints_emitted_by_count_and_verify_against_recomputed_head() {
542 let SignedWriterFixture {
544 public_key: pubkey,
545 directory: _dir,
546 writer:
547 MemoryWriter {
548 log,
549 sink,
550 handle: h,
551 },
552 } = signed_writer()
553 .capacity(64)
554 .frequency(CheckpointFrequency::Every(AuditEventCount(2)))
555 .call();
556 for i in 0..4 {
557 log.emit(life(NodeId(i)));
558 }
559 finish_writer(log, h).await; let recs = sink.records();
562 let checkpoints = recs
564 .iter()
565 .filter(|r| r.class == AuditEventClass::Checkpoint)
566 .count();
567 check!(checkpoints == 2);
568
569 let mut head = GENESIS_HEAD;
571 let mut seq = 0u64;
572 for r in &recs {
573 if r.class == AuditEventClass::Checkpoint {
574 let v: serde_json::Value = serde_json::from_slice(&r.value).unwrap();
575 let cp = Checkpoint::from_value(&v).expect("cp");
576 check!(
577 (cp.verify(&pubkey), cp.chain_head, cp.seq_high) == (true, head, Seq(seq - 1))
578 );
579 } else {
580 head = chain_hash(&head, seq, &r.value);
581 seq += 1;
582 }
583 }
584 }
585
586 #[tokio::test]
587 async fn shutdown_emits_final_checkpoint_for_pending_tail() {
588 let SignedWriterFixture {
590 public_key: pubkey,
591 directory: _dir,
592 writer:
593 MemoryWriter {
594 log,
595 sink,
596 handle: h,
597 },
598 } = signed_writer().call();
599 finish_lifecycle_writer(log, h, &[NodeId(1), NodeId(2), NodeId(3)]).await;
600
601 let recs = sink.records();
602 let cps: Vec<_> = recs
603 .iter()
604 .filter(|r| r.class == AuditEventClass::Checkpoint)
605 .collect();
606 check!(cps.len() == 1); let v: serde_json::Value = serde_json::from_slice(&cps[0].value).unwrap();
608 let cp = Checkpoint::from_value(&v).unwrap();
609 check!((cp.verify(&pubkey), cp.seq_high) == (true, Seq(2)));
610 }
611
612 struct DeadTimer(StdMonotonicClock);
614
615 impl Timer for DeadTimer {
616 fn clock(&self) -> &dyn MonotonicClock { &self.0 }
617
618 fn at(&self, _deadline: MonotonicInstant) -> Result<TimerFuture, TimeError> {
619 Err(TimeError::InstantOverflow)
620 }
621 }
622
623 #[tokio::test]
624 async fn writer_stops_when_a_ticker_cannot_be_armed() {
625 let TempSpool { dir: _dir, spool } = roomy_spool();
626 let sink = Arc::new(MemorySink::default());
627 let (log, rx) = AuditLog::new(16);
628 let mut params = params(
629 Arc::<MemorySink>::clone(&sink),
630 spool,
631 Arc::new(AuditStats::new()),
632 );
633 params.timer = Arc::new(DeadTimer(StdMonotonicClock::new()));
634 let handle = spawn_writer(rx, params);
635
636 log.emit(life(NodeId(1)));
640 handle.await.unwrap();
641 check!(sink.records().is_empty());
642 drop(log);
643 }
644
645 #[tokio::test]
646 async fn failed_checkpoint_increments_dropped() {
647 let TempSpool { dir: _dir, spool } = roomy_spool();
648 let sink = Arc::new(TestFailableSink::default());
649 sink.allow_n(1);
650 let stats = Arc::new(AuditStats::new());
651 let (log, rx) = AuditLog::new(4);
652 let mut params = params(sink, spool, Arc::clone(&stats));
653 params.signer = Some(test_signer().signer);
654 params.checkpoint_every_n = 1;
655 params.spool = None;
656 let handle = spawn_writer(rx, params);
657
658 log.emit(life(NodeId(1)));
659 finish_writer(log, handle).await;
660
661 check!(stats.dropped() >= 1);
662 }
663
664 #[test]
665 fn fired_reports_timer_outcome() {
666 check!(fired(Ok(()), "test"));
667 check!(!fired(Err(TimeError::InstantOverflow), "test"));
668 }
669
670 #[tokio::test]
671 async fn pending_loss_marker_advances_since_checkpoint_without_spool() {
672 let TestSigner { signer, .. } = test_signer();
673 let sink = Arc::new(MemorySink::default());
674 let PendingLossWriter {
675 log,
676 receiver: rx,
677 _losses,
678 directory: _dir,
679 mut params,
680 } = pending_loss_writer(&sink, || Arc::new(AuditStats::new()), signer);
681 params.spool = None;
682 let handle = spawn_writer(rx, params);
683
684 log.emit(life(NodeId(1)));
685 await_until("loss marker, event, and checkpoint reached sink", || {
686 sink.records()
687 .iter()
688 .any(|r| r.class == AuditEventClass::Checkpoint)
689 })
690 .await;
691 finish_writer(log, handle).await;
692
693 let recs = sink.records();
694 check!(
695 recs.iter()
696 .filter(|r| r.class == AuditEventClass::RecordsLost)
697 .count()
698 == 1
699 );
700 check!(
701 recs.iter()
702 .filter(|r| r.class == AuditEventClass::ApplicationLifecycle)
703 .count()
704 == 1
705 );
706 check!(
707 recs.iter()
708 .filter(|r| r.class == AuditEventClass::Checkpoint)
709 .count()
710 >= 1
711 );
712 }
713
714 #[tokio::test]
715 async fn pending_loss_marker_advances_since_checkpoint_with_spool() {
716 let TestSigner { signer, .. } = test_signer();
717 let SinkStats { sink, stats } = failed_sink_stats();
718 let PendingLossWriter {
719 log,
720 receiver: rx,
721 _losses,
722 directory: _dir,
723 params,
724 } = pending_loss_writer(&sink, || Arc::clone(&stats), signer);
725 let handle = spawn_writer(rx, params);
726
727 log.emit(life(NodeId(1)));
728 await_until("loss marker, event, and checkpoint spooled", || {
729 stats.spooled() >= 3
730 })
731 .await;
732 finish_writer(log, handle).await;
733 }
734}