Skip to main content

platform/aid/
key_cache.rs

1//! AIS 签名公钥缓存
2//!
3//! 为验证器提供本地 SQLite 缓存,避免重复从 AIS 注册响应中重新获取 Ed25519 公钥。
4//! 缓存项以 key_id 为索引,存储对应的 Ed25519 verifying key(32 bytes)和过期时间。
5
6use 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/// AIS 签名公钥缓存管理器
17///
18/// 缓存 AIS Ed25519 verifying key,以 key_id 为索引。
19/// 用于 signaling 等服务在验证 AIdCredential 时按 key_id 查找对应公钥。
20#[derive(Debug, Clone)]
21pub struct KeyCache {
22    pool: SqlitePool,
23    last_cleanup_time: Arc<Mutex<u64>>,
24}
25
26impl KeyCache {
27    /// 创建新的密钥缓存实例并初始化 SQLite 表
28    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    /// 初始化缓存数据库表
60    async fn init_tables(&self) -> Result<(), AidError> {
61        // key_cache 表:存储 Ed25519 verifying key(base64 编码,32 bytes)和过期时间
62        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    /// 从缓存中获取 Ed25519 verifying key
88    ///
89    /// 返回 `(VerifyingKey, expires_at)`;若缓存不存在或已过期则返回 `None`。
90    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                // 缓存条目已过期,删除并返回 None
120                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                // base64 解码 → 32 bytes → VerifyingKey
130                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    /// 将 Ed25519 verifying key 写入缓存
155    ///
156    /// `expires_at` 为 Unix 秒时间戳,传 0 表示永不过期。
157    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    /// 清理所有已过期的缓存条目
190    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    /// 返回缓存中的条目总数(用于监控/调试)
210    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    /// 内部:每小时触发一次过期清理
224    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        // 设置已过期的时间戳
303        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        // 过期条目应返回 None
314        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}