relay_actions/cache/
records.rs

1use anyhow::{Context, Result, bail};
2use serde::{Deserialize, Serialize};
3use std::{
4    collections::HashMap,
5    path::{Path, PathBuf},
6    time::{SystemTime, UNIX_EPOCH},
7};
8
9use crate::cache::AgentsCache;
10
11use super::Cache;
12use relay_lib::prelude::{AgentId, KeyRecord, UserId};
13
14const FILE: &str = "agent_keys.ron";
15
16pub type AgentKeyCache = HashMap<AgentId, CachedRecord>;
17pub type UserKeys = HashMap<UserId, CachedRecord>;
18pub type UserKeyCache = HashMap<AgentId, UserKeys>;
19
20#[derive(Debug, Default)]
21pub struct RecordsCache {
22    pub base: PathBuf,
23    pub agent_keys: AgentKeyCache,
24    pub user_keys: UserKeyCache,
25}
26
27impl Cache for RecordsCache {
28    fn load(base: &Path) -> Result<Self> {
29        let content = std::fs::read_to_string(base.join(FILE)).unwrap_or_default();
30        let agent_keys: AgentKeyCache = ron::from_str(&content).unwrap_or_default();
31
32        Ok(Self {
33            base: base.to_path_buf(),
34            agent_keys,
35            user_keys: HashMap::new(),
36        })
37    }
38
39    fn save(&self, base: &Path) -> Result<()> {
40        let agent_path = base.join(FILE);
41        std::fs::create_dir_all(base)?;
42        std::fs::write(agent_path, ron::to_string(&self.agent_keys)?)?;
43
44        let user_dir = base.join("user_keys");
45        std::fs::create_dir_all(&user_dir)?;
46
47        for (agent_id, users) in &self.user_keys {
48            let path = user_dir.join(format!("{}.json", agent_id));
49            std::fs::write(path, ron::to_string(users)?)?;
50        }
51
52        Ok(())
53    }
54
55    fn cleanup(&mut self) -> Result<()> {
56        self.agent_keys.retain(|_, v| v.is_valid());
57
58        for users in self.user_keys.values_mut() {
59            users.retain(|_, record| record.is_valid());
60        }
61
62        Ok(())
63    }
64}
65
66impl RecordsCache {
67    fn load_user_cache(&mut self, agent: &AgentId) -> Result<()> {
68        if self.user_keys.contains_key(agent) {
69            return Ok(());
70        }
71
72        let path = self.base.join("user_keys").join(format!("{}.ron", agent));
73
74        let users: UserKeys = if let Ok(data) = std::fs::read_to_string(&path) {
75            ron::from_str(&data)?
76        } else {
77            HashMap::new()
78        };
79
80        self.user_keys.insert(agent.clone(), users);
81        Ok(())
82    }
83
84    pub fn agent(&mut self, agents: &mut AgentsCache, agent: &AgentId) -> Result<KeyRecord> {
85        if let Some(entry) = self.agent_keys.get(agent)
86            && entry.is_valid()
87        {
88            return Ok(entry.record.clone());
89        }
90
91        let base_url = agents.url(agent).context("Agent not found in cache")?;
92        let resp = reqwest::blocking::Client::new()
93            .get(base_url.join("/keys/agent")?)
94            .send()
95            .context("Failed to send request to fetch agent record")?;
96
97        if !resp.status().is_success() {
98            let status = resp.status();
99            let body = resp.text().unwrap_or_default();
100            bail!("Failed to fetch agent record: {} - {}", status, body);
101        }
102
103        let record = resp
104            .json::<KeyRecord>()
105            .context("Failed to parse agent record response")?;
106
107        let cached = CachedRecord {
108            record: record.clone(),
109            fetched_at: now_secs(),
110            ttl_secs: 3600,
111        };
112
113        self.agent_keys.insert(agent.clone(), cached);
114
115        Ok(record)
116    }
117
118    pub fn user(
119        &mut self,
120        agents: &mut AgentsCache,
121        agent: &AgentId,
122        user: &UserId,
123    ) -> Result<KeyRecord> {
124        self.load_user_cache(agent)?;
125
126        let users = self.user_keys.get_mut(agent).unwrap();
127
128        if let Some(entry) = users.get(user)
129            && entry.is_valid()
130        {
131            return Ok(entry.record.clone());
132        }
133
134        let base_url = agents.url(agent)?;
135
136        let resp = reqwest::blocking::Client::new()
137            .get(base_url.join(&format!("/keys/user/{}", user))?)
138            .send()
139            .context("Failed to send request to fetch user record")?;
140
141        if !resp.status().is_success() {
142            let status = resp.status();
143            let body = resp.text().unwrap_or_default();
144            bail!("Failed to fetch user record: {} - {}", status, body);
145        }
146
147        let record = resp
148            .json::<KeyRecord>()
149            .context("Failed to parse user record response")?;
150
151        users.insert(
152            user.clone(),
153            CachedRecord {
154                record: record.clone(),
155                fetched_at: now_secs(),
156                ttl_secs: 600,
157            },
158        );
159
160        Ok(record)
161    }
162
163    pub fn invalidate_agent(&mut self, agent: &AgentId) {
164        self.agent_keys.remove(agent);
165        self.user_keys.retain(|a, _| a != agent);
166    }
167
168    pub fn clear(&mut self) {
169        self.agent_keys.clear();
170        self.user_keys.clear();
171    }
172}
173
174#[derive(Debug, Clone, Serialize, Deserialize)]
175pub struct CachedRecord {
176    pub record: KeyRecord,
177    pub fetched_at: u64,
178    pub ttl_secs: u64,
179}
180
181impl CachedRecord {
182    pub fn is_valid(&self) -> bool {
183        let now = now_secs();
184        now < self.fetched_at + self.ttl_secs
185    }
186}
187
188fn now_secs() -> u64 {
189    SystemTime::now()
190        .duration_since(UNIX_EPOCH)
191        .unwrap()
192        .as_secs()
193}
relay_actions/cache/records.rs

relay_actions/cache/
records.rs
