#![allow(missing_docs)]
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use thiserror::Error;
const CALIBRATION_TMP_PREFIX: &str = ".tmp.keyhog-calibration-";
use crate::state_file::{self, CALIBRATION_CACHE_FILE_BYTES};
use crate::STALE_TMP_CUTOFF_SECS;
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct BetaCounters {
pub alpha: u32,
pub beta: u32,
}
impl Default for BetaCounters {
fn default() -> Self {
Self { alpha: 1, beta: 1 }
}
}
impl BetaCounters {
pub fn posterior_mean(&self) -> f64 {
let total = self.alpha as f64 + self.beta as f64;
if total == 0.0 {
0.5
} else {
self.alpha as f64 / total
}
}
pub fn observations(&self) -> u32 {
self.alpha
.saturating_sub(1)
.saturating_add(self.beta.saturating_sub(1))
}
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct OnDisk {
version: u32,
detectors: HashMap<String, BetaCounters>,
}
const SCHEMA_VERSION: u32 = 1;
#[derive(Debug, Error)]
pub enum CalibrationLoadError {
#[error("calibration cache '{}' could not be read: {source}", path.display())]
Read {
path: PathBuf,
source: std::io::Error,
},
#[error("calibration cache '{}' is not valid JSON: {source}", path.display())]
Parse {
path: PathBuf,
source: serde_json::Error,
},
#[error(
"calibration cache '{}' has schema version {found}; expected {expected}",
path.display()
)]
SchemaVersion {
path: PathBuf,
found: u32,
expected: u32,
},
#[error(
"calibration cache '{}' has invalid counters for detector '{}': alpha={alpha}, beta={beta}; counters must be >= 1",
path.display(),
detector_id
)]
InvalidCounters {
path: PathBuf,
detector_id: String,
alpha: u32,
beta: u32,
},
#[error("calibration cache '{}' contains an empty detector id", path.display())]
EmptyDetectorId { path: PathBuf },
#[error(
"calibration cache '{}' exceeds {cap} byte cap; delete the cache file and rerun calibration",
path.display()
)]
TooLarge { path: PathBuf, cap: u64 },
}
#[derive(Debug, Default)]
pub struct Calibration {
inner: RwLock<HashMap<String, BetaCounters>>,
}
impl Calibration {
fn empty() -> Self {
Self::default()
}
pub fn try_load(path: &Path) -> Result<Option<Self>, CalibrationLoadError> {
sweep_stale_calibration_tmp_files(path);
let bytes = match state_file::read_capped(
path,
CALIBRATION_CACHE_FILE_BYTES,
"calibration cache",
) {
Ok(b) => b,
Err(source) if source.kind() == std::io::ErrorKind::NotFound => return Ok(None),
Err(source) if source.kind() == std::io::ErrorKind::InvalidData => {
return Err(CalibrationLoadError::TooLarge {
path: path.to_path_buf(),
cap: CALIBRATION_CACHE_FILE_BYTES,
});
}
Err(source) => {
return Err(CalibrationLoadError::Read {
path: path.to_path_buf(),
source,
});
}
};
let on_disk: OnDisk =
serde_json::from_slice(&bytes).map_err(|source| CalibrationLoadError::Parse {
path: path.to_path_buf(),
source,
})?;
if on_disk.version != SCHEMA_VERSION {
return Err(CalibrationLoadError::SchemaVersion {
path: path.to_path_buf(),
found: on_disk.version,
expected: SCHEMA_VERSION,
});
}
validate_on_disk(path, &on_disk)?;
Ok(Some(Self {
inner: RwLock::new(on_disk.detectors),
}))
}
pub(crate) fn load(path: &Path) -> Self {
match Self::try_load(path) {
Ok(Some(calibration)) => calibration,
Ok(None) => Self::empty(),
Err(error) => {
tracing::warn!(
cache = %path.display(),
error = %error,
"calibration cache could not be loaded; treating as cold start"
);
Self::empty()
}
}
}
pub fn save(&self, path: &Path) -> std::io::Result<()> {
let detectors = self.inner.read().clone();
let on_disk = OnDisk {
version: SCHEMA_VERSION,
detectors,
};
let serialized = serde_json::to_vec_pretty(&on_disk)
.map_err(|e| std::io::Error::other(format!("calibration encode: {e}")))?;
state_file::write_atomically(path, CALIBRATION_TMP_PREFIX, &serialized)
}
pub fn record_outcome(&self, detector_id: &str, true_positive: bool) {
if true_positive {
self.record_true_positive(detector_id);
} else {
self.record_false_positive(detector_id);
}
}
pub(crate) fn record_true_positive(&self, detector_id: &str) {
let mut guard = self.inner.write();
let entry = guard.entry(detector_id.to_string()).or_default();
entry.alpha = entry.alpha.saturating_add(1);
}
pub(crate) fn record_false_positive(&self, detector_id: &str) {
let mut guard = self.inner.write();
let entry = guard.entry(detector_id.to_string()).or_default();
entry.beta = entry.beta.saturating_add(1);
}
pub(crate) fn confidence_multiplier(&self, detector_id: &str) -> f64 {
self.inner
.read()
.get(detector_id)
.copied()
.unwrap_or_default() .posterior_mean()
}
pub fn counters(&self, detector_id: &str) -> BetaCounters {
self.inner
.read()
.get(detector_id)
.copied()
.unwrap_or_default() }
pub fn entries(&self) -> Vec<(String, BetaCounters)> {
let mut out: Vec<_> = self
.inner
.read()
.iter()
.map(|(k, v)| (k.clone(), *v))
.collect();
out.sort_by(|a, b| a.0.cmp(&b.0));
out
}
#[doc(hidden)]
pub(crate) fn test_seed_counters(&self, id: &str, alpha: u32, beta: u32) {
let mut guard = self.inner.write();
let entry = guard.entry(id.to_string()).or_default();
entry.alpha = alpha;
entry.beta = beta;
}
}
pub fn calibration_default_cache_path() -> Option<PathBuf> {
crate::keyhog_cache_root().map(|d| d.join("calibration.json"))
}
fn validate_on_disk(path: &Path, on_disk: &OnDisk) -> Result<(), CalibrationLoadError> {
for (detector_id, counters) in &on_disk.detectors {
if detector_id.trim().is_empty() {
return Err(CalibrationLoadError::EmptyDetectorId {
path: path.to_path_buf(),
});
}
if counters.alpha == 0 || counters.beta == 0 {
return Err(CalibrationLoadError::InvalidCounters {
path: path.to_path_buf(),
detector_id: detector_id.clone(),
alpha: counters.alpha,
beta: counters.beta,
});
}
}
Ok(())
}
fn sweep_stale_calibration_tmp_files(cache_path: &Path) {
let swept = state_file::sweep_stale_tmp_siblings(
cache_path,
&[CALIBRATION_TMP_PREFIX],
STALE_TMP_CUTOFF_SECS,
);
if swept > 0 {
if let Some(parent) = cache_path.parent() {
tracing::debug!(
count = swept,
dir = %parent.display(),
"swept stale calibration cache tmp files left by an interrupted save"
);
}
}
}