platform/aid/
key_cache.rs1use crate::aid::credential::error::AidError;
7use base64::prelude::*;
8use ed25519_dalek::VerifyingKey;
9use sqlx::Row;
10use sqlx::sqlite::{SqliteConnectOptions, SqlitePool, SqlitePoolOptions};
11use std::path::Path;
12use std::str::FromStr;
13use std::sync::{Arc, Mutex};
14use std::time::{Duration, SystemTime, UNIX_EPOCH};
15
16#[derive(Debug, Clone)]
21pub struct KeyCache {
22 pool: SqlitePool,
23 last_cleanup_time: Arc<Mutex<u64>>,
24}
25
26impl KeyCache {
27 pub async fn new<P: AsRef<Path>>(cache_db_file: P) -> Result<Self, AidError> {
29 let database_url = format!("sqlite:{}", cache_db_file.as_ref().display());
30 let options = SqliteConnectOptions::from_str(&database_url)
31 .map_err(|e| AidError::DecodeFailure(format!("invalid database URL: {e}")))?
32 .create_if_missing(true)
33 .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
34 .synchronous(sqlx::sqlite::SqliteSynchronous::Normal)
35 .busy_timeout(Duration::from_secs(5));
36
37 let pool = SqlitePoolOptions::new()
38 .max_connections(5)
39 .connect_with(options)
40 .await
41 .map_err(|e| {
42 AidError::DecodeFailure(format!("failed to open key cache database: {e}"))
43 })?;
44
45 let cache = Self {
46 pool,
47 last_cleanup_time: Arc::new(Mutex::new(0)),
48 };
49
50 cache.init_tables().await?;
51
52 crate::recording::info!(
53 "AIS 公钥缓存初始化完成,数据库:{}",
54 cache_db_file.as_ref().display()
55 );
56 Ok(cache)
57 }
58
59 async fn init_tables(&self) -> Result<(), AidError> {
61 sqlx::query(
63 r#"
64 CREATE TABLE IF NOT EXISTS key_cache (
65 key_id INTEGER PRIMARY KEY,
66 pubkey TEXT NOT NULL,
67 cached_at INTEGER NOT NULL,
68 expires_at INTEGER NOT NULL
69 )
70 "#,
71 )
72 .execute(&self.pool)
73 .await
74 .map_err(|e| AidError::DecodeFailure(format!("failed to create key_cache table: {e}")))?;
75
76 sqlx::query("CREATE INDEX IF NOT EXISTS idx_cache_expires_at ON key_cache(expires_at)")
77 .execute(&self.pool)
78 .await
79 .map_err(|e| {
80 AidError::DecodeFailure(format!("failed to create key_cache index: {e}"))
81 })?;
82
83 crate::recording::debug!("AIS 公钥缓存数据库表已就绪");
84 Ok(())
85 }
86
87 pub async fn get_cached_key(
91 &self,
92 key_id: u32,
93 ) -> Result<Option<(VerifyingKey, u64)>, AidError> {
94 self.maybe_cleanup().await;
95
96 let now = SystemTime::now()
97 .duration_since(UNIX_EPOCH)
98 .unwrap_or_default()
99 .as_secs();
100
101 let result = sqlx::query("SELECT pubkey, expires_at FROM key_cache WHERE key_id = ?1")
102 .bind(key_id as i64)
103 .fetch_optional(&self.pool)
104 .await
105 .map_err(|e| {
106 crate::recording::error!("查询 AIS 公钥缓存失败:key_id={}, {}", key_id, e);
107 AidError::DecodeFailure(format!("key cache query error: {e}"))
108 })?;
109
110 match result {
111 Some(row) => {
112 let pubkey_b64: String = row.try_get("pubkey").map_err(|e| {
113 AidError::DecodeFailure(format!("failed to read pubkey column: {e}"))
114 })?;
115 let expires_at: i64 = row.try_get("expires_at").map_err(|e| {
116 AidError::DecodeFailure(format!("failed to read expires_at column: {e}"))
117 })?;
118
119 if expires_at > 0 && (expires_at as u64) <= now {
121 crate::recording::debug!("AIS 公钥缓存已过期,删除:key_id={}", key_id);
122 let _ = sqlx::query("DELETE FROM key_cache WHERE key_id = ?1")
123 .bind(key_id as i64)
124 .execute(&self.pool)
125 .await;
126 return Ok(None);
127 }
128
129 let pubkey_bytes = BASE64_STANDARD.decode(&pubkey_b64).map_err(|e| {
131 AidError::DecodeFailure(format!("failed to base64 decode pubkey: {e}"))
132 })?;
133
134 let pubkey_array: [u8; 32] = pubkey_bytes.try_into().map_err(|_| {
135 AidError::DecodeFailure(
136 "pubkey 长度无效,期望 32 bytes(Ed25519 VerifyingKey)".to_string(),
137 )
138 })?;
139
140 let verifying_key = VerifyingKey::from_bytes(&pubkey_array).map_err(|e| {
141 AidError::DecodeFailure(format!("invalid Ed25519 verifying key: {e}"))
142 })?;
143
144 crate::recording::debug!("命中 AIS 公钥缓存:key_id={}", key_id);
145 Ok(Some((verifying_key, expires_at as u64)))
146 }
147 None => {
148 crate::recording::debug!("AIS 公钥缓存未命中:key_id={}", key_id);
149 Ok(None)
150 }
151 }
152 }
153
154 pub async fn cache_key(
158 &self,
159 key_id: u32,
160 verifying_key: &VerifyingKey,
161 expires_at: u64,
162 ) -> Result<(), AidError> {
163 let now = SystemTime::now()
164 .duration_since(UNIX_EPOCH)
165 .unwrap_or_default()
166 .as_secs();
167
168 let pubkey_b64 = BASE64_STANDARD.encode(verifying_key.as_bytes());
169
170 sqlx::query(
171 "REPLACE INTO key_cache (key_id, pubkey, cached_at, expires_at) VALUES (?1, ?2, ?3, ?4)",
172 )
173 .bind(key_id as i64)
174 .bind(&pubkey_b64)
175 .bind(now as i64)
176 .bind(expires_at as i64)
177 .execute(&self.pool)
178 .await
179 .map_err(|e| AidError::DecodeFailure(format!("failed to write key cache: {e}")))?;
180
181 crate::recording::debug!(
182 "AIS 公钥已写入缓存:key_id={}, expires_at={}",
183 key_id,
184 expires_at
185 );
186 Ok(())
187 }
188
189 pub async fn cleanup_expired_keys(&self) -> Result<u32, AidError> {
191 let now = SystemTime::now()
192 .duration_since(UNIX_EPOCH)
193 .unwrap_or_default()
194 .as_secs();
195
196 let result = sqlx::query("DELETE FROM key_cache WHERE expires_at > 0 AND expires_at < ?1")
197 .bind(now as i64)
198 .execute(&self.pool)
199 .await
200 .map_err(|e| AidError::DecodeFailure(format!("failed to cleanup expired keys: {e}")))?;
201
202 let deleted = result.rows_affected() as u32;
203 if deleted > 0 {
204 crate::recording::info!("AIS 公钥缓存:已清理 {} 条过期记录", deleted);
205 }
206 Ok(deleted)
207 }
208
209 pub async fn get_cached_key_count(&self) -> Result<u32, AidError> {
211 let row = sqlx::query("SELECT COUNT(*) as count FROM key_cache")
212 .fetch_one(&self.pool)
213 .await
214 .map_err(|e| AidError::DecodeFailure(format!("failed to count cached keys: {e}")))?;
215
216 let count: i64 = row
217 .try_get("count")
218 .map_err(|e| AidError::DecodeFailure(format!("failed to read count: {e}")))?;
219
220 Ok(count as u32)
221 }
222
223 async fn maybe_cleanup(&self) {
225 let now = SystemTime::now()
226 .duration_since(UNIX_EPOCH)
227 .unwrap_or_default()
228 .as_secs();
229
230 let should_cleanup = {
231 let last = self.last_cleanup_time.lock().unwrap();
232 now.saturating_sub(*last) >= 3600
233 };
234
235 if should_cleanup {
236 crate::recording::debug!("触发 AIS 公钥缓存定时清理");
237 match self.cleanup_expired_keys().await {
238 Ok(n) if n > 0 => crate::recording::info!("AIS 公钥缓存清理完成,删除 {} 条", n),
239 Ok(_) => {}
240 Err(e) => crate::recording::error!("AIS 公钥缓存清理失败:{}", e),
241 }
242 let mut last = self.last_cleanup_time.lock().unwrap();
243 *last = now;
244 }
245 }
246}
247
248#[cfg(test)]
249mod tests {
250 use super::*;
251 use ed25519_dalek::SigningKey;
252 use rand::rngs::OsRng;
253 use tempfile::tempdir;
254
255 fn generate_test_key() -> VerifyingKey {
256 let signing_key = SigningKey::generate(&mut OsRng);
257 signing_key.verifying_key()
258 }
259
260 #[tokio::test]
261 async fn test_cache_creation() {
262 let temp_dir = tempdir().unwrap();
263 let cache = KeyCache::new(temp_dir.path().join("test.db")).await;
264 assert!(cache.is_ok());
265 }
266
267 #[tokio::test]
268 async fn test_key_caching_and_retrieval() {
269 let temp_dir = tempdir().unwrap();
270 let cache = KeyCache::new(temp_dir.path().join("test.db"))
271 .await
272 .unwrap();
273
274 let verifying_key = generate_test_key();
275 let expires_at = SystemTime::now()
276 .duration_since(UNIX_EPOCH)
277 .unwrap()
278 .as_secs()
279 + 3600;
280
281 cache
282 .cache_key(1, &verifying_key, expires_at)
283 .await
284 .unwrap();
285 assert_eq!(cache.get_cached_key_count().await.unwrap(), 1);
286
287 let cached = cache.get_cached_key(1).await.unwrap();
288 assert!(cached.is_some());
289 let (retrieved_key, retrieved_expires) = cached.unwrap();
290 assert_eq!(verifying_key.as_bytes(), retrieved_key.as_bytes());
291 assert_eq!(expires_at, retrieved_expires);
292 }
293
294 #[tokio::test]
295 async fn test_cache_expiration() {
296 let temp_dir = tempdir().unwrap();
297 let cache = KeyCache::new(temp_dir.path().join("test.db"))
298 .await
299 .unwrap();
300
301 let verifying_key = generate_test_key();
302 let expires_at = SystemTime::now()
304 .duration_since(UNIX_EPOCH)
305 .unwrap()
306 .as_secs()
307 .saturating_sub(1);
308
309 cache
310 .cache_key(1, &verifying_key, expires_at)
311 .await
312 .unwrap();
313 let cached = cache.get_cached_key(1).await.unwrap();
315 assert!(cached.is_none());
316 }
317
318 #[tokio::test]
319 async fn test_cache_cleanup() {
320 let temp_dir = tempdir().unwrap();
321 let cache = KeyCache::new(temp_dir.path().join("test.db"))
322 .await
323 .unwrap();
324
325 let verifying_key = generate_test_key();
326 let expired_at = SystemTime::now()
327 .duration_since(UNIX_EPOCH)
328 .unwrap()
329 .as_secs()
330 .saturating_sub(1);
331
332 cache
333 .cache_key(1, &verifying_key, expired_at)
334 .await
335 .unwrap();
336 assert_eq!(cache.get_cached_key_count().await.unwrap(), 1);
337
338 let cleaned = cache.cleanup_expired_keys().await.unwrap();
339 assert_eq!(cleaned, 1);
340 assert_eq!(cache.get_cached_key_count().await.unwrap(), 0);
341 }
342}