assay_core/runtime/
mandate_store.rs

1//! MandateStore: SQLite-backed mandate consumption tracking.
2//!
3//! Provides atomic, idempotent mandate consumption with:
4//! - Single-use / max_uses constraint enforcement
5//! - Nonce replay prevention
6//! - tool_call_id idempotency
7
8use super::schema::MANDATE_SCHEMA;
9use chrono::{DateTime, Utc};
10use rusqlite::{params, Connection, OptionalExtension};
11use sha2::{Digest, Sha256};
12use std::path::Path;
13use std::sync::{Arc, Mutex};
14use thiserror::Error;
15
16/// Authorization receipt returned after successful consumption.
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct AuthzReceipt {
19    pub mandate_id: String,
20    pub use_id: String,
21    pub use_count: u32,
22    pub consumed_at: DateTime<Utc>,
23    pub tool_call_id: String,
24    /// True if this was a new consumption, false if idempotent retry.
25    /// Used to avoid emitting duplicate lifecycle events on retries.
26    pub was_new: bool,
27}
28
29/// Authorization errors.
30#[derive(Debug, Error, PartialEq, Eq)]
31pub enum AuthzError {
32    #[error("Mandate not found: {mandate_id}")]
33    MandateNotFound { mandate_id: String },
34
35    #[error("Mandate already used (single_use=true)")]
36    AlreadyUsed,
37
38    #[error("Max uses exceeded: {current} > {max}")]
39    MaxUsesExceeded { max: u32, current: u32 },
40
41    #[error("Nonce replay detected: {nonce}")]
42    NonceReplay { nonce: String },
43
44    #[error("Mandate metadata conflict for {mandate_id}: stored {field} differs")]
45    MandateConflict { mandate_id: String, field: String },
46
47    #[error("Invalid mandate constraints: single_use=true with max_uses={max_uses}")]
48    InvalidConstraints { max_uses: u32 },
49
50    #[error("Mandate revoked at {revoked_at}")]
51    Revoked { revoked_at: DateTime<Utc> },
52
53    #[error("Database error: {0}")]
54    Database(String),
55}
56
57impl From<rusqlite::Error> for AuthzError {
58    fn from(e: rusqlite::Error) -> Self {
59        AuthzError::Database(e.to_string())
60    }
61}
62
63/// Mandate metadata for upsert.
64#[derive(Debug, Clone)]
65pub struct MandateMetadata {
66    pub mandate_id: String,
67    pub mandate_kind: String,
68    pub audience: String,
69    pub issuer: String,
70    pub expires_at: Option<DateTime<Utc>>,
71    pub single_use: bool,
72    pub max_uses: Option<u32>,
73    pub canonical_digest: String,
74    pub key_id: String,
75}
76
77/// Parameters for consume_mandate.
78#[derive(Debug, Clone)]
79pub struct ConsumeParams<'a> {
80    pub mandate_id: &'a str,
81    pub tool_call_id: &'a str,
82    pub nonce: Option<&'a str>,
83    pub audience: &'a str,
84    pub issuer: &'a str,
85    pub tool_name: &'a str,
86    pub operation_class: &'a str,
87    pub source_run_id: Option<&'a str>,
88}
89
90/// SQLite-backed mandate store.
91#[derive(Clone)]
92pub struct MandateStore {
93    conn: Arc<Mutex<Connection>>,
94}
95
96impl MandateStore {
97    /// Open a file-backed store.
98    pub fn open(path: &Path) -> Result<Self, AuthzError> {
99        let conn = Connection::open(path)?;
100        Self::init_connection(&conn)?;
101        Ok(Self {
102            conn: Arc::new(Mutex::new(conn)),
103        })
104    }
105
106    /// Create an in-memory store (for testing).
107    pub fn memory() -> Result<Self, AuthzError> {
108        let conn = Connection::open_in_memory()?;
109        Self::init_connection(&conn)?;
110        Ok(Self {
111            conn: Arc::new(Mutex::new(conn)),
112        })
113    }
114
115    /// Create store from existing connection (for multi-connection tests).
116    pub fn from_connection(conn: Connection) -> Result<Self, AuthzError> {
117        Self::init_connection(&conn)?;
118        Ok(Self {
119            conn: Arc::new(Mutex::new(conn)),
120        })
121    }
122
123    fn init_connection(conn: &Connection) -> Result<(), AuthzError> {
124        conn.execute("PRAGMA foreign_keys = ON", [])?;
125        // WAL mode for file-backed DBs (no-op for in-memory)
126        let _ = conn.execute("PRAGMA journal_mode = WAL", []);
127        // Busy timeout: wait up to 5s for locks (Windows needs this for concurrency)
128        // Use let _ = because PRAGMA returns a result set
129        let _ = conn.execute("PRAGMA busy_timeout = 5000", []);
130        conn.execute_batch(MANDATE_SCHEMA)?;
131        Ok(())
132    }
133
134    /// Upsert mandate metadata. Idempotent for same content, errors on conflict.
135    pub fn upsert_mandate(&self, meta: &MandateMetadata) -> Result<(), AuthzError> {
136        // Validate constraints: single_use implies max_uses == 1
137        if meta.single_use {
138            if let Some(max) = meta.max_uses {
139                if max != 1 {
140                    return Err(AuthzError::InvalidConstraints { max_uses: max });
141                }
142            }
143        }
144
145        let conn = self.conn.lock().unwrap();
146
147        // Insert with ON CONFLICT DO NOTHING
148        conn.execute(
149            r#"
150            INSERT INTO mandates (
151                mandate_id, mandate_kind, audience, issuer, expires_at,
152                single_use, max_uses, use_count, canonical_digest, key_id
153            ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, 0, ?8, ?9)
154            ON CONFLICT(mandate_id) DO NOTHING
155            "#,
156            params![
157                meta.mandate_id,
158                meta.mandate_kind,
159                meta.audience,
160                meta.issuer,
161                meta.expires_at.map(|t| t.to_rfc3339()),
162                meta.single_use as i32,
163                meta.max_uses.map(|m| m as i64),
164                meta.canonical_digest,
165                meta.key_id,
166            ],
167        )?;
168
169        // Verify consistency if already existed
170        let stored: Option<(String, String, String, String, String)> = conn
171            .query_row(
172                r#"
173                SELECT mandate_kind, audience, issuer, canonical_digest, key_id
174                FROM mandates WHERE mandate_id = ?
175                "#,
176                [&meta.mandate_id],
177                |row| {
178                    Ok((
179                        row.get(0)?,
180                        row.get(1)?,
181                        row.get(2)?,
182                        row.get(3)?,
183                        row.get(4)?,
184                    ))
185                },
186            )
187            .optional()?;
188
189        if let Some((kind, aud, iss, digest, key)) = stored {
190            if kind != meta.mandate_kind {
191                return Err(AuthzError::MandateConflict {
192                    mandate_id: meta.mandate_id.clone(),
193                    field: "mandate_kind".to_string(),
194                });
195            }
196            if aud != meta.audience {
197                return Err(AuthzError::MandateConflict {
198                    mandate_id: meta.mandate_id.clone(),
199                    field: "audience".to_string(),
200                });
201            }
202            if iss != meta.issuer {
203                return Err(AuthzError::MandateConflict {
204                    mandate_id: meta.mandate_id.clone(),
205                    field: "issuer".to_string(),
206                });
207            }
208            if digest != meta.canonical_digest {
209                return Err(AuthzError::MandateConflict {
210                    mandate_id: meta.mandate_id.clone(),
211                    field: "canonical_digest".to_string(),
212                });
213            }
214            if key != meta.key_id {
215                return Err(AuthzError::MandateConflict {
216                    mandate_id: meta.mandate_id.clone(),
217                    field: "key_id".to_string(),
218                });
219            }
220        }
221
222        Ok(())
223    }
224
225    /// Consume mandate atomically. Idempotent on tool_call_id.
226    pub fn consume_mandate(&self, params: &ConsumeParams<'_>) -> Result<AuthzReceipt, AuthzError> {
227        let conn = self.conn.lock().unwrap();
228
229        // BEGIN IMMEDIATE acquires write lock immediately
230        conn.execute("BEGIN IMMEDIATE", [])?;
231
232        let result = self.consume_mandate_inner(&conn, params);
233
234        match &result {
235            Ok(_) => {
236                conn.execute("COMMIT", [])?;
237            }
238            Err(_) => {
239                let _ = conn.execute("ROLLBACK", []);
240            }
241        }
242
243        result
244    }
245
246    fn consume_mandate_inner(
247        &self,
248        conn: &Connection,
249        params: &ConsumeParams<'_>,
250    ) -> Result<AuthzReceipt, AuthzError> {
251        let ConsumeParams {
252            mandate_id,
253            tool_call_id,
254            nonce,
255            audience,
256            issuer,
257            tool_name,
258            operation_class,
259            source_run_id,
260        } = params;
261        // Step 1: Idempotency check
262        let existing: Option<(String, i64, String)> = conn
263            .query_row(
264                "SELECT use_id, use_count, consumed_at FROM mandate_uses WHERE tool_call_id = ?",
265                [tool_call_id],
266                |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
267            )
268            .optional()?;
269
270        if let Some((use_id, use_count, consumed_at)) = existing {
271            // Return existing receipt (idempotent retry)
272            return Ok(AuthzReceipt {
273                mandate_id: mandate_id.to_string(),
274                use_id,
275                use_count: use_count as u32,
276                consumed_at: DateTime::parse_from_rfc3339(&consumed_at)
277                    .map(|dt| dt.with_timezone(&Utc))
278                    .unwrap_or_else(|_| Utc::now()),
279                tool_call_id: tool_call_id.to_string(),
280                was_new: false, // Idempotent retry
281            });
282        }
283
284        // Step 2: Nonce replay check (atomic INSERT, not SELECT+INSERT)
285        if let Some(n) = nonce {
286            let insert_result = conn.execute(
287                "INSERT INTO nonces (audience, issuer, nonce, mandate_id) VALUES (?1, ?2, ?3, ?4)",
288                params![audience, issuer, n, mandate_id],
289            );
290
291            if let Err(e) = insert_result {
292                if e.to_string().contains("UNIQUE constraint failed") {
293                    return Err(AuthzError::NonceReplay {
294                        nonce: n.to_string(),
295                    });
296                }
297                return Err(e.into());
298            }
299        }
300
301        // Step 3: Get mandate metadata + current use count
302        let row: Option<(i64, i32, Option<i64>)> = conn
303            .query_row(
304                "SELECT use_count, single_use, max_uses FROM mandates WHERE mandate_id = ?",
305                [mandate_id],
306                |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
307            )
308            .optional()?;
309
310        let (current_count, single_use, max_uses) = match row {
311            Some(r) => r,
312            None => {
313                return Err(AuthzError::MandateNotFound {
314                    mandate_id: mandate_id.to_string(),
315                });
316            }
317        };
318
319        let new_count = current_count + 1;
320
321        // Step 4: Check constraints
322        if single_use != 0 && current_count > 0 {
323            return Err(AuthzError::AlreadyUsed);
324        }
325
326        if let Some(max) = max_uses {
327            if new_count > max {
328                return Err(AuthzError::MaxUsesExceeded {
329                    max: max as u32,
330                    current: new_count as u32,
331                });
332            }
333        }
334
335        // Step 5: Increment count + insert use record
336        conn.execute(
337            "UPDATE mandates SET use_count = ?1 WHERE mandate_id = ?2",
338            params![new_count, mandate_id],
339        )?;
340
341        // use_id is content-addressed (deterministic) per SPEC-Mandate-v1.0.4 §7.4
342        // use_id = sha256(mandate_id + ":" + tool_call_id + ":" + use_count)
343        let use_id = compute_use_id(mandate_id, tool_call_id, new_count as u32);
344        let consumed_at = Utc::now();
345
346        conn.execute(
347            r#"
348            INSERT INTO mandate_uses (
349                use_id, mandate_id, tool_call_id, use_count, consumed_at,
350                tool_name, operation_class, nonce, source_run_id
351            ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
352            "#,
353            params![
354                use_id,
355                mandate_id,
356                tool_call_id,
357                new_count,
358                consumed_at.to_rfc3339(),
359                tool_name,
360                operation_class,
361                nonce,
362                source_run_id,
363            ],
364        )?;
365
366        Ok(AuthzReceipt {
367            mandate_id: mandate_id.to_string(),
368            use_id,
369            use_count: new_count as u32,
370            consumed_at,
371            tool_call_id: tool_call_id.to_string(),
372            was_new: true, // First consumption
373        })
374    }
375
376    /// Get current use count for a mandate (for testing/debugging).
377    pub fn get_use_count(&self, mandate_id: &str) -> Result<Option<u32>, AuthzError> {
378        let conn = self.conn.lock().unwrap();
379        let count: Option<i64> = conn
380            .query_row(
381                "SELECT use_count FROM mandates WHERE mandate_id = ?",
382                [mandate_id],
383                |row| row.get(0),
384            )
385            .optional()?;
386        Ok(count.map(|c| c as u32))
387    }
388
389    /// Count use records for a mandate (for testing).
390    pub fn count_uses(&self, mandate_id: &str) -> Result<u32, AuthzError> {
391        let conn = self.conn.lock().unwrap();
392        let count: i64 = conn.query_row(
393            "SELECT COUNT(*) FROM mandate_uses WHERE mandate_id = ?",
394            [mandate_id],
395            |row| row.get(0),
396        )?;
397        Ok(count as u32)
398    }
399
400    /// Check if nonce exists (for testing).
401    pub fn nonce_exists(
402        &self,
403        audience: &str,
404        issuer: &str,
405        nonce: &str,
406    ) -> Result<bool, AuthzError> {
407        let conn = self.conn.lock().unwrap();
408        let exists: i64 = conn.query_row(
409            "SELECT COUNT(*) FROM nonces WHERE audience = ? AND issuer = ? AND nonce = ?",
410            params![audience, issuer, nonce],
411            |row| row.get(0),
412        )?;
413        Ok(exists > 0)
414    }
415
416    // =========================================================================
417    // Revocation API (P0-A)
418    // =========================================================================
419
420    /// Insert or update a revocation record.
421    ///
422    /// Idempotent: re-inserting with same mandate_id updates the record.
423    pub fn upsert_revocation(&self, r: &RevocationRecord) -> Result<(), AuthzError> {
424        let conn = self.conn.lock().unwrap();
425        conn.execute(
426            r#"
427            INSERT INTO mandate_revocations (mandate_id, revoked_at, reason, revoked_by, source, event_id)
428            VALUES (?1, ?2, ?3, ?4, ?5, ?6)
429            ON CONFLICT(mandate_id) DO UPDATE SET
430                revoked_at = excluded.revoked_at,
431                reason = excluded.reason,
432                revoked_by = excluded.revoked_by,
433                source = excluded.source,
434                event_id = excluded.event_id
435            "#,
436            params![
437                r.mandate_id,
438                r.revoked_at.to_rfc3339(),
439                r.reason,
440                r.revoked_by,
441                r.source,
442                r.event_id,
443            ],
444        )?;
445        Ok(())
446    }
447
448    /// Get revoked_at timestamp for a mandate (if revoked).
449    pub fn get_revoked_at(&self, mandate_id: &str) -> Result<Option<DateTime<Utc>>, AuthzError> {
450        let conn = self.conn.lock().unwrap();
451        let s: Option<String> = conn
452            .query_row(
453                "SELECT revoked_at FROM mandate_revocations WHERE mandate_id = ?1",
454                [mandate_id],
455                |row| row.get(0),
456            )
457            .optional()?;
458
459        match s {
460            Some(ts) => {
461                let dt = DateTime::parse_from_rfc3339(&ts)
462                    .map_err(|e| {
463                        AuthzError::Database(format!("Invalid revoked_at timestamp: {e}"))
464                    })?
465                    .with_timezone(&Utc);
466                Ok(Some(dt))
467            }
468            None => Ok(None),
469        }
470    }
471
472    /// Check if a mandate is revoked (convenience method).
473    pub fn is_revoked(&self, mandate_id: &str) -> Result<bool, AuthzError> {
474        Ok(self.get_revoked_at(mandate_id)?.is_some())
475    }
476}
477
478/// Revocation record for upsert.
479#[derive(Debug, Clone)]
480pub struct RevocationRecord {
481    pub mandate_id: String,
482    pub revoked_at: DateTime<Utc>,
483    pub reason: Option<String>,
484    pub revoked_by: Option<String>,
485    pub source: Option<String>,
486    pub event_id: Option<String>,
487}
488
489/// Compute deterministic use_id per SPEC-Mandate-v1.0.4 §7.4.
490///
491/// ```text
492/// use_id = "sha256:" + hex(SHA256(mandate_id + ":" + tool_call_id + ":" + use_count))
493/// ```
494pub fn compute_use_id(mandate_id: &str, tool_call_id: &str, use_count: u32) -> String {
495    let input = format!("{}:{}:{}", mandate_id, tool_call_id, use_count);
496    let hash = Sha256::digest(input.as_bytes());
497    format!("sha256:{}", hex::encode(hash))
498}
499
500#[cfg(test)]
501mod tests {
502    use super::*;
503
504    fn test_metadata() -> MandateMetadata {
505        MandateMetadata {
506            mandate_id: "sha256:test123".to_string(),
507            mandate_kind: "intent".to_string(),
508            audience: "org/app".to_string(),
509            issuer: "auth.org.com".to_string(),
510            expires_at: None,
511            single_use: false,
512            max_uses: None,
513            canonical_digest: "sha256:digest123".to_string(),
514            key_id: "sha256:key123".to_string(),
515        }
516    }
517
518    fn consume(
519        store: &MandateStore,
520        mandate_id: &str,
521        tool_call_id: &str,
522        nonce: Option<&str>,
523        audience: &str,
524        issuer: &str,
525    ) -> Result<AuthzReceipt, AuthzError> {
526        store.consume_mandate(&ConsumeParams {
527            mandate_id,
528            tool_call_id,
529            nonce,
530            audience,
531            issuer,
532            tool_name: "test_tool",
533            operation_class: "read",
534            source_run_id: None,
535        })
536    }
537
538    // === A) Schema/migrations ===
539
540    #[test]
541    fn test_store_bootstraps_schema() {
542        let store = MandateStore::memory().unwrap();
543        let conn = store.conn.lock().unwrap();
544
545        // Check tables exist
546        let tables: Vec<String> = conn
547            .prepare("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
548            .unwrap()
549            .query_map([], |row| row.get(0))
550            .unwrap()
551            .filter_map(|r| r.ok())
552            .collect();
553
554        assert!(tables.contains(&"mandates".to_string()));
555        assert!(tables.contains(&"mandate_uses".to_string()));
556        assert!(tables.contains(&"nonces".to_string()));
557    }
558
559    #[test]
560    fn test_store_sets_foreign_keys() {
561        let store = MandateStore::memory().unwrap();
562        let conn = store.conn.lock().unwrap();
563
564        let fk: i32 = conn
565            .query_row("PRAGMA foreign_keys", [], |row| row.get(0))
566            .unwrap();
567        assert_eq!(fk, 1);
568    }
569
570    // === B) Upsert invariants ===
571
572    #[test]
573    fn test_upsert_mandate_inserts_new() {
574        let store = MandateStore::memory().unwrap();
575        let meta = test_metadata();
576
577        store.upsert_mandate(&meta).unwrap();
578
579        let count = store.get_use_count(&meta.mandate_id).unwrap();
580        assert_eq!(count, Some(0));
581    }
582
583    #[test]
584    fn test_upsert_mandate_is_noop_on_same_content() {
585        let store = MandateStore::memory().unwrap();
586        let meta = test_metadata();
587
588        store.upsert_mandate(&meta).unwrap();
589        store.upsert_mandate(&meta).unwrap(); // Should not error
590
591        let count = store.get_use_count(&meta.mandate_id).unwrap();
592        assert_eq!(count, Some(0));
593    }
594
595    #[test]
596    fn test_upsert_mandate_rejects_conflicting_metadata() {
597        let store = MandateStore::memory().unwrap();
598        let meta = test_metadata();
599        store.upsert_mandate(&meta).unwrap();
600
601        // Try to upsert with different audience
602        let mut conflict = meta.clone();
603        conflict.audience = "different/app".to_string();
604
605        let result = store.upsert_mandate(&conflict);
606        assert!(matches!(
607            result,
608            Err(AuthzError::MandateConflict { field, .. }) if field == "audience"
609        ));
610    }
611
612    // === C) Consume - idempotency & counts ===
613
614    #[test]
615    fn test_consume_fails_if_mandate_missing() {
616        let store = MandateStore::memory().unwrap();
617
618        let result = consume(
619            &store,
620            "sha256:nonexistent",
621            "tc_1",
622            None,
623            "org/app",
624            "auth.org.com",
625        );
626
627        assert!(matches!(result, Err(AuthzError::MandateNotFound { .. })));
628    }
629
630    #[test]
631    fn test_consume_first_time_returns_use_count_1() {
632        let store = MandateStore::memory().unwrap();
633        let meta = test_metadata();
634        store.upsert_mandate(&meta).unwrap();
635
636        let receipt = consume(
637            &store,
638            &meta.mandate_id,
639            "tc_1",
640            None,
641            &meta.audience,
642            &meta.issuer,
643        )
644        .unwrap();
645
646        assert_eq!(receipt.use_count, 1);
647        assert!(!receipt.use_id.is_empty());
648        assert_eq!(receipt.tool_call_id, "tc_1");
649        assert_eq!(store.get_use_count(&meta.mandate_id).unwrap(), Some(1));
650        assert_eq!(store.count_uses(&meta.mandate_id).unwrap(), 1);
651    }
652
653    #[test]
654    fn test_consume_is_idempotent_for_same_tool_call_id() {
655        let store = MandateStore::memory().unwrap();
656        let meta = test_metadata();
657        store.upsert_mandate(&meta).unwrap();
658
659        let receipt1 = consume(
660            &store,
661            &meta.mandate_id,
662            "tc_1",
663            None,
664            &meta.audience,
665            &meta.issuer,
666        )
667        .unwrap();
668        let receipt2 = consume(
669            &store,
670            &meta.mandate_id,
671            "tc_1",
672            None,
673            &meta.audience,
674            &meta.issuer,
675        )
676        .unwrap();
677
678        // Same receipt (idempotent)
679        assert_eq!(receipt1.use_id, receipt2.use_id);
680        assert_eq!(receipt1.use_count, receipt2.use_count);
681
682        // was_new distinguishes first vs retry
683        assert!(receipt1.was_new, "First consume should be was_new=true");
684        assert!(!receipt2.was_new, "Retry should be was_new=false");
685
686        // Count didn't increment
687        assert_eq!(store.get_use_count(&meta.mandate_id).unwrap(), Some(1));
688        assert_eq!(store.count_uses(&meta.mandate_id).unwrap(), 1);
689    }
690
691    #[test]
692    fn test_consume_increments_for_different_tool_call_ids() {
693        let store = MandateStore::memory().unwrap();
694        let meta = test_metadata();
695        store.upsert_mandate(&meta).unwrap();
696
697        let r1 = consume(
698            &store,
699            &meta.mandate_id,
700            "tc_1",
701            None,
702            &meta.audience,
703            &meta.issuer,
704        )
705        .unwrap();
706        let r2 = consume(
707            &store,
708            &meta.mandate_id,
709            "tc_2",
710            None,
711            &meta.audience,
712            &meta.issuer,
713        )
714        .unwrap();
715
716        assert_eq!(r1.use_count, 1);
717        assert_eq!(r2.use_count, 2);
718        assert_eq!(store.get_use_count(&meta.mandate_id).unwrap(), Some(2));
719        assert_eq!(store.count_uses(&meta.mandate_id).unwrap(), 2);
720    }
721
722    // === D) Constraints - single_use/max_uses ===
723
724    #[test]
725    fn test_single_use_allows_first_then_rejects_second() {
726        let store = MandateStore::memory().unwrap();
727        let mut meta = test_metadata();
728        meta.single_use = true;
729        meta.max_uses = Some(1);
730        store.upsert_mandate(&meta).unwrap();
731
732        let r1 = consume(
733            &store,
734            &meta.mandate_id,
735            "tc_1",
736            None,
737            &meta.audience,
738            &meta.issuer,
739        );
740        assert!(r1.is_ok());
741
742        let r2 = consume(
743            &store,
744            &meta.mandate_id,
745            "tc_2",
746            None,
747            &meta.audience,
748            &meta.issuer,
749        );
750        assert!(matches!(r2, Err(AuthzError::AlreadyUsed)));
751
752        // use_count stayed at 1 (rollback worked)
753        assert_eq!(store.get_use_count(&meta.mandate_id).unwrap(), Some(1));
754    }
755
756    #[test]
757    fn test_max_uses_allows_up_to_n_then_rejects() {
758        let store = MandateStore::memory().unwrap();
759        let mut meta = test_metadata();
760        meta.max_uses = Some(3);
761        store.upsert_mandate(&meta).unwrap();
762
763        for i in 1..=3 {
764            let r = consume(
765                &store,
766                &meta.mandate_id,
767                &format!("tc_{}", i),
768                None,
769                &meta.audience,
770                &meta.issuer,
771            );
772            assert!(r.is_ok(), "Call {} should succeed", i);
773        }
774
775        let r4 = consume(
776            &store,
777            &meta.mandate_id,
778            "tc_4",
779            None,
780            &meta.audience,
781            &meta.issuer,
782        );
783        assert!(matches!(
784            r4,
785            Err(AuthzError::MaxUsesExceeded { max: 3, current: 4 })
786        ));
787
788        // use_count stayed at 3
789        assert_eq!(store.get_use_count(&meta.mandate_id).unwrap(), Some(3));
790    }
791
792    #[test]
793    fn test_max_uses_null_is_unlimited() {
794        let store = MandateStore::memory().unwrap();
795        let meta = test_metadata(); // max_uses = None
796        store.upsert_mandate(&meta).unwrap();
797
798        for i in 1..=20 {
799            let r = consume(
800                &store,
801                &meta.mandate_id,
802                &format!("tc_{}", i),
803                None,
804                &meta.audience,
805                &meta.issuer,
806            );
807            assert!(r.is_ok(), "Call {} should succeed", i);
808        }
809
810        assert_eq!(store.get_use_count(&meta.mandate_id).unwrap(), Some(20));
811    }
812
813    #[test]
814    fn test_single_use_true_and_max_uses_gt_1_is_invalid() {
815        let store = MandateStore::memory().unwrap();
816        let mut meta = test_metadata();
817        meta.single_use = true;
818        meta.max_uses = Some(10); // Invalid: single_use with max > 1
819
820        let result = store.upsert_mandate(&meta);
821        assert!(matches!(
822            result,
823            Err(AuthzError::InvalidConstraints { max_uses: 10 })
824        ));
825    }
826
827    // === E) Nonce replay prevention ===
828
829    #[test]
830    fn test_nonce_inserted_on_first_consume() {
831        let store = MandateStore::memory().unwrap();
832        let meta = test_metadata();
833        store.upsert_mandate(&meta).unwrap();
834
835        consume(
836            &store,
837            &meta.mandate_id,
838            "tc_1",
839            Some("nonce_123"),
840            &meta.audience,
841            &meta.issuer,
842        )
843        .unwrap();
844
845        assert!(store
846            .nonce_exists(&meta.audience, &meta.issuer, "nonce_123")
847            .unwrap());
848    }
849
850    #[test]
851    fn test_nonce_replay_rejected() {
852        let store = MandateStore::memory().unwrap();
853        let meta = test_metadata();
854        store.upsert_mandate(&meta).unwrap();
855
856        let r1 = consume(
857            &store,
858            &meta.mandate_id,
859            "tc_1",
860            Some("nonce_123"),
861            &meta.audience,
862            &meta.issuer,
863        );
864        assert!(r1.is_ok());
865
866        // Different tool_call_id, same nonce
867        let r2 = consume(
868            &store,
869            &meta.mandate_id,
870            "tc_2",
871            Some("nonce_123"),
872            &meta.audience,
873            &meta.issuer,
874        );
875        assert!(matches!(r2, Err(AuthzError::NonceReplay { nonce }) if nonce == "nonce_123"));
876
877        // use_count stayed at 1 (rollback worked)
878        assert_eq!(store.get_use_count(&meta.mandate_id).unwrap(), Some(1));
879    }
880
881    #[test]
882    fn test_nonce_scope_is_audience_and_issuer() {
883        let store = MandateStore::memory().unwrap();
884
885        // Create two mandates with different audience
886        let meta1 = test_metadata();
887        let mut meta2 = test_metadata();
888        meta2.mandate_id = "sha256:test456".to_string();
889        meta2.audience = "different/app".to_string();
890        meta2.canonical_digest = "sha256:digest456".to_string();
891
892        store.upsert_mandate(&meta1).unwrap();
893        store.upsert_mandate(&meta2).unwrap();
894
895        // Same nonce, different audience should be allowed
896        let r1 = consume(
897            &store,
898            &meta1.mandate_id,
899            "tc_1",
900            Some("shared_nonce"),
901            &meta1.audience,
902            &meta1.issuer,
903        );
904        assert!(r1.is_ok());
905
906        // Same nonce but different audience
907        let r2 = consume(
908            &store,
909            &meta2.mandate_id,
910            "tc_2",
911            Some("shared_nonce"),
912            &meta2.audience,
913            &meta2.issuer,
914        );
915        assert!(
916            r2.is_ok(),
917            "Same nonce with different audience should be allowed"
918        );
919    }
920
921    // === F) Multi-call invariants (serialized via mutex) ===
922
923    #[test]
924    fn test_multicall_produces_monotonic_counts_no_gaps() {
925        let store = MandateStore::memory().unwrap();
926        let meta = test_metadata();
927        store.upsert_mandate(&meta).unwrap();
928
929        let mut counts = Vec::new();
930        for i in 1..=50 {
931            let r = consume(
932                &store,
933                &meta.mandate_id,
934                &format!("tc_{}", i),
935                None,
936                &meta.audience,
937                &meta.issuer,
938            )
939            .unwrap();
940            counts.push(r.use_count);
941        }
942
943        // Verify monotonic: 1, 2, 3, ..., 50
944        let expected: Vec<u32> = (1..=50).collect();
945        assert_eq!(counts, expected);
946        assert_eq!(store.get_use_count(&meta.mandate_id).unwrap(), Some(50));
947        assert_eq!(store.count_uses(&meta.mandate_id).unwrap(), 50);
948    }
949
950    #[test]
951    fn test_multicall_idempotent_same_tool_call_id() {
952        let store = MandateStore::memory().unwrap();
953        let meta = test_metadata();
954        store.upsert_mandate(&meta).unwrap();
955
956        let mut receipts = Vec::new();
957        for _ in 0..20 {
958            let r = consume(
959                &store,
960                &meta.mandate_id,
961                "tc_same",
962                None,
963                &meta.audience,
964                &meta.issuer,
965            )
966            .unwrap();
967            receipts.push(r);
968        }
969
970        // All receipts should be identical
971        let first = &receipts[0];
972        for r in &receipts {
973            assert_eq!(r.use_id, first.use_id);
974            assert_eq!(r.use_count, first.use_count);
975        }
976
977        // Only one actual use
978        assert_eq!(store.get_use_count(&meta.mandate_id).unwrap(), Some(1));
979        assert_eq!(store.count_uses(&meta.mandate_id).unwrap(), 1);
980    }
981
982    // === H) Revocation API ===
983
984    #[test]
985    fn test_revocation_roundtrip() {
986        let store = MandateStore::memory().unwrap();
987
988        let revoked_at = Utc::now();
989        let record = RevocationRecord {
990            mandate_id: "sha256:revoked123".to_string(),
991            revoked_at,
992            reason: Some("User requested".to_string()),
993            revoked_by: Some("admin@example.com".to_string()),
994            source: Some("assay://myorg/myapp".to_string()),
995            event_id: Some("evt_revoke_001".to_string()),
996        };
997
998        store.upsert_revocation(&record).unwrap();
999
1000        let got = store.get_revoked_at(&record.mandate_id).unwrap();
1001        assert!(got.is_some());
1002        // Compare within 1 second tolerance (RFC3339 loses sub-second precision)
1003        let diff = (got.unwrap() - revoked_at).num_seconds().abs();
1004        assert!(diff <= 1, "revoked_at timestamps differ by {}s", diff);
1005    }
1006
1007    #[test]
1008    fn test_revocation_is_revoked_helper() {
1009        let store = MandateStore::memory().unwrap();
1010
1011        assert!(!store.is_revoked("sha256:not_revoked").unwrap());
1012
1013        store
1014            .upsert_revocation(&RevocationRecord {
1015                mandate_id: "sha256:is_revoked".to_string(),
1016                revoked_at: Utc::now(),
1017                reason: None,
1018                revoked_by: None,
1019                source: None,
1020                event_id: None,
1021            })
1022            .unwrap();
1023
1024        assert!(store.is_revoked("sha256:is_revoked").unwrap());
1025    }
1026
1027    #[test]
1028    fn test_revocation_upsert_is_idempotent() {
1029        let store = MandateStore::memory().unwrap();
1030
1031        let record = RevocationRecord {
1032            mandate_id: "sha256:idem".to_string(),
1033            revoked_at: Utc::now(),
1034            reason: Some("First".to_string()),
1035            revoked_by: None,
1036            source: None,
1037            event_id: None,
1038        };
1039
1040        store.upsert_revocation(&record).unwrap();
1041        store.upsert_revocation(&record).unwrap(); // Should not fail
1042        store.upsert_revocation(&record).unwrap();
1043
1044        assert!(store.is_revoked(&record.mandate_id).unwrap());
1045    }
1046}
assay_core/runtime/mandate_store.rs

assay_core/runtime/
mandate_store.rs
