pq_map/discover/
probe_cache.rs

1use std::collections::BTreeMap;
2use std::fs;
3use std::path::PathBuf;
4
5use chrono::{DateTime, Utc};
6use serde::{Deserialize, Serialize};
7
8use crate::model::{Endpoint, PortScan};
9
10use super::cache;
11
12pub const SNAPSHOT_VERSION: u32 = 2;
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct Snapshot {
16    pub version: u32,
17    pub fetched_at: DateTime<Utc>,
18    pub probes: BTreeMap<String, Endpoint>,
19    #[serde(default)]
20    pub ports: BTreeMap<String, BTreeMap<u16, PortScan>>,
21}
22
23fn cache_dir() -> PathBuf {
24    if let Ok(dir) = std::env::var("PROBE_CACHE_DIR") {
25        return PathBuf::from(dir);
26    }
27    let base = std::env::var("XDG_CACHE_HOME")
28        .ok()
29        .filter(|s| !s.is_empty())
30        .map(PathBuf::from)
31        .or_else(|| {
32            std::env::var("HOME")
33                .ok()
34                .map(|h| PathBuf::from(h).join(".cache"))
35        });
36    base.unwrap_or_else(std::env::temp_dir)
37        .join("pq-map")
38        .join("probe")
39}
40
41pub fn load(domain: &str) -> Option<Snapshot> {
42    cache::migrate_legacy_file(&cache_dir(), domain);
43    let file = cache::latest_file(&cache_dir(), domain)?;
44    let body = fs::read_to_string(&file).ok()?;
45    let snap: Snapshot = serde_json::from_str(&body).ok()?;
46    if snap.version != SNAPSHOT_VERSION {
47        return None;
48    }
49    match cache::age(&file) {
50        Some(age) => eprintln!(
51            "probes: loading snapshot {} ({} min old)",
52            file.display(),
53            age / 60
54        ),
55        None => eprintln!("probes: loading snapshot {}", file.display()),
56    }
57    Some(snap)
58}
59
60pub fn store(
61    domain: &str,
62    probes: &BTreeMap<String, Endpoint>,
63    ports: &BTreeMap<String, BTreeMap<u16, PortScan>>,
64) {
65    let snapshot = Snapshot {
66        version: SNAPSHOT_VERSION,
67        fetched_at: Utc::now(),
68        probes: probes.clone(),
69        ports: ports.clone(),
70    };
71    let Ok(body) = serde_json::to_string_pretty(&snapshot) else {
72        eprintln!("probes: caching failed: serialization error");
73        return;
74    };
75    cache::migrate_legacy_file(&cache_dir(), domain);
76    match cache::domain_store(&cache_dir(), domain, &body) {
77        Ok(path) => eprintln!("probes: snapshot cached: {}", path.display()),
78        Err(e) => eprintln!("probes: caching failed: {e}"),
79    }
80}
81
82pub fn is_fresh(endpoint: &Endpoint, ttl: u64, now: DateTime<Utc>) -> bool {
83    endpoint
84        .observed_at
85        .map(|t| ((now - t).num_seconds().max(0) as u64) <= ttl)
86        .unwrap_or(false)
87}
88
89pub fn scan_fresh(scan: &PortScan, ttl: u64, now: DateTime<Utc>) -> bool {
90    is_fresh_time(scan.observed_at, ttl, now)
91}
92
93fn is_fresh_time(t: Option<DateTime<Utc>>, ttl: u64, now: DateTime<Utc>) -> bool {
94    t.map(|t| ((now - t).num_seconds().max(0) as u64) <= ttl)
95        .unwrap_or(false)
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101    use crate::model::ProbeStatus;
102
103    fn timeout_endpoint() -> Endpoint {
104        let mut ep = Endpoint::failed(ProbeStatus::Timeout { secs: 8 });
105        ep.observed_at = Some(Utc::now());
106        ep
107    }
108
109    #[test]
110    fn snapshot_roundtrip_keeps_failure_endpoints() {
111        let dir = std::env::temp_dir().join("pq-scope-probe-cache-test");
112        fs::create_dir_all(&dir).unwrap();
113        unsafe { std::env::set_var("PROBE_CACHE_DIR", &dir) };
114
115        let mut probes = BTreeMap::new();
116        probes.insert("vpn.example.com".to_string(), timeout_endpoint());
117
118        let mut ports = BTreeMap::new();
119        let mut scans = BTreeMap::new();
120        scans.insert(
121            465u16,
122            crate::model::PortScan {
123                port: 465,
124                state: crate::model::ScanState::Open,
125                served: None,
126                negotiated: None,
127                observed_at: Some(Utc::now()),
128            },
129        );
130        ports.insert("mx1.example.com".to_string(), scans);
131
132        store("roundtrip.example.com", &probes, &ports);
133        let loaded = load("roundtrip.example.com").expect("snapshot loads");
134        let ep = loaded.probes.get("vpn.example.com").unwrap();
135        assert_eq!(ep.status, ProbeStatus::Timeout { secs: 8 });
136        assert!(is_fresh(ep, 86_400, Utc::now()));
137        let scans = loaded.ports.get("mx1.example.com").unwrap();
138        let scan = scans.get(&465).unwrap();
139        assert_eq!(scan.state, crate::model::ScanState::Open);
140        assert!(scan_fresh(scan, 86_400, Utc::now()));
141
142        fs::remove_dir_all(&dir).unwrap();
143        unsafe { std::env::remove_var("PROBE_CACHE_DIR") };
144    }
145
146    #[test]
147    fn is_fresh_respects_ttl_and_missing_time() {
148        let now = Utc::now();
149        let mut ep = timeout_endpoint();
150        ep.observed_at = Some(now - chrono::Duration::hours(23));
151        assert!(is_fresh(&ep, 86_400, now));
152
153        ep.observed_at = Some(now - chrono::Duration::hours(25));
154        assert!(!is_fresh(&ep, 86_400, now));
155
156        ep.observed_at = None;
157        assert!(
158            !is_fresh(&ep, 86_400, now),
159            "synthesized endpoints never reusable"
160        );
161    }
162}
pq_map/discover/probe_cache.rs

pq_map/discover/
probe_cache.rs
