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
41fn cache_file(domain: &str) -> PathBuf {
42    cache_dir().join(format!("{domain}.json"))
43}
44
45pub fn load(domain: &str) -> Option<Snapshot> {
46    let body = fs::read_to_string(cache_file(domain)).ok()?;
47    let snap: Snapshot = serde_json::from_str(&body).ok()?;
48    (snap.version == SNAPSHOT_VERSION).then_some(snap)
49}
50
51pub fn store(
52    domain: &str,
53    probes: &BTreeMap<String, Endpoint>,
54    ports: &BTreeMap<String, BTreeMap<u16, PortScan>>,
55) {
56    let snapshot = Snapshot {
57        version: SNAPSHOT_VERSION,
58        fetched_at: Utc::now(),
59        probes: probes.clone(),
60        ports: ports.clone(),
61    };
62    let Ok(body) = serde_json::to_string_pretty(&snapshot) else {
63        eprintln!("probes: caching failed: serialization error");
64        return;
65    };
66    let file = cache_file(domain);
67    if let Some(dir) = file.parent() {
68        if let Err(e) = fs::create_dir_all(dir) {
69            eprintln!("probes: caching failed: {e}");
70            return;
71        }
72    }
73    match cache::store(&file, &body) {
74        Ok(()) => eprintln!("probes: snapshot cached: {}", file.display()),
75        Err(e) => eprintln!("probes: caching failed: {e}"),
76    }
77}
78
79pub fn is_fresh(endpoint: &Endpoint, ttl: u64, now: DateTime<Utc>) -> bool {
80    endpoint
81        .observed_at
82        .map(|t| ((now - t).num_seconds().max(0) as u64) <= ttl)
83        .unwrap_or(false)
84}
85
86pub fn scan_fresh(scan: &PortScan, ttl: u64, now: DateTime<Utc>) -> bool {
87    is_fresh_time(scan.observed_at, ttl, now)
88}
89
90fn is_fresh_time(t: Option<DateTime<Utc>>, ttl: u64, now: DateTime<Utc>) -> bool {
91    t.map(|t| ((now - t).num_seconds().max(0) as u64) <= ttl)
92        .unwrap_or(false)
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use crate::model::ProbeStatus;
99
100    fn timeout_endpoint() -> Endpoint {
101        let mut ep = Endpoint::failed(ProbeStatus::Timeout { secs: 8 });
102        ep.observed_at = Some(Utc::now());
103        ep
104    }
105
106    #[test]
107    fn snapshot_roundtrip_keeps_failure_endpoints() {
108        let dir = std::env::temp_dir().join("pq-scope-probe-cache-test");
109        fs::create_dir_all(&dir).unwrap();
110        unsafe { std::env::set_var("PROBE_CACHE_DIR", &dir) };
111
112        let mut probes = BTreeMap::new();
113        probes.insert("vpn.nbp.pl".to_string(), timeout_endpoint());
114
115        let mut ports = BTreeMap::new();
116        let mut scans = BTreeMap::new();
117        scans.insert(
118            465u16,
119            crate::model::PortScan {
120                port: 465,
121                state: crate::model::ScanState::Open,
122                served: None,
123                negotiated: None,
124                observed_at: Some(Utc::now()),
125            },
126        );
127        ports.insert("mx1.nbp.pl".to_string(), scans);
128
129        store("roundtrip.nbp.pl", &probes, &ports);
130        let loaded = load("roundtrip.nbp.pl").expect("snapshot loads");
131        let ep = loaded.probes.get("vpn.nbp.pl").unwrap();
132        assert_eq!(ep.status, ProbeStatus::Timeout { secs: 8 });
133        assert!(is_fresh(ep, 86_400, Utc::now()));
134        let scans = loaded.ports.get("mx1.nbp.pl").unwrap();
135        let scan = scans.get(&465).unwrap();
136        assert_eq!(scan.state, crate::model::ScanState::Open);
137        assert!(scan_fresh(scan, 86_400, Utc::now()));
138
139        fs::remove_dir_all(&dir).unwrap();
140        unsafe { std::env::remove_var("PROBE_CACHE_DIR") };
141    }
142
143    #[test]
144    fn is_fresh_respects_ttl_and_missing_time() {
145        let now = Utc::now();
146        let mut ep = timeout_endpoint();
147        ep.observed_at = Some(now - chrono::Duration::hours(23));
148        assert!(is_fresh(&ep, 86_400, now));
149
150        ep.observed_at = Some(now - chrono::Duration::hours(25));
151        assert!(!is_fresh(&ep, 86_400, now));
152
153        ep.observed_at = None;
154        assert!(
155            !is_fresh(&ep, 86_400, now),
156            "synthesized endpoints never reusable"
157        );
158    }
159}
pq_map/discover/probe_cache.rs

pq_map/discover/
probe_cache.rs
