Skip to main content

keyhog_scanner/engine/
recovery.rs

1use crate::hw_probe::ScanBackend;
2
3const MAX_RECOVERY_REASON_BYTES: usize = 4096;
4const MISSING_RECOVERY_REASON: &str = "backend fault without diagnostic";
5
6/// One exact source-byte interval completed after the selected backend faulted.
7#[derive(Clone, Debug, Eq, PartialEq)]
8pub struct RecoveredInputRange {
9    pub chunk_index: usize,
10    pub byte_start: usize,
11    pub byte_end: usize,
12}
13
14impl RecoveredInputRange {
15    pub fn new(chunk_index: usize, byte_start: usize, byte_end: usize) -> Self {
16        Self {
17            chunk_index,
18            byte_start,
19            byte_end,
20        }
21    }
22
23    #[must_use]
24    pub fn len(&self) -> usize {
25        self.byte_end.saturating_sub(self.byte_start)
26    }
27
28    #[must_use]
29    pub fn is_empty(&self) -> bool {
30        self.byte_start >= self.byte_end
31    }
32}
33
34/// Complete, non-secret receipt for automatic recovery of stable input bytes.
35#[derive(Clone, Debug, Eq, PartialEq)]
36pub struct BackendRecoveryReceipt {
37    pub failed_backend: ScanBackend,
38    pub recovery_backend: ScanBackend,
39    pub ranges: Vec<RecoveredInputRange>,
40    pub reason: String,
41}
42
43impl BackendRecoveryReceipt {
44    pub fn new(
45        failed_backend: ScanBackend,
46        recovery_backend: ScanBackend,
47        ranges: Vec<RecoveredInputRange>,
48        reason: String,
49    ) -> Self {
50        Self {
51            failed_backend,
52            recovery_backend,
53            ranges: canonicalize_ranges(ranges),
54            reason: sanitize_recovery_reason(reason),
55        }
56    }
57    pub(crate) fn phase1_admission(
58        backend: ScanBackend,
59        chunks: &[keyhog_core::Chunk],
60        error: super::Phase1AdmissionPlanIdentityError,
61    ) -> Self {
62        let reason = match error {
63            super::Phase1AdmissionPlanIdentityError::Malformed => "malformed phase-one admission plan identity; discarded the untrusted plan and recomputed exact admission",
64            super::Phase1AdmissionPlanIdentityError::Mismatch => "phase-one admission plan identity mismatch; discarded the untrusted plan and recomputed exact admission",
65        };
66        let ranges = chunks
67            .iter()
68            .enumerate()
69            .filter(|(_, chunk)| !chunk.data.is_empty())
70            .map(|(chunk_index, chunk)| RecoveredInputRange::new(chunk_index, 0, chunk.data.len()))
71            .collect();
72        Self::new(backend, backend, ranges, reason.to_string())
73    }
74
75    #[must_use]
76    pub fn is_phase1_admission_recovery(&self) -> bool {
77        self.reason == "phase-one admission plan identity mismatch; discarded the untrusted plan and recomputed exact admission"
78            || self.reason == "malformed phase-one admission plan identity; discarded the untrusted plan and recomputed exact admission"
79    }
80
81    #[must_use]
82    pub fn recovered_bytes(&self) -> u64 {
83        self.ranges
84            .iter()
85            // LAW10: this is a diagnostic byte counter only; saturation cannot alter recovery candidates or findings.
86            .map(|range| u64::try_from(range.len()).unwrap_or(u64::MAX))
87            .fold(0u64, u64::saturating_add)
88    }
89
90    #[must_use]
91    pub fn recovered_chunks(&self) -> usize {
92        let mut previous = None;
93        self.ranges
94            .iter()
95            .filter(|range| {
96                let distinct = previous != Some(range.chunk_index);
97                previous = Some(range.chunk_index);
98                distinct
99            })
100            .count()
101    }
102}
103
104fn sanitize_recovery_reason(reason: String) -> String {
105    let mut sanitized = String::with_capacity(reason.len().min(MAX_RECOVERY_REASON_BYTES));
106    for ch in reason.chars() {
107        let ch = if ch.is_control() { '\u{fffd}' } else { ch };
108        if sanitized.len().saturating_add(ch.len_utf8()) > MAX_RECOVERY_REASON_BYTES {
109            break;
110        }
111        sanitized.push(ch);
112    }
113    if sanitized.is_empty() {
114        MISSING_RECOVERY_REASON.to_string()
115    } else {
116        sanitized
117    }
118}
119
120/// Result of one fallible coalesced dispatch, including any completed recovery.
121pub struct CoalescedScanOutcome {
122    pub matches: Vec<Vec<keyhog_core::RawMatch>>,
123    pub recovery: Option<BackendRecoveryReceipt>,
124    /// GPU MoE recoveries emitted by this exact dispatch.
125    pub gpu_recovery_receipts: u64,
126}
127
128// Tests live in `tests/unit/engine_recovery.rs` (KH-1308).
129
130pub(crate) fn canonicalize_ranges(
131    mut ranges: Vec<RecoveredInputRange>,
132) -> Vec<RecoveredInputRange> {
133    ranges.retain(|range| !range.is_empty());
134    ranges.sort_unstable_by_key(|range| (range.chunk_index, range.byte_start, range.byte_end));
135    let mut write = 0;
136    for read in 0..ranges.len() {
137        if write != 0
138            && ranges[write - 1].chunk_index == ranges[read].chunk_index
139            && ranges[read].byte_start <= ranges[write - 1].byte_end
140        {
141            ranges[write - 1].byte_end = ranges[write - 1].byte_end.max(ranges[read].byte_end);
142            continue;
143        }
144        ranges.swap(write, read);
145        write += 1;
146    }
147    ranges.truncate(write);
148    ranges
149}