platform/storage/nonce/
sqlite_nonce_storage.rs

1//! SQLite Nonce 存储实现
2//!
3//! 提供基于 SQLite 的 Nonce 存储功能实现，使用 sqlx 提供异步支持
4
5use anyhow::Result;
6use nonce_auth::NonceError;
7use nonce_auth::storage::{NonceEntry, NonceStorage, StorageStats};
8use sqlx::sqlite::{SqliteConnectOptions, SqlitePool, SqlitePoolOptions};
9use std::path::Path;
10use std::str::FromStr;
11use std::sync::Arc;
12use std::time::{Duration, SystemTime, UNIX_EPOCH};
13use tokio::sync::RwLock;
14
15use super::db_nonce_entry::DbNonceEntry;
16
17/// A sqlx-based implementation of NonceStorage for nonce-auth
18pub struct SqliteNonceStorage {
19    pool: Arc<SqlitePool>,
20    cleanup_lock: Arc<RwLock<()>>,
21}
22
23impl SqliteNonceStorage {
24    /// 创建新的 Nonce 存储实例（同步方法，适用于非 async 上下文）
25    pub fn new<P: AsRef<Path>>(db_path: P) -> Result<Self> {
26        let db_file = db_path.as_ref().join("nonce.db");
27
28        // 创建新的运行时来初始化（仅用于非 async 上下文）
29        let rt = tokio::runtime::Runtime::new()?;
30        let pool = rt.block_on(Self::init_pool(&db_file))?;
31
32        Ok(Self {
33            pool: Arc::new(pool),
34            cleanup_lock: Arc::new(RwLock::new(())),
35        })
36    }
37
38    /// 创建新的 Nonce 存储实例（异步方法，适用于 async 上下文）
39    pub async fn new_async<P: AsRef<Path>>(db_path: P) -> Result<Self> {
40        let db_file = db_path.as_ref().join("nonce.db");
41        let pool = Self::init_pool(&db_file).await?;
42
43        Ok(Self {
44            pool: Arc::new(pool),
45            cleanup_lock: Arc::new(RwLock::new(())),
46        })
47    }
48
49    async fn init_pool<P: AsRef<Path>>(db_file: P) -> Result<SqlitePool> {
50        let options =
51            SqliteConnectOptions::from_str(&format!("sqlite:{}", db_file.as_ref().display()))?
52                .create_if_missing(true)
53                .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
54                .synchronous(sqlx::sqlite::SqliteSynchronous::Normal)
55                .busy_timeout(Duration::from_secs(5));
56
57        let pool = SqlitePoolOptions::new()
58            .max_connections(10)
59            .connect_with(options)
60            .await?;
61
62        // 初始化数据库表
63        sqlx::query(
64            "CREATE TABLE IF NOT EXISTS nonce_entries (
65                id INTEGER PRIMARY KEY AUTOINCREMENT,
66                nonce TEXT NOT NULL,
67                context TEXT,
68                expires_at INTEGER NOT NULL,
69                created_at INTEGER NOT NULL,
70                UNIQUE(nonce, context)
71            )",
72        )
73        .execute(&pool)
74        .await?;
75
76        // 创建索引
77        sqlx::query(
78            "CREATE INDEX IF NOT EXISTS idx_nonce_context ON nonce_entries(nonce, context)",
79        )
80        .execute(&pool)
81        .await?;
82
83        sqlx::query("CREATE INDEX IF NOT EXISTS idx_expires_at ON nonce_entries(expires_at)")
84            .execute(&pool)
85            .await?;
86
87        Ok(pool)
88    }
89
90    /// Get a nonce entry from storage
91    pub async fn get(
92        &self,
93        nonce: &str,
94        context: Option<&str>,
95    ) -> Result<Option<DbNonceEntry>, NonceError> {
96        let result = if let Some(ctx) = context {
97            sqlx::query_as::<_, (i64, String, Option<String>, i64, i64)>(
98                "SELECT id, nonce, context, expires_at, created_at FROM nonce_entries WHERE nonce = ? AND context = ?",
99            )
100            .bind(nonce)
101            .bind(ctx)
102            .fetch_optional(&*self.pool)
103            .await
104        } else {
105            sqlx::query_as::<_, (i64, String, Option<String>, i64, i64)>(
106                "SELECT id, nonce, context, expires_at, created_at FROM nonce_entries WHERE nonce = ? AND context IS NULL",
107            )
108            .bind(nonce)
109            .fetch_optional(&*self.pool)
110            .await
111        };
112
113        match result {
114            Ok(Some((id, nonce, context, expires_at, created_at))) => Ok(Some(DbNonceEntry {
115                id: Some(id),
116                nonce,
117                context,
118                expires_at,
119                created_at,
120            })),
121            Ok(None) => Ok(None),
122            Err(e) => Err(NonceError::from_storage_error(e)),
123        }
124    }
125}
126
127#[async_trait::async_trait]
128impl NonceStorage for SqliteNonceStorage {
129    async fn get(
130        &self,
131        nonce: &str,
132        context: Option<&str>,
133    ) -> Result<Option<NonceEntry>, NonceError> {
134        let db_entry = self.get(nonce, context).await?;
135
136        // 检查是否过期
137        if let Some(entry) = db_entry {
138            let now = SystemTime::now()
139                .duration_since(UNIX_EPOCH)
140                .unwrap()
141                .as_secs() as i64;
142
143            if entry.expires_at <= now {
144                // 已过期，返回 None
145                return Ok(None);
146            }
147
148            Ok(Some(NonceEntry {
149                nonce: entry.nonce,
150                context: entry.context,
151                created_at: entry.created_at,
152            }))
153        } else {
154            Ok(None)
155        }
156    }
157
158    async fn set(
159        &self,
160        nonce: &str,
161        context: Option<&str>,
162        ttl: Duration,
163    ) -> Result<(), NonceError> {
164        let now = SystemTime::now()
165            .duration_since(UNIX_EPOCH)
166            .unwrap()
167            .as_secs() as i64;
168
169        let expires_at = now + ttl.as_secs() as i64;
170
171        let result = sqlx::query(
172            "INSERT INTO nonce_entries (nonce, context, expires_at, created_at) VALUES (?, ?, ?, ?)",
173        )
174        .bind(nonce)
175        .bind(context)
176        .bind(expires_at)
177        .bind(now)
178        .execute(&*self.pool)
179        .await;
180
181        match result {
182            Ok(_) => Ok(()),
183            Err(sqlx::Error::Database(e)) if e.message().contains("UNIQUE") => {
184                Err(NonceError::DuplicateNonce)
185            }
186            Err(e) => Err(NonceError::from_storage_error(e)),
187        }
188    }
189
190    async fn exists(&self, nonce: &str, context: Option<&str>) -> Result<bool, NonceError> {
191        Ok(self.get(nonce, context).await?.is_some())
192    }
193
194    async fn cleanup_expired(&self, current_time: i64) -> Result<usize, NonceError> {
195        // 使用锁防止并发清理
196        let _lock = self.cleanup_lock.write().await;
197
198        let result = sqlx::query("DELETE FROM nonce_entries WHERE expires_at < ?")
199            .bind(current_time)
200            .execute(&*self.pool)
201            .await
202            .map_err(NonceError::from_storage_error)?;
203
204        Ok(result.rows_affected() as usize)
205    }
206
207    async fn get_stats(&self) -> Result<StorageStats, NonceError> {
208        let total: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM nonce_entries")
209            .fetch_one(&*self.pool)
210            .await
211            .map_err(NonceError::from_storage_error)?;
212
213        Ok(StorageStats {
214            total_records: total.0 as usize,
215            backend_info: "SQLite (sqlx async)".to_string(),
216        })
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223    use tempfile::tempdir;
224
225    #[tokio::test]
226    async fn test_basic_storage() {
227        let temp_dir = tempdir().unwrap();
228        let storage = SqliteNonceStorage::new_async(temp_dir.path())
229            .await
230            .unwrap();
231
232        // Set nonce with TTL
233        storage
234            .set(
235                "test_nonce",
236                Some("test_context"),
237                Duration::from_secs(3600),
238            )
239            .await
240            .unwrap();
241
242        // Get should return the nonce
243        let result = storage
244            .get("test_nonce", Some("test_context"))
245            .await
246            .unwrap();
247        assert!(result.is_some());
248        assert_eq!(result.unwrap().nonce, "test_nonce");
249
250        // Duplicate should fail
251        assert!(matches!(
252            storage
253                .set(
254                    "test_nonce",
255                    Some("test_context"),
256                    Duration::from_secs(3600)
257                )
258                .await,
259            Err(NonceError::DuplicateNonce)
260        ));
261    }
262
263    #[tokio::test]
264    async fn test_cleanup() {
265        let temp_dir = tempdir().unwrap();
266        let storage = SqliteNonceStorage::new_async(temp_dir.path())
267            .await
268            .unwrap();
269
270        let before_insert = SystemTime::now()
271            .duration_since(UNIX_EPOCH)
272            .unwrap()
273            .as_secs();
274        println!("Before insert: {before_insert}");
275
276        // 添加即将过期的 nonce（1 秒TTL）
277        storage
278            .set("expired", None, Duration::from_secs(1))
279            .await
280            .unwrap();
281
282        // 添加未过期的 nonce (1小时TTL)
283        storage
284            .set("valid", None, Duration::from_secs(3600))
285            .await
286            .unwrap();
287
288        // 等待确保第一个过期（等待 2 秒）
289        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
290
291        let after_sleep = SystemTime::now()
292            .duration_since(UNIX_EPOCH)
293            .unwrap()
294            .as_secs();
295        println!("After sleep: {after_sleep}");
296        println!("Time elapsed: {} seconds", after_sleep - before_insert);
297
298        // 通过 NonceStorage trait 调用（会检查过期）
299        use nonce_auth::storage::NonceStorage;
300        let expired_entry = NonceStorage::get(&storage, "expired", None).await.unwrap();
301        println!("Expired entry after trait call: {expired_entry:?}");
302
303        // 验证：过期的应该查不到（trait方法内部会检查过期）
304        assert!(
305            expired_entry.is_none(),
306            "Expected expired nonce to return None"
307        );
308
309        let valid_entry = NonceStorage::get(&storage, "valid", None).await.unwrap();
310        assert!(valid_entry.is_some(), "Expected valid nonce to be present");
311    }
312
313    #[tokio::test]
314    async fn test_stats() {
315        let temp_dir = tempdir().unwrap();
316        let storage = SqliteNonceStorage::new_async(temp_dir.path())
317            .await
318            .unwrap();
319
320        // 添加一些 nonce
321        storage
322            .set("n1", None, Duration::from_secs(1))
323            .await
324            .unwrap();
325        storage
326            .set("n2", None, Duration::from_secs(3600))
327            .await
328            .unwrap();
329
330        let stats = storage.get_stats().await.unwrap();
331        assert_eq!(stats.total_records, 2);
332        assert!(stats.backend_info.contains("SQLite"));
333    }
334}
platform/storage/nonce/sqlite_nonce_storage.rs

platform/storage/nonce/
sqlite_nonce_storage.rs
