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