1use std::fs;
9use std::path::Path;
10
11use chio_core::capability::scope::MonetaryAmount;
12use chio_kernel::{ApprovalStoreError, BatchApproval, BatchApprovalStore};
13use r2d2::Pool;
14use r2d2_sqlite::SqliteConnectionManager;
15use rusqlite::{params, OptionalExtension};
16
17pub struct SqliteBatchApprovalStore {
18 pool: Pool<SqliteConnectionManager>,
19}
20
21const BATCH_APPROVAL_STORE_SUPPORTED_SCHEMA_VERSION: i32 = 0;
23const BATCH_APPROVAL_STORE_SCHEMA_KEY: &str = "batch_approval";
26const BATCH_APPROVAL_STORE_LEGACY_ANCHOR_TABLES: &[&str] = &["chio_hitl_batches"];
29
30impl SqliteBatchApprovalStore {
31 pub fn open(path: impl AsRef<Path>) -> Result<Self, ApprovalStoreError> {
32 let path = path.as_ref();
33 if let Some(parent) = crate::sqlite_parent_dir_to_create(path) {
37 fs::create_dir_all(&parent)
38 .map_err(|e| ApprovalStoreError::Backend(format!("create dir: {e}")))?;
39 }
40 let manager = SqliteConnectionManager::file(path);
41 let pool = Pool::builder()
42 .max_size(4)
43 .build(manager)
44 .map_err(|e| ApprovalStoreError::Backend(format!("pool build: {e}")))?;
45 let store = Self { pool };
46 store.run_migrations()?;
47 Ok(store)
48 }
49
50 pub fn open_in_memory() -> Result<Self, ApprovalStoreError> {
51 let manager = SqliteConnectionManager::memory();
52 let pool = Pool::builder()
53 .max_size(1)
54 .build(manager)
55 .map_err(|e| ApprovalStoreError::Backend(format!("pool build: {e}")))?;
56 let store = Self { pool };
57 store.run_migrations()?;
58 Ok(store)
59 }
60
61 fn run_migrations(&self) -> Result<(), ApprovalStoreError> {
62 let conn = self
63 .pool
64 .get()
65 .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
66 crate::check_schema_version(
67 &conn,
68 BATCH_APPROVAL_STORE_SCHEMA_KEY,
69 BATCH_APPROVAL_STORE_SUPPORTED_SCHEMA_VERSION,
70 BATCH_APPROVAL_STORE_LEGACY_ANCHOR_TABLES,
71 )
72 .map_err(|error| ApprovalStoreError::Backend(error.to_string()))?;
73 conn.execute_batch(
74 r#"
75 PRAGMA journal_mode = WAL;
76 PRAGMA synchronous = FULL;
77 PRAGMA busy_timeout = 5000;
78
79 CREATE TABLE IF NOT EXISTS chio_hitl_batches (
80 batch_id TEXT PRIMARY KEY,
81 approver_hex TEXT NOT NULL,
82 subject_id TEXT NOT NULL,
83 server_pattern TEXT NOT NULL,
84 tool_pattern TEXT NOT NULL,
85 per_call_currency TEXT,
86 per_call_units INTEGER,
87 total_currency TEXT,
88 total_units INTEGER,
89 max_calls INTEGER,
90 not_before INTEGER NOT NULL,
91 not_after INTEGER NOT NULL,
92 used_calls INTEGER NOT NULL DEFAULT 0,
93 used_total_units INTEGER NOT NULL DEFAULT 0,
94 revoked INTEGER NOT NULL DEFAULT 0
95 );
96 CREATE INDEX IF NOT EXISTS idx_chio_hitl_batches_subject
97 ON chio_hitl_batches(subject_id, revoked);
98 "#,
99 )
100 .map_err(|e| ApprovalStoreError::Backend(format!("migration: {e}")))?;
101 crate::stamp_schema_version(
102 &conn,
103 BATCH_APPROVAL_STORE_SCHEMA_KEY,
104 BATCH_APPROVAL_STORE_SUPPORTED_SCHEMA_VERSION,
105 )
106 .map_err(|error| ApprovalStoreError::Backend(error.to_string()))?;
107 Ok(())
108 }
109}
110
111fn row_to_batch(row: &rusqlite::Row<'_>) -> rusqlite::Result<BatchApproval> {
112 let per_call_currency: Option<String> = row.get(5)?;
113 let per_call_units: Option<i64> = row.get(6)?;
114 let total_currency: Option<String> = row.get(7)?;
115 let total_units: Option<i64> = row.get(8)?;
116 let max_calls: Option<i64> = row.get(9)?;
117 let revoked: i64 = row.get(14)?;
118 Ok(BatchApproval {
119 batch_id: row.get(0)?,
120 approver_hex: row.get(1)?,
121 subject_id: row.get(2)?,
122 server_pattern: row.get(3)?,
123 tool_pattern: row.get(4)?,
124 max_amount_per_call: match (per_call_currency, per_call_units) {
125 (Some(currency), Some(units)) => Some(MonetaryAmount {
126 currency,
127 units: units.max(0) as u64,
128 }),
129 _ => None,
130 },
131 max_total_amount: match (total_currency, total_units) {
132 (Some(currency), Some(units)) => Some(MonetaryAmount {
133 currency,
134 units: units.max(0) as u64,
135 }),
136 _ => None,
137 },
138 max_calls: max_calls.map(|v| v.max(0) as u32),
139 not_before: row.get::<_, i64>(10)?.max(0) as u64,
140 not_after: row.get::<_, i64>(11)?.max(0) as u64,
141 used_calls: row.get::<_, i64>(12)?.max(0) as u32,
142 used_total_units: row.get::<_, i64>(13)?.max(0) as u64,
143 revoked: revoked != 0,
144 })
145}
146
147fn pattern_matches(pattern: &str, value: &str) -> bool {
148 if pattern == "*" {
149 return true;
150 }
151 if let Some(prefix) = pattern.strip_suffix('*') {
152 return value.starts_with(prefix);
153 }
154 pattern == value
155}
156
157fn amount_fits(batch: &BatchApproval, amount: Option<&MonetaryAmount>) -> bool {
158 let Some(amt) = amount else {
159 return batch.max_amount_per_call.is_none() && batch.max_total_amount.is_none();
160 };
161 if let Some(per_call) = &batch.max_amount_per_call {
162 if amt.currency != per_call.currency || amt.units > per_call.units {
163 return false;
164 }
165 }
166 if let Some(total) = &batch.max_total_amount {
167 if amt.currency != total.currency
168 || batch.used_total_units.saturating_add(amt.units) > total.units
169 {
170 return false;
171 }
172 }
173 true
174}
175
176impl BatchApprovalStore for SqliteBatchApprovalStore {
177 fn store(&self, batch: &BatchApproval) -> Result<(), ApprovalStoreError> {
178 let conn = self
179 .pool
180 .get()
181 .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
182 conn.execute(
183 r#"INSERT OR REPLACE INTO chio_hitl_batches (
184 batch_id, approver_hex, subject_id, server_pattern, tool_pattern,
185 per_call_currency, per_call_units, total_currency, total_units,
186 max_calls, not_before, not_after, used_calls, used_total_units, revoked
187 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)"#,
188 params![
189 batch.batch_id,
190 batch.approver_hex,
191 batch.subject_id,
192 batch.server_pattern,
193 batch.tool_pattern,
194 batch
195 .max_amount_per_call
196 .as_ref()
197 .map(|a| a.currency.clone()),
198 batch.max_amount_per_call.as_ref().map(|a| a.units as i64),
199 batch.max_total_amount.as_ref().map(|a| a.currency.clone()),
200 batch.max_total_amount.as_ref().map(|a| a.units as i64),
201 batch.max_calls.map(|c| c as i64),
202 batch.not_before as i64,
203 batch.not_after as i64,
204 batch.used_calls as i64,
205 batch.used_total_units as i64,
206 batch.revoked as i64,
207 ],
208 )
209 .map_err(|e| ApprovalStoreError::Backend(format!("insert batch: {e}")))?;
210 Ok(())
211 }
212
213 fn find_matching(
214 &self,
215 subject_id: &str,
216 server_id: &str,
217 tool_name: &str,
218 amount: Option<&MonetaryAmount>,
219 now: u64,
220 ) -> Result<Option<BatchApproval>, ApprovalStoreError> {
221 let conn = self
222 .pool
223 .get()
224 .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
225 let mut stmt = conn
228 .prepare(
229 r#"SELECT batch_id, approver_hex, subject_id, server_pattern, tool_pattern,
230 per_call_currency, per_call_units, total_currency, total_units,
231 max_calls, not_before, not_after, used_calls, used_total_units, revoked
232 FROM chio_hitl_batches
233 WHERE subject_id = ?1 AND revoked = 0
234 AND not_before <= ?2 AND not_after > ?2"#,
235 )
236 .map_err(|e| ApprovalStoreError::Backend(format!("prepare: {e}")))?;
237 let rows = stmt
238 .query_map(params![subject_id, now as i64], row_to_batch)
239 .map_err(|e| ApprovalStoreError::Backend(format!("query: {e}")))?;
240 for row in rows {
241 let batch = row.map_err(|e| ApprovalStoreError::Backend(format!("row: {e}")))?;
242 if !pattern_matches(&batch.server_pattern, server_id) {
243 continue;
244 }
245 if !pattern_matches(&batch.tool_pattern, tool_name) {
246 continue;
247 }
248 if let Some(max) = batch.max_calls {
249 if batch.used_calls >= max {
250 continue;
251 }
252 }
253 if !amount_fits(&batch, amount) {
254 continue;
255 }
256 return Ok(Some(batch));
257 }
258 Ok(None)
259 }
260
261 fn record_usage(
262 &self,
263 batch_id: &str,
264 amount: Option<&MonetaryAmount>,
265 ) -> Result<(), ApprovalStoreError> {
266 let conn = self
267 .pool
268 .get()
269 .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
270 let added_units = amount.map(|a| a.units as i64).unwrap_or(0);
271 let rows = conn.execute(
272 "UPDATE chio_hitl_batches SET used_calls = used_calls + 1, used_total_units = used_total_units + ?2 WHERE batch_id = ?1",
273 params![batch_id, added_units],
274 )
275 .map_err(|e| ApprovalStoreError::Backend(format!("update usage: {e}")))?;
276 if rows == 0 {
277 return Err(ApprovalStoreError::NotFound(batch_id.to_string()));
278 }
279 Ok(())
280 }
281
282 fn revoke(&self, batch_id: &str) -> Result<(), ApprovalStoreError> {
283 let conn = self
284 .pool
285 .get()
286 .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
287 let rows = conn
288 .execute(
289 "UPDATE chio_hitl_batches SET revoked = 1 WHERE batch_id = ?1",
290 params![batch_id],
291 )
292 .map_err(|e| ApprovalStoreError::Backend(format!("revoke: {e}")))?;
293 if rows == 0 {
294 return Err(ApprovalStoreError::NotFound(batch_id.to_string()));
295 }
296 Ok(())
297 }
298
299 fn get(&self, batch_id: &str) -> Result<Option<BatchApproval>, ApprovalStoreError> {
300 let conn = self
301 .pool
302 .get()
303 .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
304 let batch: Option<BatchApproval> = conn
305 .query_row(
306 r#"SELECT batch_id, approver_hex, subject_id, server_pattern, tool_pattern,
307 per_call_currency, per_call_units, total_currency, total_units,
308 max_calls, not_before, not_after, used_calls, used_total_units, revoked
309 FROM chio_hitl_batches WHERE batch_id = ?1"#,
310 params![batch_id],
311 row_to_batch,
312 )
313 .optional()
314 .map_err(|e| ApprovalStoreError::Backend(format!("get: {e}")))?;
315 Ok(batch)
316 }
317}