1use async_trait::async_trait;
16use chrono::{DateTime, Utc};
17use sqlx::{Row, SqlitePool};
18use tracing::{debug, info};
19
20use crate::{Result, SecretsError};
21
22pub use zlayer_types::secrets::client_keys::{ActorKind, ClientPublicKey, PUBLIC_KEY_LEN};
23
24const SCHEMA: &str = r"
27CREATE TABLE IF NOT EXISTS client_public_keys (
28 key_id TEXT PRIMARY KEY,
29 actor_kind TEXT NOT NULL CHECK(actor_kind IN ('user','api_key')),
30 actor_id TEXT NOT NULL,
31 public_key BLOB NOT NULL,
32 label TEXT,
33 created_at TEXT NOT NULL,
34 last_used_at TEXT,
35 revoked_at TEXT
36);
37CREATE INDEX IF NOT EXISTS idx_client_public_keys_actor
38 ON client_public_keys(actor_kind, actor_id) WHERE revoked_at IS NULL;
39";
40
41#[async_trait]
43pub trait ClientKeyStore: Send + Sync {
44 async fn register(
46 &self,
47 actor_kind: ActorKind,
48 actor_id: &str,
49 public_key: &[u8],
50 label: Option<&str>,
51 ) -> Result<ClientPublicKey>;
52
53 async fn get(&self, key_id: &str) -> Result<Option<ClientPublicKey>>;
56
57 async fn list_by_actor(
59 &self,
60 actor_kind: ActorKind,
61 actor_id: &str,
62 ) -> Result<Vec<ClientPublicKey>>;
63
64 async fn revoke(&self, key_id: &str) -> Result<()>;
66
67 async fn touch_last_used(&self, key_id: &str) -> Result<()>;
69}
70
71pub struct PersistentClientKeyStore {
76 pool: SqlitePool,
77}
78
79impl PersistentClientKeyStore {
80 pub async fn new(pool: SqlitePool) -> Result<Self> {
86 sqlx::query(SCHEMA)
87 .execute(&pool)
88 .await
89 .map_err(|e| SecretsError::Storage(format!("Failed to initialize schema: {e}")))?;
90
91 info!("Initialized client public keys schema");
92 Ok(Self { pool })
93 }
94
95 #[must_use]
98 pub fn pool(&self) -> &SqlitePool {
99 &self.pool
100 }
101
102 fn generate_key_id() -> String {
104 let bytes: [u8; 16] = rand::random();
105 format!("ck_{}", hex::encode(bytes))
106 }
107
108 fn format_timestamp(ts: DateTime<Utc>) -> String {
112 ts.to_rfc3339_opts(chrono::SecondsFormat::Millis, true)
113 }
114
115 fn parse_timestamp(s: &str) -> Result<DateTime<Utc>> {
117 DateTime::parse_from_rfc3339(s)
118 .map(|dt| dt.with_timezone(&Utc))
119 .map_err(|e| SecretsError::Storage(format!("invalid timestamp {s:?}: {e}")))
120 }
121
122 fn row_to_key(row: &sqlx::sqlite::SqliteRow) -> Result<ClientPublicKey> {
124 let key_id: String = row
125 .try_get("key_id")
126 .map_err(|e| SecretsError::Storage(format!("Failed to read key_id: {e}")))?;
127 let actor_kind_str: String = row
128 .try_get("actor_kind")
129 .map_err(|e| SecretsError::Storage(format!("Failed to read actor_kind: {e}")))?;
130 let actor_kind = ActorKind::from_str(&actor_kind_str)?;
131 let actor_id: String = row
132 .try_get("actor_id")
133 .map_err(|e| SecretsError::Storage(format!("Failed to read actor_id: {e}")))?;
134 let public_key: Vec<u8> = row
135 .try_get("public_key")
136 .map_err(|e| SecretsError::Storage(format!("Failed to read public_key: {e}")))?;
137 let label: Option<String> = row
138 .try_get("label")
139 .map_err(|e| SecretsError::Storage(format!("Failed to read label: {e}")))?;
140 let created_at_str: String = row
141 .try_get("created_at")
142 .map_err(|e| SecretsError::Storage(format!("Failed to read created_at: {e}")))?;
143 let last_used_at_str: Option<String> = row
144 .try_get("last_used_at")
145 .map_err(|e| SecretsError::Storage(format!("Failed to read last_used_at: {e}")))?;
146 let revoked_at_str: Option<String> = row
147 .try_get("revoked_at")
148 .map_err(|e| SecretsError::Storage(format!("Failed to read revoked_at: {e}")))?;
149
150 let created_at = Self::parse_timestamp(&created_at_str)?;
151 let last_used_at = match last_used_at_str {
152 Some(s) => Some(Self::parse_timestamp(&s)?),
153 None => None,
154 };
155 let revoked_at = match revoked_at_str {
156 Some(s) => Some(Self::parse_timestamp(&s)?),
157 None => None,
158 };
159
160 Ok(ClientPublicKey {
161 key_id,
162 actor_kind,
163 actor_id,
164 public_key,
165 label,
166 created_at,
167 last_used_at,
168 revoked_at,
169 })
170 }
171}
172
173#[async_trait]
174impl ClientKeyStore for PersistentClientKeyStore {
175 async fn register(
176 &self,
177 actor_kind: ActorKind,
178 actor_id: &str,
179 public_key: &[u8],
180 label: Option<&str>,
181 ) -> Result<ClientPublicKey> {
182 if public_key.len() != PUBLIC_KEY_LEN {
183 return Err(SecretsError::Storage(format!(
184 "invalid public key length: expected {PUBLIC_KEY_LEN} bytes, got {}",
185 public_key.len()
186 )));
187 }
188
189 let key_id = Self::generate_key_id();
190 let created_at = Utc::now();
191 let created_at_str = Self::format_timestamp(created_at);
192 let public_key_vec = public_key.to_vec();
193
194 sqlx::query(
195 "INSERT INTO client_public_keys \
196 (key_id, actor_kind, actor_id, public_key, label, created_at, last_used_at, revoked_at) \
197 VALUES (?, ?, ?, ?, ?, ?, NULL, NULL)",
198 )
199 .bind(&key_id)
200 .bind(actor_kind.as_str())
201 .bind(actor_id)
202 .bind(&public_key_vec)
203 .bind(label)
204 .bind(&created_at_str)
205 .execute(&self.pool)
206 .await
207 .map_err(|e| SecretsError::Storage(format!("Failed to insert client public key: {e}")))?;
208
209 debug!(
210 "Registered client public key {} for {} {}",
211 key_id,
212 actor_kind.as_str(),
213 actor_id
214 );
215
216 Ok(ClientPublicKey {
217 key_id,
218 actor_kind,
219 actor_id: actor_id.to_string(),
220 public_key: public_key_vec,
221 label: label.map(str::to_string),
222 created_at,
223 last_used_at: None,
224 revoked_at: None,
225 })
226 }
227
228 async fn get(&self, key_id: &str) -> Result<Option<ClientPublicKey>> {
229 let row = sqlx::query(
230 "SELECT key_id, actor_kind, actor_id, public_key, label, created_at, last_used_at, revoked_at \
231 FROM client_public_keys WHERE key_id = ?",
232 )
233 .bind(key_id)
234 .fetch_optional(&self.pool)
235 .await
236 .map_err(|e| SecretsError::Storage(format!("Failed to query client public key: {e}")))?;
237
238 match row {
239 Some(row) => Ok(Some(Self::row_to_key(&row)?)),
240 None => Ok(None),
241 }
242 }
243
244 async fn list_by_actor(
245 &self,
246 actor_kind: ActorKind,
247 actor_id: &str,
248 ) -> Result<Vec<ClientPublicKey>> {
249 let rows = sqlx::query(
250 "SELECT key_id, actor_kind, actor_id, public_key, label, created_at, last_used_at, revoked_at \
251 FROM client_public_keys \
252 WHERE actor_kind = ? AND actor_id = ? AND revoked_at IS NULL \
253 ORDER BY created_at DESC",
254 )
255 .bind(actor_kind.as_str())
256 .bind(actor_id)
257 .fetch_all(&self.pool)
258 .await
259 .map_err(|e| SecretsError::Storage(format!("Failed to list client public keys: {e}")))?;
260
261 let mut out = Vec::with_capacity(rows.len());
262 for row in &rows {
263 out.push(Self::row_to_key(row)?);
264 }
265 Ok(out)
266 }
267
268 async fn revoke(&self, key_id: &str) -> Result<()> {
269 let now = Self::format_timestamp(Utc::now());
270
271 let result = sqlx::query("UPDATE client_public_keys SET revoked_at = ? WHERE key_id = ?")
272 .bind(&now)
273 .bind(key_id)
274 .execute(&self.pool)
275 .await
276 .map_err(|e| {
277 SecretsError::Storage(format!("Failed to revoke client public key: {e}"))
278 })?;
279
280 if result.rows_affected() == 0 {
281 return Err(SecretsError::NotFound {
282 name: key_id.to_string(),
283 });
284 }
285
286 debug!("Revoked client public key {}", key_id);
287 Ok(())
288 }
289
290 async fn touch_last_used(&self, key_id: &str) -> Result<()> {
291 let now = Self::format_timestamp(Utc::now());
292
293 let result = sqlx::query("UPDATE client_public_keys SET last_used_at = ? WHERE key_id = ?")
294 .bind(&now)
295 .bind(key_id)
296 .execute(&self.pool)
297 .await
298 .map_err(|e| SecretsError::Storage(format!("Failed to update last_used_at: {e}")))?;
299
300 if result.rows_affected() == 0 {
301 return Err(SecretsError::NotFound {
302 name: key_id.to_string(),
303 });
304 }
305
306 Ok(())
307 }
308}
309
310#[cfg(test)]
311mod tests {
312 use super::*;
313
314 async fn create_test_store() -> PersistentClientKeyStore {
315 let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
316 PersistentClientKeyStore::new(pool).await.unwrap()
317 }
318
319 #[tokio::test]
320 async fn register_and_get_roundtrip() {
321 let store = create_test_store().await;
322 let pk = [7u8; 32];
323 let registered = store
324 .register(ActorKind::User, "user-123", &pk, Some("laptop"))
325 .await
326 .unwrap();
327
328 assert!(registered.key_id.starts_with("ck_"));
329 assert_eq!(registered.key_id.len(), 3 + 32);
330 assert_eq!(registered.actor_kind, ActorKind::User);
331 assert_eq!(registered.actor_id, "user-123");
332 assert_eq!(registered.public_key, pk.to_vec());
333 assert_eq!(registered.label.as_deref(), Some("laptop"));
334 assert!(registered.last_used_at.is_none());
335 assert!(registered.revoked_at.is_none());
336
337 let fetched = store.get(®istered.key_id).await.unwrap().unwrap();
338 assert_eq!(fetched.key_id, registered.key_id);
339 assert_eq!(fetched.actor_kind, ActorKind::User);
340 assert_eq!(fetched.actor_id, "user-123");
341 assert_eq!(fetched.public_key, pk.to_vec());
342 assert_eq!(fetched.label.as_deref(), Some("laptop"));
343
344 assert!(store.get("ck_does_not_exist").await.unwrap().is_none());
346 }
347
348 #[tokio::test]
349 async fn duplicate_key_id_errors() {
350 let store = create_test_store().await;
351 let pk = [3u8; 32];
352
353 let registered = store
354 .register(ActorKind::ApiKey, "api-1", &pk, None)
355 .await
356 .unwrap();
357
358 let result = sqlx::query(
361 "INSERT INTO client_public_keys \
362 (key_id, actor_kind, actor_id, public_key, label, created_at, last_used_at, revoked_at) \
363 VALUES (?, 'api_key', 'api-1', ?, NULL, '2026-01-01T00:00:00Z', NULL, NULL)",
364 )
365 .bind(®istered.key_id)
366 .bind(pk.to_vec())
367 .execute(store.pool())
368 .await;
369
370 assert!(result.is_err(), "duplicate key_id must be rejected");
371 }
372
373 #[tokio::test]
374 async fn list_by_actor_returns_only_active() {
375 let store = create_test_store().await;
376 let pk = [9u8; 32];
377
378 let a = store
379 .register(ActorKind::User, "u1", &pk, Some("a"))
380 .await
381 .unwrap();
382 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
385 let _b = store
386 .register(ActorKind::User, "u1", &pk, Some("b"))
387 .await
388 .unwrap();
389 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
390 let c = store
391 .register(ActorKind::User, "u1", &pk, Some("c"))
392 .await
393 .unwrap();
394 let _other = store
396 .register(ActorKind::User, "u2", &pk, Some("other"))
397 .await
398 .unwrap();
399 let _other_kind = store
401 .register(ActorKind::ApiKey, "u1", &pk, Some("api"))
402 .await
403 .unwrap();
404
405 store.revoke(&a.key_id).await.unwrap();
407
408 let active = store.list_by_actor(ActorKind::User, "u1").await.unwrap();
409 assert_eq!(active.len(), 2);
410 assert_eq!(active[0].key_id, c.key_id);
412 assert!(active.iter().all(|k| k.revoked_at.is_none()));
413 assert!(active.iter().all(|k| k.actor_id == "u1"));
414 assert!(active
415 .iter()
416 .all(|k| matches!(k.actor_kind, ActorKind::User)));
417 }
418
419 #[tokio::test]
420 async fn revoke_hides_from_list_but_get_still_finds_with_revoked_at() {
421 let store = create_test_store().await;
422 let pk = [1u8; 32];
423
424 let key = store
425 .register(ActorKind::ApiKey, "svc-1", &pk, None)
426 .await
427 .unwrap();
428
429 let pre = store
431 .list_by_actor(ActorKind::ApiKey, "svc-1")
432 .await
433 .unwrap();
434 assert_eq!(pre.len(), 1);
435 assert_eq!(pre[0].key_id, key.key_id);
436
437 store.revoke(&key.key_id).await.unwrap();
438
439 let post = store
441 .list_by_actor(ActorKind::ApiKey, "svc-1")
442 .await
443 .unwrap();
444 assert!(post.is_empty());
445
446 let fetched = store.get(&key.key_id).await.unwrap().unwrap();
448 assert!(fetched.revoked_at.is_some());
449
450 let missing = store.revoke("ck_nope").await;
452 assert!(matches!(missing, Err(SecretsError::NotFound { .. })));
453 }
454
455 #[tokio::test]
456 async fn invalid_public_key_length_rejected() {
457 let store = create_test_store().await;
458
459 let too_short = [0u8; 16];
460 let err = store
461 .register(ActorKind::User, "u1", &too_short, None)
462 .await
463 .unwrap_err();
464 assert!(matches!(err, SecretsError::Storage(_)));
465
466 let too_long = [0u8; 64];
467 let err = store
468 .register(ActorKind::User, "u1", &too_long, None)
469 .await
470 .unwrap_err();
471 assert!(matches!(err, SecretsError::Storage(_)));
472
473 let empty: &[u8] = &[];
474 let err = store
475 .register(ActorKind::User, "u1", empty, None)
476 .await
477 .unwrap_err();
478 assert!(matches!(err, SecretsError::Storage(_)));
479
480 let list = store.list_by_actor(ActorKind::User, "u1").await.unwrap();
482 assert!(list.is_empty());
483 }
484
485 #[tokio::test]
486 async fn touch_last_used_updates_timestamp() {
487 let store = create_test_store().await;
488 let pk = [2u8; 32];
489
490 let key = store
491 .register(ActorKind::User, "u1", &pk, None)
492 .await
493 .unwrap();
494 assert!(key.last_used_at.is_none());
495
496 store.touch_last_used(&key.key_id).await.unwrap();
497
498 let fetched = store.get(&key.key_id).await.unwrap().unwrap();
499 assert!(fetched.last_used_at.is_some());
500
501 let err = store.touch_last_used("ck_nope").await.unwrap_err();
503 assert!(matches!(err, SecretsError::NotFound { .. }));
504 }
505
506 #[test]
507 fn actor_kind_str_roundtrip() {
508 assert_eq!(ActorKind::User.as_str(), "user");
509 assert_eq!(ActorKind::ApiKey.as_str(), "api_key");
510 assert_eq!(ActorKind::from_str("user").unwrap(), ActorKind::User);
511 assert_eq!(ActorKind::from_str("api_key").unwrap(), ActorKind::ApiKey);
512 assert!(matches!(
513 ActorKind::from_str("garbage"),
514 Err(SecretsError::Storage(_))
515 ));
516 }
517}