1#![deny(missing_docs)]
31
32pub(crate) mod instant;
33
34#[cfg(feature = "test-support")]
35pub mod test_support;
36
37use instant::StoredInstant;
38use std::num::{NonZeroI64, NonZeroUsize};
39use std::sync::Arc;
40
41use async_trait::async_trait;
42use jiff::{SignedDuration, Timestamp};
43use serde::{Deserialize, Serialize};
44use sqlx::postgres::PgPoolOptions;
45use sqlx::{PgPool, Postgres, Row, Transaction};
46use tokio::sync::broadcast;
47
48use tollgate_core::{
49 AccountId, AccountSnapshot, AccountStatus, BudgetSchedule, BudgetView, CapacityClass,
50 CostTable, CostUnits, EnforcementMode, FencingToken, Generation, KeyId, LeaseGrant, LeaseId,
51 Period, PermissionBits, PolicyRevision, Principal, PublishableSnapshot, ResolvedLimits,
52 Rollover, UsageEvent,
53};
54use tollgate_store::{
55 AccountConfig, AccountView, AdminAuthority, AdminReceipt, AdminState, AdminStore,
56 AllocateError, Allocation, BudgetError, Conservation, CreateAccountError, GrantPolicy,
57 IngestError, IngestReport, KeyDirectory, KeyError, KeyRecord, KeySnapshotError, KeySummary,
58 LeaseAllocator, PUSH_CHANNEL_CAPACITY, PublishSnapshotError, ReclaimBatch, ReclaimedLease,
59 Revocation, RolledAccount, RolloverBatch, SetStatusError, SnapshotPush, SnapshotResolution,
60 SnapshotSource, StatusChange, StoreError, StoreHealth, UsageSink, pushes_exceed_capacity,
61 validate_key_page_limit,
62};
63
64const STATE_ACTIVE: i16 = 0;
65const STATE_RELEASED: i16 = 1;
66const STATE_EXPIRED: i16 = 2;
67
68#[derive(Debug, Clone, Copy)]
72struct StoredId(u128);
73
74impl Serialize for StoredId {
75 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
76 where
77 S: serde::Serializer,
78 {
79 match u64::try_from(self.0) {
80 Ok(value) => serializer.serialize_u64(value),
81 Err(_) => serializer.collect_str(&format_args!("{:032x}", self.0)),
82 }
83 }
84}
85
86impl<'de> Deserialize<'de> for StoredId {
87 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
88 where
89 D: serde::Deserializer<'de>,
90 {
91 struct StoredIdVisitor;
92
93 impl serde::de::Visitor<'_> for StoredIdVisitor {
94 type Value = StoredId;
95
96 fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97 formatter.write_str("a legacy u64 number or a canonical 128-bit identifier string")
98 }
99
100 fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
101 Ok(StoredId(u128::from(value)))
102 }
103
104 fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
105 where
106 E: serde::de::Error,
107 {
108 value
109 .parse::<AccountId>()
110 .map(|id| StoredId(id.0))
111 .map_err(E::custom)
112 }
113 }
114
115 deserializer.deserialize_any(StoredIdVisitor)
116 }
117}
118
119fn digest_from(bytes: &[u8], key_id: KeyId) -> Result<[u8; 32], StoreError> {
124 <[u8; 32]>::try_from(bytes).map_err(|_| {
125 StoreError(format!(
126 "credential {key_id} has a {}-byte digest in storage, expected 32",
127 bytes.len()
128 ))
129 })
130}
131
132#[async_trait]
133impl KeyDirectory for PostgresStore {
134 async fn credential_activity(
135 &self,
136 keys: &[KeyId],
137 ) -> Result<Vec<tollgate_store::CredentialActivity>, StoreError> {
138 use tollgate_store::{CredentialActivity, CredentialActivityState};
139 let mut result = Vec::with_capacity(keys.len());
140 for chunk in keys.chunks(tollgate_store::MAX_INGEST_BATCH) {
143 let ids: Vec<_> = chunk.iter().map(|id| id_bytes(id.0)).collect();
144 let rows = sqlx::query(
145 "SELECT evidence.key_id IS NOT NULL, evidence.last_committed_at_us
149 FROM UNNEST($1::bytea[]) WITH ORDINALITY AS requested(key_id, ordinal)
150 LEFT JOIN LATERAL (
151 SELECT k.key_id, a.last_committed_at_us
152 FROM tollgate_credential_keys k
153 LEFT JOIN tollgate_credential_activity a ON a.key_id = k.key_id
154 WHERE k.key_id = requested.key_id LIMIT 1
155 ) AS evidence ON true
156 ORDER BY requested.ordinal",
157 )
158 .bind(&ids)
159 .fetch_all(&self.pool)
160 .await
161 .map_err(storage)?;
162 if rows.len() != chunk.len() {
163 return Err(StoreError("incomplete credential activity read".into()));
164 }
165 for (&key_id, row) in chunk.iter().zip(rows) {
166 let state = if !row.get::<bool, _>(0) {
167 CredentialActivityState::Unknown
168 } else if let Some(at) = row.get::<Option<i64>, _>(1) {
169 CredentialActivityState::Committed {
170 last_committed_at: micros_ts(at, "credential activity")?,
171 }
172 } else {
173 CredentialActivityState::Unobserved
174 };
175 result.push(CredentialActivity { key_id, state });
176 }
177 }
178 Ok(result)
179 }
180
181 async fn insert_key(&self, record: KeyRecord) -> Result<(), KeyError> {
182 self.insert_credential(record, None)
183 .await
184 .map(|receipt| receipt.outcome)
185 }
186
187 async fn revoke_key(&self, key_id: KeyId, now: Timestamp) -> Result<Revocation, KeyError> {
188 self.revoke_key_audited(key_id, now)
189 .await
190 .map(|receipt| receipt.outcome)
191 }
192
193 async fn revoke_key_audited(
194 &self,
195 key_id: KeyId,
196 now: Timestamp,
197 ) -> Result<AdminReceipt<Revocation>, KeyError> {
198 let mut tx = self.pool.begin().await.map_err(storage)?;
199 let result = async {
200 let row = sqlx::query(
204 "SELECT account_id, revoked_at_us IS NOT NULL
205 FROM tollgate_credential_keys WHERE key_id = $1 FOR UPDATE",
206 )
207 .bind(id_bytes(key_id.0))
208 .fetch_optional(&mut *tx)
209 .await
210 .map_err(storage)?
211 .ok_or(KeyError::UnknownKey)?;
212 let account_bytes: Vec<u8> = row.get(0);
213 let account_bytes: [u8; 16] = account_bytes
214 .try_into()
215 .map_err(|_| StoreError("credential account identifier is not 16 bytes".into()))?;
216 let account_id = AccountId(u128::from_be_bytes(account_bytes));
217 let revoked: bool = row.get(1);
218 let before = AdminState::Credential {
219 account_id,
220 key_id,
221 revoked,
222 };
223 let outcome = if revoked {
224 Revocation::AlreadyRetired
225 } else {
226 sqlx::query(
227 "UPDATE tollgate_credential_keys SET revoked_at_us = $2 WHERE key_id = $1",
228 )
229 .bind(id_bytes(key_id.0))
230 .bind(ts_micros(now))
231 .execute(&mut *tx)
232 .await
233 .map_err(storage)?;
234 Revocation::Retired
235 };
236 Ok(AdminReceipt::new(
237 outcome,
238 before,
239 AdminState::Credential {
240 account_id,
241 key_id,
242 revoked: true,
243 },
244 ))
245 }
246 .await;
247 finish_transaction(tx, result).await
248 }
249
250 async fn publish_key_snapshot(
251 &self,
252 account: AccountId,
253 key: KeyId,
254 snapshot: PublishableSnapshot,
255 ) -> Result<AdminReceipt<()>, KeySnapshotError> {
256 self.publish_key_snapshot_with_generation(account, key, snapshot, false)
257 .await
258 }
259
260 async fn publish_key_snapshot_next(
261 &self,
262 account: AccountId,
263 key: KeyId,
264 snapshot: PublishableSnapshot,
265 ) -> Result<AdminReceipt<()>, KeySnapshotError> {
266 self.publish_key_snapshot_with_generation(account, key, snapshot, true)
267 .await
268 }
269
270 async fn remove_key_snapshot(
271 &self,
272 account: AccountId,
273 key: KeyId,
274 ) -> Result<AdminReceipt<()>, KeySnapshotError> {
275 let mut tx = self.pool.begin().await.map_err(storage)?;
276 let result = async {
277 let (principal, _revoked) = lock_account_key(&mut tx, account, key).await?;
278 Ok::<_, KeySnapshotError>((principal, remove_in_tx(&mut tx, principal).await?))
279 }
280 .await;
281 let (principal, receipt) = finish_transaction(tx, result).await?;
282 self.announce_removal(principal, &receipt);
283 Ok(receipt)
284 }
285
286 async fn active_keys(&self, now: Timestamp) -> Result<Vec<KeyRecord>, StoreError> {
287 let cutoff = StoredInstant::from(now);
288 let rows = sqlx::query(
294 "SELECT key_id, account_id, principal, digest, not_after_floor_us, not_after_submicro_ns, not_after_is_lower_bound
295 FROM tollgate_credential_keys
296 WHERE revoked_at_us IS NULL
297 AND (not_after_floor_us IS NULL OR not_after_submicro_ns IS NULL
298 OR (not_after_floor_us, not_after_submicro_ns) > ($1, $2))
299 ORDER BY key_id",
300 )
301 .bind(cutoff.micros)
302 .bind(cutoff.submicro_nanos)
303 .fetch_all(&self.pool)
304 .await
305 .map_err(storage)?;
306
307 rows.into_iter().map(credential_from_row).collect()
308 }
309
310 async fn account_keys(
311 &self,
312 account: AccountId,
313 after: Option<KeyId>,
314 limit: NonZeroUsize,
315 ) -> Result<Vec<KeySummary>, KeyError> {
316 validate_key_page_limit(limit)?;
317 let sql = if after.is_some() {
322 "SELECT page.* FROM tollgate_accounts AS account
323 LEFT JOIN LATERAL (
324 SELECT key_id, not_after_floor_us, not_after_submicro_ns,
325 not_after_is_lower_bound, revoked_at_us
326 FROM tollgate_credential_keys
327 WHERE account_id = account.account_id AND key_id > $3
328 ORDER BY key_id LIMIT $2
329 ) AS page ON TRUE
330 WHERE account.account_id = $1 ORDER BY page.key_id"
331 } else {
332 "SELECT page.* FROM tollgate_accounts AS account
333 LEFT JOIN LATERAL (
334 SELECT key_id, not_after_floor_us, not_after_submicro_ns,
335 not_after_is_lower_bound, revoked_at_us
336 FROM tollgate_credential_keys
337 WHERE account_id = account.account_id
338 ORDER BY key_id LIMIT $2
339 ) AS page ON TRUE
340 WHERE account.account_id = $1 ORDER BY page.key_id"
341 };
342 let mut query = sqlx::query(sql)
343 .bind(id_bytes(account.0))
344 .bind(i64::try_from(limit.get()).unwrap_or(i64::MAX));
345 if let Some(cursor) = after {
346 query = query.bind(id_bytes(cursor.0));
347 }
348 let rows = query.fetch_all(&self.pool).await.map_err(storage)?;
349 if rows.is_empty() {
350 return Err(KeyError::UnknownAccount);
351 }
352 rows.into_iter()
353 .filter(|row| row.get::<Option<&[u8]>, _>(0).is_some())
354 .map(|row| summary_from_row(row).map_err(KeyError::Storage))
355 .collect()
356 }
357
358 async fn insert_key_within(
359 &self,
360 record: KeyRecord,
361 max_active: NonZeroUsize,
362 now: Timestamp,
363 ) -> Result<(), KeyError> {
364 self.insert_key_within_audited(record, max_active, now)
365 .await
366 .map(|receipt| receipt.outcome)
367 }
368
369 async fn insert_key_within_audited(
370 &self,
371 record: KeyRecord,
372 max_active: NonZeroUsize,
373 now: Timestamp,
374 ) -> Result<AdminReceipt<()>, KeyError> {
375 self.insert_credential(record, Some((max_active, now)))
376 .await
377 }
378}
379
380impl PostgresStore {
381 async fn publish_key_snapshot_with_generation(
382 &self,
383 account: AccountId,
384 key: KeyId,
385 snapshot: PublishableSnapshot,
386 allocate_generation: bool,
387 ) -> Result<AdminReceipt<()>, KeySnapshotError> {
388 let generation = if allocate_generation {
389 None
390 } else {
391 Some(i64::try_from(snapshot.generation.0).map_err(|_| {
392 StoreError("snapshot generation exceeds PostgreSQL BIGINT range".into())
393 })?)
394 };
395 let mut tx = self.pool.begin().await.map_err(storage)?;
396 let result = async {
397 let (principal, revoked) = lock_account_key(&mut tx, account, key).await?;
398 if revoked {
399 return Err(KeySnapshotError::Retired { key_id: key });
400 }
401 if snapshot.key_id != Some(key) {
402 return Err(PublishSnapshotError::CredentialMismatch { key_id: key }.into());
403 }
404 let published = publish_in_tx(&mut tx, principal, generation, snapshot).await?;
405 Ok((principal, published))
406 }
407 .await;
408 let (principal, (written, published, before, after)) =
409 finish_transaction(tx, result).await?;
410 if written {
411 self.push_to_subscribers(SnapshotPush {
412 principal,
413 resolution: SnapshotResolution::Present(published),
414 });
415 }
416 Ok(AdminReceipt::new((), before, after))
417 }
418
419 async fn create_account_as(
422 &self,
423 config: AccountConfig,
424 origin: AdminAuthority,
425 ) -> Result<tollgate_store::AdminReceipt<()>, CreateAccountError> {
426 let result = sqlx::query(
427 "INSERT INTO tollgate_accounts
428 (account_id, balance, deposited, status, capacity_class, next_fence,
429 usage_recorded, settlement_loss, overage_recorded, origin, status_set_by)
430 VALUES ($1, $2, $2, $3, $4, 1, 0, 0, 0, $5, $5)
431 ON CONFLICT (account_id) DO NOTHING",
432 )
433 .bind(id_bytes(config.account_id.0))
434 .bind(to_i64(config.initial_balance, "balance").map_err(CreateAccountError::Storage)?)
435 .bind(config.status.as_str())
436 .bind(config.capacity_class.as_str())
437 .bind(origin.as_str())
438 .execute(&self.pool)
439 .await
440 .map_err(|e| CreateAccountError::Storage(storage(e)))?;
441 if result.rows_affected() == 0 {
442 return Err(CreateAccountError::AlreadyExists);
443 }
444 Ok(AdminReceipt::new(
445 (),
446 AdminState::Absent,
447 AdminState::AccountCreated {
448 initial_balance: config.initial_balance,
449 status: config.status,
450 capacity_class: config.capacity_class,
451 origin,
452 },
453 ))
454 }
455
456 async fn set_status_as(
460 &self,
461 account: AccountId,
462 status: AccountStatus,
463 authority: AdminAuthority,
464 ) -> Result<tollgate_store::AdminReceipt<StatusChange>, SetStatusError> {
465 let mut tx = self.pool.begin().await.map_err(storage)?;
466 let result = async {
467 let row = sqlx::query(
472 "SELECT status, origin, status_set_by
473 FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE",
474 )
475 .bind(id_bytes(account.0))
476 .fetch_optional(&mut *tx)
477 .await
478 .map_err(storage)?
479 .ok_or(SetStatusError::UnknownAccount)?;
480
481 let before = decode_status(row.get::<String, _>(0))?;
482 let origin = decode_authority(row.get::<String, _>(1))?;
483 let before_set_by = decode_authority(row.get::<String, _>(2))?;
484 if authority == AdminAuthority::Provisioner && origin != AdminAuthority::Provisioner {
485 return Err(SetStatusError::NotProvisioned);
486 }
487 if before == AccountStatus::Closed && status != AccountStatus::Closed {
488 return Err(SetStatusError::AccountClosed);
491 }
492 if authority == AdminAuthority::Provisioner
493 && before != status
494 && before_set_by == AdminAuthority::Operator
495 {
496 return Err(SetStatusError::OperatorHold);
497 }
498 let set_by = if authority == AdminAuthority::Provisioner && before == status {
501 before_set_by
502 } else {
503 authority
504 };
505
506 sqlx::query(
507 "UPDATE tollgate_accounts SET status = $2, status_set_by = $3 WHERE account_id = $1",
508 )
509 .bind(id_bytes(account.0))
510 .bind(status.as_str())
511 .bind(set_by.as_str())
512 .execute(&mut *tx)
513 .await
514 .map_err(storage)?;
515
516 let (republished, unreadable) =
517 republish_patched_snapshots(&mut tx, account, "{status}", status.as_str()).await?;
518 Ok(((before, before_set_by, set_by), republished, unreadable))
519 }
520 .await;
521
522 let ((before, before_set_by, set_by), republished, unreadable) =
523 finish_transaction(tx, result).await?;
524 if pushes_exceed_capacity(republished.len()) {
526 tracing::warn!(
527 %account,
528 principals = republished.len(),
529 capacity = PUSH_CHANNEL_CAPACITY,
530 "status change emitted more pushes than the channel holds; subscribers will resync"
531 );
532 }
533 let count = republished.len();
534 for (principal, snapshot) in republished {
535 self.push_to_subscribers(SnapshotPush {
536 principal,
537 resolution: SnapshotResolution::Present(snapshot),
538 });
539 }
540 Ok(AdminReceipt::new(
541 StatusChange {
542 republished: count + unreadable,
545 unreadable,
546 },
547 AdminState::Status {
548 status: before,
549 set_by: before_set_by,
550 },
551 AdminState::Status { status, set_by },
552 ))
553 }
554
555 async fn insert_credential(
560 &self,
561 record: KeyRecord,
562 bound: Option<(NonZeroUsize, Timestamp)>,
563 ) -> Result<AdminReceipt<()>, KeyError> {
564 let mut tx = self
565 .pool
566 .begin()
567 .await
568 .map_err(|e| KeyError::Storage(storage(e)))?;
569 let account = sqlx::query(ACCOUNT_LOCK_SQL)
577 .bind(id_bytes(record.account_id.0))
578 .fetch_optional(&mut *tx)
579 .await
580 .map_err(|e| KeyError::Storage(storage(e)))?;
581 if account.is_none() {
582 return Err(KeyError::UnknownAccount);
583 }
584
585 if let Some((max_active, now)) = bound {
586 let existing: Option<i32> = sqlx::query_scalar(
593 "SELECT 1 FROM tollgate_credential_keys WHERE key_id = $1 OR principal = $2",
594 )
595 .bind(id_bytes(record.key_id.0))
596 .bind(id_bytes(record.principal.0))
597 .fetch_optional(&mut *tx)
598 .await
599 .map_err(|e| KeyError::Storage(storage(e)))?;
600 if existing.is_some() {
601 return Err(KeyError::AlreadyExists);
602 }
603
604 let cutoff = StoredInstant::from(now);
605 let live: i64 = sqlx::query_scalar(LIVE_KEY_COUNT_SQL)
606 .bind(id_bytes(record.account_id.0))
607 .bind(cutoff.micros)
608 .bind(cutoff.submicro_nanos)
609 .fetch_one(&mut *tx)
610 .await
611 .map_err(|e| KeyError::Storage(storage(e)))?;
612 if u128::from(live.max(0).unsigned_abs())
613 >= u128::try_from(max_active.get()).unwrap_or(u128::MAX)
614 {
615 return Err(KeyError::ActiveKeyLimit { limit: max_active });
616 }
617 }
618
619 let expiry = record.not_after.map(StoredInstant::from);
620 let result = sqlx::query(
621 "INSERT INTO tollgate_credential_keys
622 (key_id, account_id, principal, digest, not_after_floor_us,
623 not_after_submicro_ns, not_after_is_lower_bound, revoked_at_us)
624 VALUES ($1, $2, $3, $4, $5, $6, FALSE, NULL)
625 ON CONFLICT (key_id) DO NOTHING",
626 )
627 .bind(id_bytes(record.key_id.0))
628 .bind(id_bytes(record.account_id.0))
629 .bind(id_bytes(record.principal.0))
630 .bind(record.digest.to_vec())
631 .bind(expiry.map(|expiry| expiry.micros))
632 .bind(expiry.map(|expiry| expiry.submicro_nanos))
633 .execute(&mut *tx)
634 .await;
635 let outcome = match result {
636 Ok(done) if done.rows_affected() == 0 => Err(KeyError::AlreadyExists),
637 Ok(_) => Ok(()),
638 Err(sqlx::Error::Database(e)) if e.is_unique_violation() => {
639 Err(KeyError::AlreadyExists)
640 }
641 Err(e) => Err(KeyError::Storage(storage(e))),
651 };
652 outcome?;
653 tx.commit()
654 .await
655 .map_err(|e| KeyError::Storage(storage(e)))?;
656 Ok(AdminReceipt::new(
657 (),
658 AdminState::Absent,
659 AdminState::Credential {
660 account_id: record.account_id,
661 key_id: record.key_id,
662 revoked: false,
663 },
664 ))
665 }
666}
667
668fn summary_from_row(row: sqlx::postgres::PgRow) -> Result<KeySummary, StoreError> {
669 let bytes: Vec<u8> = row.get(0);
670 let fixed: [u8; 16] = bytes
671 .as_slice()
672 .try_into()
673 .map_err(|_| StoreError("credential identifier is not 16 bytes".into()))?;
674 let not_after = match (row.get::<Option<i64>, _>(1), row.get::<Option<i16>, _>(2)) {
675 (None, None) if !row.get::<bool, _>(3) => None,
676 (Some(micros), Some(submicro_nanos)) => Some(
677 StoredInstant {
678 micros,
679 submicro_nanos,
680 }
681 .timestamp()?,
682 ),
683 _ => return Err(StoreError("incomplete stored credential expiry".into())),
684 };
685 Ok(KeySummary {
686 key_id: KeyId(u128::from_be_bytes(fixed)),
687 not_after,
688 revoked_at: row
689 .get::<Option<i64>, _>(4)
690 .map(|micros| {
691 StoredInstant {
692 micros,
693 submicro_nanos: 0,
694 }
695 .timestamp()
696 })
697 .transpose()?,
698 })
699}
700
701fn credential_from_row(row: sqlx::postgres::PgRow) -> Result<KeyRecord, StoreError> {
702 let id = |index| -> Result<u128, StoreError> {
703 let bytes: Vec<u8> = row.get(index);
704 let fixed: [u8; 16] = bytes
705 .as_slice()
706 .try_into()
707 .map_err(|_| StoreError("credential identifier is not 16 bytes".into()))?;
708 Ok(u128::from_be_bytes(fixed))
709 };
710 let key_id = KeyId(id(0)?);
711 let not_after = match (row.get::<Option<i64>, _>(4), row.get::<Option<i16>, _>(5)) {
712 (None, None) if !row.get::<bool, _>(6) => None,
713 (Some(micros), Some(submicro_nanos)) => Some(
714 StoredInstant {
715 micros,
716 submicro_nanos,
717 }
718 .timestamp()?,
719 ),
720 _ => return Err(StoreError("incomplete stored credential expiry".into())),
721 };
722 Ok(KeyRecord {
723 key_id,
724 account_id: AccountId(id(1)?),
725 principal: Principal(id(2)?),
726 digest: digest_from(row.get::<Vec<u8>, _>(3).as_slice(), key_id)?,
727 not_after,
728 })
729}
730
731#[async_trait]
732impl tollgate_store::KeySource for PostgresStore {
733 async fn active_keys_page(
734 &self,
735 now: Timestamp,
736 after: Option<KeyId>,
737 limit: NonZeroUsize,
738 ) -> Result<tollgate_store::KeyPage, StoreError> {
739 validate_key_page_limit(limit)?;
740 let cutoff = StoredInstant::from(now);
741 let mut tx = self.pool.begin().await.map_err(storage)?;
744 sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY")
745 .execute(&mut *tx)
746 .await
747 .map_err(storage)?;
748 let revision: i64 =
749 sqlx::query_scalar("SELECT revision FROM tollgate_credential_revision WHERE singleton")
750 .fetch_one(&mut *tx)
751 .await
752 .map_err(storage)?;
753 let revision = u64::try_from(revision)
754 .map_err(|_| StoreError("credential revision is negative".into()))?;
755 let sql = if after.is_some() {
757 "SELECT key_id, account_id, principal, digest, not_after_floor_us, not_after_submicro_ns, not_after_is_lower_bound
758 FROM tollgate_credential_keys WHERE revoked_at_us IS NULL
759 AND (not_after_floor_us IS NULL OR not_after_submicro_ns IS NULL
760 OR (not_after_floor_us, not_after_submicro_ns) > ($1, $2)) AND key_id > $3
761 ORDER BY key_id LIMIT $4"
762 } else {
763 "SELECT key_id, account_id, principal, digest, not_after_floor_us, not_after_submicro_ns, not_after_is_lower_bound
764 FROM tollgate_credential_keys WHERE revoked_at_us IS NULL
765 AND (not_after_floor_us IS NULL OR not_after_submicro_ns IS NULL
766 OR (not_after_floor_us, not_after_submicro_ns) > ($1, $2)) AND key_id >= $3
767 ORDER BY key_id LIMIT $4"
768 };
769 let rows = sqlx::query(sql)
770 .bind(cutoff.micros)
771 .bind(cutoff.submicro_nanos)
772 .bind(id_bytes(after.unwrap_or(KeyId(0)).0))
773 .bind((limit.get() + 1) as i64)
774 .fetch_all(&mut *tx)
775 .await
776 .map_err(storage)?;
777 tx.commit().await.map_err(storage)?;
778 let mut records = rows
781 .into_iter()
782 .map(credential_from_row)
783 .map(|row| row.map(tollgate_store::CredentialRecord::from))
784 .collect::<Result<Vec<_>, _>>()?;
785 let next_after = if records.len() > limit.get() {
786 records.pop();
787 records.last().map(|record| record.key_id)
788 } else {
789 None
790 };
791 tollgate_store::KeyPage::try_new(revision, now, after, limit, records, next_after)
792 }
793}
794
795#[cfg(test)]
796mod stored_id_tests {
797 use super::StoredId;
798
799 #[test]
800 fn malformed_storage_id_explains_both_accepted_representations() {
801 let error = serde_json::from_value::<StoredId>(serde_json::Value::Bool(true)).unwrap_err();
802 assert!(
803 error
804 .to_string()
805 .contains("a legacy u64 number or a canonical 128-bit identifier string"),
806 "unexpected diagnostic: {error}"
807 );
808 }
809}
810
811#[derive(Serialize)]
827struct StoredSnapshotRef<'a> {
828 account_id: StoredId,
829 key_id: Option<StoredId>,
830 status: &'a AccountStatus,
831 capacity_class: &'a CapacityClass,
834 enforcement_mode: &'a EnforcementMode,
835 valid_until: &'a Timestamp,
836 permissions: &'a PermissionBits,
837 limits: &'a ResolvedLimits,
838 cost_table: &'a Arc<CostTable>,
839 budget: Option<&'a BudgetView>,
843 policy_revision: &'a PolicyRevision,
848}
849
850impl<'a> From<&'a AccountSnapshot> for StoredSnapshotRef<'a> {
851 fn from(snapshot: &'a AccountSnapshot) -> Self {
852 StoredSnapshotRef {
853 account_id: StoredId(snapshot.account_id.0),
854 key_id: snapshot.key_id.map(|id| StoredId(id.0)),
855 status: &snapshot.status,
856 capacity_class: &snapshot.capacity_class,
857 enforcement_mode: &snapshot.enforcement_mode,
858 valid_until: &snapshot.valid_until,
859 permissions: &snapshot.permissions,
860 limits: &snapshot.limits,
861 cost_table: &snapshot.cost_table,
862 budget: snapshot.budget.as_ref(),
863 policy_revision: &snapshot.policy_revision,
864 }
865 }
866}
867
868#[derive(Deserialize)]
872struct StoredSnapshot {
873 account_id: StoredId,
874 key_id: Option<StoredId>,
875 status: AccountStatus,
876 #[serde(default)]
882 capacity_class: CapacityClass,
883 #[serde(default)]
894 enforcement_mode: EnforcementMode,
895 valid_until: Timestamp,
896 permissions: PermissionBits,
897 limits: ResolvedLimits,
898 cost_table: Arc<CostTable>,
899 #[serde(default)]
905 budget: Option<BudgetView>,
906 #[serde(default)]
912 policy_revision: PolicyRevision,
913}
914
915impl StoredSnapshot {
916 fn into_snapshot(self, generation: Generation) -> AccountSnapshot {
922 let builder = AccountSnapshot::builder(
923 AccountId(self.account_id.0),
924 generation,
925 self.status,
926 self.valid_until,
927 self.permissions,
928 self.limits,
929 self.cost_table,
930 )
931 .enforcement_mode(self.enforcement_mode)
932 .capacity_class(self.capacity_class)
933 .policy_revision(self.policy_revision);
934 match self.key_id {
935 Some(key_id) => builder.key_id(tollgate_core::KeyId(key_id.0)).build(),
936 None => builder.build(),
937 }
938 }
939}
940
941const ACTIVE_LEASE_SUM_SQL: &str =
959 "SELECT COALESCE(SUM(granted), 0)::BIGINT, COALESCE(SUM(used), 0)::BIGINT
960 FROM tollgate_leases WHERE account_id = $1 AND state = 0";
961
962pub(crate) const ACCOUNT_LOCK_SQL: &str =
981 "SELECT 1 FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE";
982
983pub(crate) const LIVE_KEY_COUNT_SQL: &str = "SELECT count(*) FROM tollgate_credential_keys
986 WHERE account_id = $1
987 AND revoked_at_us IS NULL
988 AND (not_after_floor_us IS NULL OR not_after_submicro_ns IS NULL
989 OR (not_after_floor_us, not_after_submicro_ns) > ($2, $3))";
990
991const RECLAIM_DUE_LEASES_SQL: &str = "SELECT lease_id, account_id, granted, used
992 FROM tollgate_leases
993 WHERE state = 0 AND (expires_at_floor_us, expires_at_submicro_ns) <= ($1, $2)
994 ORDER BY expires_at_floor_us, expires_at_submicro_ns
995 LIMIT $3 FOR UPDATE SKIP LOCKED";
996
997const DUE_PERIODS_SQL: &str = "WITH due AS (
1026 SELECT account_id, allowance_balance AS prior, budget_allowance AS allowance,
1027 period_start_us AS crossed_from
1028 FROM tollgate_accounts
1029 WHERE budget_allowance IS NOT NULL
1030 AND budget_period = $1 AND period_start_us < $2
1031 ORDER BY period_start_us LIMIT $3 FOR UPDATE SKIP LOCKED
1032 )
1033 UPDATE tollgate_accounts AS account SET
1034 deposited = account.deposited + due.allowance,
1035 expired = account.expired + due.prior,
1036 balance = account.balance - due.prior + due.allowance,
1037 allowance_balance = due.allowance,
1038 period_start_us = $2
1039 FROM due
1040 WHERE account.account_id = due.account_id
1041 RETURNING due.account_id, due.allowance, due.prior, due.crossed_from";
1042
1043fn id_bytes(id: u128) -> Vec<u8> {
1044 id.to_be_bytes().to_vec()
1045}
1046
1047fn id_from(bytes: &[u8]) -> u128 {
1048 let mut buf = [0u8; 16];
1049 buf.copy_from_slice(bytes);
1050 u128::from_be_bytes(buf)
1051}
1052
1053fn to_i64(units: CostUnits, what: &str) -> Result<i64, StoreError> {
1054 i64::try_from(units.get()).map_err(|_| StoreError(format!("{what} exceeds i64 range")))
1055}
1056
1057fn stored_fence(value: i64) -> Result<FencingToken, StoreError> {
1060 u64::try_from(value)
1061 .ok()
1062 .filter(|value| *value > 0)
1063 .map(FencingToken)
1064 .ok_or_else(|| StoreError(format!("stored fencing token is not positive: {value}")))
1065}
1066
1067fn to_units(value: i64, what: &str) -> Result<CostUnits, StoreError> {
1071 u64::try_from(value)
1072 .map(CostUnits)
1073 .map_err(|_| StoreError(format!("{what} is negative in storage: {value}")))
1074}
1075
1076fn ts_micros(ts: Timestamp) -> i64 {
1077 ts.as_microsecond()
1078}
1079
1080fn micros_ts(value: i64, what: &str) -> Result<Timestamp, StoreError> {
1083 tollgate_store::clock::timestamp_from_micros(value).map_err(|e| {
1084 StoreError(format!(
1085 "{what} is not a representable instant: {value} ({e})"
1086 ))
1087 })
1088}
1089
1090fn storage(e: sqlx::Error) -> StoreError {
1091 StoreError(format!("postgres: {e}"))
1092}
1093
1094fn alloc_storage(e: sqlx::Error) -> AllocateError {
1095 AllocateError::Storage(storage(e))
1096}
1097
1098async fn finish_transaction<T, E>(
1111 tx: Transaction<'_, Postgres>,
1112 result: Result<T, E>,
1113) -> Result<T, E>
1114where
1115 E: From<StoreError> + std::fmt::Display,
1116{
1117 match result {
1118 Ok(value) => {
1119 tx.commit().await.map_err(|e| E::from(storage(e)))?;
1120 Ok(value)
1121 }
1122 Err(error) => {
1123 if let Err(rollback_error) = tx.rollback().await {
1124 return Err(E::from(StoreError(format!(
1125 "operation failed ({error}); transaction rollback failed ({})",
1126 storage(rollback_error)
1127 ))));
1128 }
1129 Err(error)
1130 }
1131 }
1132}
1133
1134pub struct PostgresStore {
1171 pool: PgPool,
1172 policy: GrantPolicy,
1173 push: broadcast::Sender<SnapshotPush>,
1174}
1175
1176#[derive(Debug, Clone, Copy)]
1181pub struct PoolConfig {
1182 pub max_connections: u32,
1185 pub acquire_timeout: std::time::Duration,
1188}
1189
1190impl Default for PoolConfig {
1191 fn default() -> Self {
1192 PoolConfig {
1193 max_connections: 16,
1194 acquire_timeout: std::time::Duration::from_secs(5),
1195 }
1196 }
1197}
1198
1199impl PoolConfig {
1200 pub fn validate(&self) -> Result<(), StoreError> {
1206 if self.max_connections == 0 {
1207 return Err(StoreError("max_connections must be positive".into()));
1208 }
1209 if self.acquire_timeout.is_zero() {
1210 return Err(StoreError("acquire_timeout must be positive".into()));
1211 }
1212 Ok(())
1213 }
1214}
1215
1216impl PostgresStore {
1217 pub async fn connect(url: &str, policy: GrantPolicy) -> Result<Arc<Self>, StoreError> {
1219 Self::connect_with(url, policy, PoolConfig::default()).await
1220 }
1221
1222 pub async fn connect_with(
1231 url: &str,
1232 policy: GrantPolicy,
1233 pool_config: PoolConfig,
1234 ) -> Result<Arc<Self>, StoreError> {
1235 policy
1236 .validate()
1237 .map_err(|e| StoreError(format!("invalid grant policy: {e}")))?;
1238 pool_config.validate()?;
1239 let pool = PgPoolOptions::new()
1240 .max_connections(pool_config.max_connections)
1241 .acquire_timeout(pool_config.acquire_timeout)
1242 .connect(url)
1243 .await
1244 .map_err(storage)?;
1245 sqlx::migrate!("./migrations")
1246 .run(&pool)
1247 .await
1248 .map_err(|e| StoreError(format!("migrate: {e}")))?;
1249 let (push, _) = broadcast::channel(PUSH_CHANNEL_CAPACITY);
1250 Ok(Arc::new(PostgresStore { pool, policy, push }))
1251 }
1252
1253 fn push_to_subscribers(&self, push: SnapshotPush) {
1259 let principal = push.principal;
1260 let subscribers = self.push.send(push).unwrap_or(0);
1261 tracing::debug!(
1262 %principal,
1263 subscribers,
1264 "snapshot pushed to subscribers"
1265 );
1266 }
1267
1268 fn announce_removal(&self, principal: Principal, receipt: &AdminReceipt<()>) {
1270 if receipt.before != receipt.after
1271 && let AdminState::Snapshot { generation, .. } = receipt.after
1272 {
1273 self.push_to_subscribers(SnapshotPush {
1274 principal,
1275 resolution: SnapshotResolution::Revoked { generation },
1276 });
1277 }
1278 }
1279
1280 pub async fn balance(&self, account: AccountId) -> Result<CostUnits, StoreError> {
1290 let row = sqlx::query("SELECT balance FROM tollgate_accounts WHERE account_id = $1")
1291 .bind(id_bytes(account.0))
1292 .fetch_optional(&self.pool)
1293 .await
1294 .map_err(storage)?;
1295 row.map(|r| to_units(r.get::<i64, _>(0), "account balance"))
1296 .transpose()
1297 .map(|units| units.unwrap_or(CostUnits::ZERO))
1298 }
1299
1300 pub async fn usage_recorded(&self, account: AccountId) -> Result<CostUnits, StoreError> {
1307 let row = sqlx::query("SELECT usage_recorded FROM tollgate_accounts WHERE account_id = $1")
1308 .bind(id_bytes(account.0))
1309 .fetch_optional(&self.pool)
1310 .await
1311 .map_err(storage)?;
1312 row.map(|r| to_units(r.get::<i64, _>(0), "account usage_recorded"))
1313 .transpose()
1314 .map(|units| units.unwrap_or(CostUnits::ZERO))
1315 }
1316
1317 pub async fn conservation(
1332 &self,
1333 account: AccountId,
1334 ) -> Result<Option<Conservation>, StoreError> {
1335 let mut tx = self.pool.begin().await.map_err(storage)?;
1349 let result = async {
1350 sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY")
1351 .execute(&mut *tx)
1352 .await
1353 .map_err(storage)?;
1354 let account_row = sqlx::query(
1355 "SELECT deposited, balance, usage_recorded, settlement_loss, overage_recorded,
1356 expired
1357 FROM tollgate_accounts WHERE account_id = $1",
1358 )
1359 .bind(id_bytes(account.0))
1360 .fetch_optional(&mut *tx)
1361 .await
1362 .map_err(storage)?;
1363 let lease_row = sqlx::query(ACTIVE_LEASE_SUM_SQL)
1364 .bind(id_bytes(account.0))
1365 .fetch_one(&mut *tx)
1366 .await
1367 .map_err(storage)?;
1368 Ok::<_, StoreError>((account_row, lease_row))
1369 }
1370 .await;
1371 let (account_row, lease_row) = finish_transaction(tx, result).await?;
1372
1373 let Some(row) = account_row else {
1374 return Ok(None);
1375 };
1376 let active_grants = to_units(lease_row.get::<i64, _>(0), "active lease grants")?;
1377 let active_used = to_units(lease_row.get::<i64, _>(1), "active lease usage")?;
1378 let recorded = to_units(row.get::<i64, _>(2), "usage_recorded")?;
1379 Ok(Some(Conservation {
1380 deposited: to_units(row.get::<i64, _>(0), "deposited")?,
1381 overage_recorded: to_units(row.get::<i64, _>(4), "overage_recorded")?,
1382 balance: to_units(row.get::<i64, _>(1), "balance")?,
1383 active_lease_grants: active_grants,
1384 settled_usage: recorded.checked_sub(active_used).ok_or_else(|| {
1392 StoreError(format!(
1393 "active lease usage {} exceeds recorded usage {} for account {account}",
1394 active_used.get(),
1395 recorded.get()
1396 ))
1397 })?,
1398 settlement_loss: to_units(row.get::<i64, _>(3), "settlement_loss")?,
1399 expired: to_units(row.get::<i64, _>(5), "expired")?,
1400 }))
1401 }
1402}
1403
1404struct LockedLeaseRow {
1408 account_id: Vec<u8>,
1409 fencing_token: i64,
1410 granted: i64,
1411 used: i64,
1412 credited: i64,
1413 expires_at: StoredInstant,
1414 state: i16,
1415 from_allowance: i64,
1420 period_start_us: i64,
1421}
1422
1423struct ReleasedCredit {
1426 account: AccountId,
1427 restored: CostUnits,
1428 preserves_funding: bool,
1429}
1430
1431async fn lock_lease(
1433 tx: &mut Transaction<'_, Postgres>,
1434 lease_id: LeaseId,
1435) -> Result<Option<LockedLeaseRow>, sqlx::Error> {
1436 let row = sqlx::query(
1437 "SELECT account_id, fencing_token, granted, used, credited, expires_at_floor_us, state,
1438 from_allowance, period_start_us, expires_at_submicro_ns
1439 FROM tollgate_leases WHERE lease_id = $1 FOR UPDATE",
1440 )
1441 .bind(id_bytes(lease_id.0))
1442 .fetch_optional(&mut **tx)
1443 .await?;
1444 Ok(row.map(|row| LockedLeaseRow {
1445 account_id: row.get(0),
1446 fencing_token: row.get(1),
1447 granted: row.get(2),
1448 used: row.get(3),
1449 credited: row.get(4),
1450 expires_at: StoredInstant {
1451 micros: row.get(5),
1452 submicro_nanos: row.get(9),
1453 },
1454 state: row.get(6),
1455 from_allowance: row.get(7),
1456 period_start_us: row.get(8),
1457 }))
1458}
1459
1460struct Exchange {
1463 floor: CostUnits,
1467 needed: CostUnits,
1470 preserves_funding: bool,
1473}
1474
1475impl Exchange {
1476 const ACQUIRE: Exchange = Exchange {
1477 floor: CostUnits::ZERO,
1478 needed: CostUnits::ZERO,
1479 preserves_funding: true,
1480 };
1481}
1482
1483impl PostgresStore {
1484 async fn acquire_in_tx(
1486 &self,
1487 tx: &mut Transaction<'_, Postgres>,
1488 account: AccountId,
1489 requested: CostUnits,
1490 expires_at: Timestamp,
1491 exchange: Exchange,
1492 ) -> Result<Allocation, AllocateError> {
1493 let Exchange {
1494 floor,
1495 needed,
1496 preserves_funding: settlement_preserves_funding,
1497 } = exchange;
1498 let row = sqlx::query(
1499 "SELECT balance, status, next_fence, allowance_balance, period_start_us
1500 FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE",
1501 )
1502 .bind(id_bytes(account.0))
1503 .fetch_optional(&mut **tx)
1504 .await
1505 .map_err(alloc_storage)?
1506 .ok_or(AllocateError::UnknownAccount)?;
1507
1508 if decode_status(row.get::<String, _>(1)).map_err(AllocateError::Storage)?
1512 != AccountStatus::Active
1513 {
1514 return Err(AllocateError::AccountInactive);
1515 }
1516 let balance =
1517 to_units(row.get::<i64, _>(0), "account balance").map_err(AllocateError::Storage)?;
1518 let granted = match self
1524 .policy
1525 .consolidation_grant(requested, balance, floor, needed)
1526 {
1527 Some(granted) => granted,
1528 None => {
1529 if settlement_preserves_funding && !requested.is_zero() {
1532 let budget = sqlx::query(BUDGET_VIEW_SQL)
1533 .bind(id_bytes(account.0))
1534 .fetch_one(&mut **tx)
1535 .await
1536 .map_err(alloc_storage)?;
1537 let evidence = budget_view(&budget)
1538 .map_err(AllocateError::Storage)?
1539 .shortfall();
1540 return Err(match evidence.exhaustion() {
1541 Some(exhausted) => AllocateError::BalanceExhausted(exhausted),
1542 None => AllocateError::BalanceInsufficient(evidence),
1543 });
1544 }
1545 return Err(AllocateError::InsufficientBalance);
1546 }
1547 };
1548 let granted_i = to_i64(granted, "grant").map_err(AllocateError::Storage)?;
1553 let allowance_balance = row.get::<i64, _>(3);
1554 let from_allowance = granted_i.min(allowance_balance);
1555 let period_start_us = row.get::<i64, _>(4);
1556 let fence = row.get::<i64, _>(2);
1557 let fence_token = stored_fence(fence).map_err(AllocateError::Storage)?;
1558
1559 let lease_id = LeaseId(uuid::Uuid::new_v4().as_u128());
1560
1561 let budget = sqlx::query(GRANT_DEBIT_SQL)
1565 .bind(id_bytes(account.0))
1566 .bind(granted_i)
1567 .bind(from_allowance)
1568 .fetch_one(&mut **tx)
1569 .await
1570 .map_err(alloc_storage)?;
1571 let funding = budget_view(&budget)
1572 .map_err(AllocateError::Storage)?
1573 .shortfall();
1574 sqlx::query(
1575 "INSERT INTO tollgate_leases
1576 (lease_id, account_id, fencing_token, granted, used, credited, expires_at_floor_us,
1577 state, from_allowance, period_start_us, expires_at_submicro_ns, expiry_is_upper_bound)
1578 VALUES ($1, $2, $3, $4, 0, 0, $5, 0, $6, $7, $8, FALSE)",
1579 )
1580 .bind(id_bytes(lease_id.0))
1581 .bind(id_bytes(account.0))
1582 .bind(fence)
1583 .bind(granted_i)
1584 .bind(StoredInstant::from(expires_at).micros)
1585 .bind(from_allowance)
1586 .bind(period_start_us)
1587 .bind(StoredInstant::from(expires_at).submicro_nanos)
1588 .execute(&mut **tx)
1589 .await
1590 .map_err(alloc_storage)?;
1591
1592 Ok(Allocation {
1593 grant: LeaseGrant {
1594 lease_id,
1595 account_id: account,
1596 fencing_token: fence_token,
1597 units: granted,
1598 expires_at,
1599 },
1600 funding: Some(funding),
1601 })
1602 }
1603
1604 async fn release_in_tx(
1607 &self,
1608 tx: &mut Transaction<'_, Postgres>,
1609 lease_id: LeaseId,
1610 fencing_token: FencingToken,
1611 unspent: CostUnits,
1612 now: Timestamp,
1613 ) -> Result<ReleasedCredit, AllocateError> {
1614 let LockedLeaseRow {
1615 account_id,
1616 fencing_token: fence,
1617 granted,
1618 used,
1619 credited: _credited,
1620 expires_at,
1621 state,
1622 from_allowance,
1623 period_start_us,
1624 } = lock_lease(tx, lease_id)
1625 .await
1626 .map_err(alloc_storage)?
1627 .ok_or(AllocateError::UnknownLease)?;
1628
1629 if stored_fence(fence).map_err(AllocateError::Storage)? != fencing_token {
1630 return Err(AllocateError::Fenced);
1631 }
1632 let expires_at = expires_at.timestamp().map_err(AllocateError::Storage)?;
1635 if state != STATE_ACTIVE
1636 || self
1637 .policy
1638 .reclaim_cutoff(now)
1639 .is_some_and(|cutoff| expires_at <= cutoff)
1640 {
1641 return Err(AllocateError::LeaseNotActive);
1642 }
1643 let account = AccountId(id_from(&account_id));
1647 let unspent_i = to_i64(unspent, "unspent").map_err(AllocateError::Storage)?;
1648 let loss = granted
1649 .checked_sub(
1650 used.checked_add(unspent_i)
1651 .ok_or(AllocateError::InvalidRelease)?,
1652 )
1653 .filter(|l| *l >= 0)
1654 .ok_or(AllocateError::InvalidRelease)?;
1655
1656 sqlx::query("UPDATE tollgate_leases SET state = $2, credited = $3 WHERE lease_id = $1")
1657 .bind(id_bytes(lease_id.0))
1658 .bind(STATE_RELEASED)
1659 .bind(unspent_i)
1660 .execute(&mut **tx)
1661 .await
1662 .map_err(alloc_storage)?;
1663 let from_topup = granted
1669 .checked_sub(from_allowance)
1670 .filter(|t| *t >= 0)
1671 .ok_or_else(|| {
1672 AllocateError::Storage(StoreError(format!(
1673 "lease allowance funding {from_allowance} exceeds its grant {granted}"
1674 )))
1675 })?;
1676 let to_topup = from_topup.min(unspent_i);
1677 let to_allowance = unspent_i - to_topup;
1678
1679 let restored: i64 = sqlx::query_scalar(
1686 "UPDATE tollgate_accounts SET
1687 balance = balance + $2
1688 + CASE WHEN period_start_us > $5 THEN 0 ELSE $3 END,
1689 allowance_balance = allowance_balance
1690 + CASE WHEN period_start_us > $5 THEN 0 ELSE $3 END,
1691 expired = expired + CASE WHEN period_start_us > $5 THEN $3 ELSE 0 END,
1692 settlement_loss = settlement_loss + $4
1693 WHERE account_id = $1
1694 RETURNING $2 + CASE WHEN period_start_us > $5 THEN 0 ELSE $3 END",
1695 )
1696 .bind(account_id)
1697 .bind(to_topup)
1698 .bind(to_allowance)
1699 .bind(loss)
1700 .bind(period_start_us)
1701 .fetch_one(&mut **tx)
1702 .await
1703 .map_err(alloc_storage)?;
1704 Ok(ReleasedCredit {
1705 account,
1706 preserves_funding: loss == 0 && restored == unspent_i,
1707 restored: to_units(restored, "restored release credit")
1708 .map_err(AllocateError::Storage)?,
1709 })
1710 }
1711
1712 fn grant_expiry(
1716 &self,
1717 ttl: SignedDuration,
1718 now: Timestamp,
1719 ) -> Result<Timestamp, AllocateError> {
1720 if ttl <= SignedDuration::ZERO {
1721 return Err(AllocateError::InvalidTtl);
1722 }
1723 now.checked_add(ttl.min(self.policy.max_ttl))
1724 .map_err(|e| AllocateError::Storage(StoreError(format!("ttl overflow: {e}"))))
1725 }
1726}
1727
1728#[async_trait]
1729impl LeaseAllocator for PostgresStore {
1730 async fn acquire(
1731 &self,
1732 account: AccountId,
1733 requested: CostUnits,
1734 ttl: SignedDuration,
1735 now: Timestamp,
1736 ) -> Result<Allocation, AllocateError> {
1737 let expires_at = self.grant_expiry(ttl, now)?;
1738 let mut tx = self.pool.begin().await.map_err(alloc_storage)?;
1739 let result = self
1740 .acquire_in_tx(&mut tx, account, requested, expires_at, Exchange::ACQUIRE)
1741 .await;
1742 finish_transaction(tx, result).await
1743 }
1744
1745 async fn release(
1746 &self,
1747 lease_id: LeaseId,
1748 fencing_token: FencingToken,
1749 unspent: CostUnits,
1750 now: Timestamp,
1751 ) -> Result<(), AllocateError> {
1752 let mut tx = self.pool.begin().await.map_err(alloc_storage)?;
1753 let result = self
1754 .release_in_tx(&mut tx, lease_id, fencing_token, unspent, now)
1755 .await
1756 .map(|_| ());
1757 finish_transaction(tx, result).await
1758 }
1759
1760 async fn consolidate(
1761 &self,
1762 lease_id: LeaseId,
1763 fencing_token: FencingToken,
1764 unspent: CostUnits,
1765 requested: CostUnits,
1766 needed: CostUnits,
1767 ttl: SignedDuration,
1768 now: Timestamp,
1769 ) -> Result<Allocation, AllocateError> {
1770 let expires_at = self.grant_expiry(ttl, now)?;
1771 let mut tx = self.pool.begin().await.map_err(alloc_storage)?;
1772 let result = async {
1781 let released = self
1782 .release_in_tx(&mut tx, lease_id, fencing_token, unspent, now)
1783 .await?;
1784 self.acquire_in_tx(
1785 &mut tx,
1786 released.account,
1787 requested,
1788 expires_at,
1789 Exchange {
1790 floor: released.restored,
1791 needed,
1792 preserves_funding: released.preserves_funding,
1793 },
1794 )
1795 .await
1796 }
1797 .await;
1798 finish_transaction(tx, result).await
1799 }
1800
1801 async fn reclaim_expired_batch(
1802 &self,
1803 now: Timestamp,
1804 limit: NonZeroUsize,
1805 ) -> Result<ReclaimBatch, StoreError> {
1806 let limit_i = i64::try_from(limit.get())
1807 .map_err(|_| StoreError(format!("reclaim batch limit exceeds i64 range: {limit}")))?;
1808 let Some(cutoff) = self.policy.reclaim_cutoff(now) else {
1809 return ReclaimBatch::try_new(Vec::new(), limit);
1810 };
1811 let cutoff = StoredInstant::from(cutoff);
1812 let mut tx = self.pool.begin().await.map_err(storage)?;
1813 let result = async {
1814 let rows = sqlx::query(RECLAIM_DUE_LEASES_SQL)
1820 .bind(cutoff.micros)
1821 .bind(cutoff.submicro_nanos)
1822 .bind(limit_i)
1823 .fetch_all(&mut *tx)
1824 .await
1825 .map_err(storage)?;
1826
1827 let mut reclaimed = Vec::with_capacity(rows.len());
1832 let mut lease_ids = Vec::with_capacity(rows.len());
1833 let mut forfeits: std::collections::BTreeMap<Vec<u8>, i64> =
1834 std::collections::BTreeMap::new();
1835 for row in rows {
1836 let lease_bytes: Vec<u8> = row.get(0);
1837 let account_bytes: Vec<u8> = row.get(1);
1838 let granted = row.get::<i64, _>(2);
1839 let used = row.get::<i64, _>(3);
1840 let forfeited = granted.checked_sub(used).ok_or_else(|| {
1844 StoreError(format!(
1845 "reclaim remainder overflow: granted {granted}, used {used}"
1846 ))
1847 })?;
1848 let forfeited_units = to_units(forfeited, "reclaim remainder")?;
1849 let total = forfeits.entry(account_bytes.clone()).or_default();
1850 *total = total
1851 .checked_add(forfeited)
1852 .ok_or_else(|| StoreError("reclaim loss sum overflow".into()))?;
1853 lease_ids.push(lease_bytes.clone());
1854 reclaimed.push(ReclaimedLease {
1855 lease_id: LeaseId(id_from(&lease_bytes)),
1856 account_id: AccountId(id_from(&account_bytes)),
1857 forfeited: forfeited_units,
1858 });
1859 }
1860
1861 let batch = ReclaimBatch::try_new(reclaimed, limit)?;
1862 if batch.is_empty() {
1863 return Ok(batch);
1864 }
1865
1866 let (account_ids, account_forfeits): (Vec<_>, Vec<_>) = forfeits.into_iter().unzip();
1867 let expected_lease_rows = u64::try_from(lease_ids.len())
1868 .map_err(|_| StoreError("reclaim lease row count exceeds u64 range".into()))?;
1869 let expected_account_rows = u64::try_from(account_ids.len())
1870 .map_err(|_| StoreError("reclaim account row count exceeds u64 range".into()))?;
1871
1872 let locked_accounts = sqlx::query(
1878 "SELECT account_id FROM tollgate_accounts
1879 WHERE account_id = ANY($1) ORDER BY account_id FOR UPDATE",
1880 )
1881 .bind(&account_ids)
1882 .fetch_all(&mut *tx)
1883 .await
1884 .map_err(storage)?;
1885 if locked_accounts.len() != account_ids.len() {
1886 return Err(StoreError(format!(
1887 "reclaim locked {} of {} referenced account rows",
1888 locked_accounts.len(),
1889 account_ids.len()
1890 )));
1891 }
1892
1893 let updated_leases = sqlx::query(
1894 "UPDATE tollgate_leases
1895 SET state = $2, credited = 0
1896 WHERE lease_id = ANY($1) AND state = $3",
1897 )
1898 .bind(&lease_ids)
1899 .bind(STATE_EXPIRED)
1900 .bind(STATE_ACTIVE)
1901 .execute(&mut *tx)
1902 .await
1903 .map_err(storage)?;
1904 if updated_leases.rows_affected() != expected_lease_rows {
1905 return Err(StoreError(format!(
1906 "reclaim updated {} of {} locked lease rows",
1907 updated_leases.rows_affected(),
1908 lease_ids.len()
1909 )));
1910 }
1911
1912 let updated_accounts = sqlx::query(
1913 "UPDATE tollgate_accounts AS account
1914 SET settlement_loss = account.settlement_loss + delta.forfeited
1915 FROM UNNEST($1::bytea[], $2::bigint[]) AS delta(account_id, forfeited)
1916 WHERE account.account_id = delta.account_id",
1917 )
1918 .bind(&account_ids)
1919 .bind(&account_forfeits)
1920 .execute(&mut *tx)
1921 .await
1922 .map_err(storage)?;
1923 if updated_accounts.rows_affected() != expected_account_rows {
1924 return Err(StoreError(format!(
1925 "reclaim updated {} of {} locked account rows",
1926 updated_accounts.rows_affected(),
1927 account_ids.len()
1928 )));
1929 }
1930
1931 Ok(batch)
1932 }
1933 .await;
1934 finish_transaction(tx, result).await
1935 }
1936}
1937
1938#[async_trait]
1939impl UsageSink for PostgresStore {
1940 async fn ingest(
1941 &self,
1942 events: &[UsageEvent],
1943 _now: Timestamp,
1944 ) -> Result<IngestReport, IngestError> {
1945 let mut report = IngestReport {
1953 unattributed: Some(0),
1954 ..IngestReport::default()
1955 };
1956 if events.is_empty() {
1957 return Ok(report);
1958 }
1959
1960 struct PreparedEvent<'a> {
1966 event: &'a UsageEvent,
1967 request_id: Vec<u8>,
1968 account_id: Vec<u8>,
1969 key_id: Option<Vec<u8>>,
1970 lease_id: Option<Vec<u8>>,
1974 occurred_at_us: i64,
1975 }
1976 let prepared: Vec<PreparedEvent<'_>> = events
1977 .iter()
1978 .map(|event| {
1979 Ok(PreparedEvent {
1980 event,
1981 request_id: id_bytes(event.request_id.0),
1982 account_id: id_bytes(event.account_id.0),
1983 key_id: event.key_id.map(|id| id_bytes(id.0)),
1984 lease_id: event.source.lease_id().map(|id| id_bytes(id.0)),
1985 occurred_at_us: ts_micros(event.occurred_at),
1986 })
1987 })
1988 .collect::<Result<_, StoreError>>()?;
1989
1990 let mut tx = self.pool.begin().await.map_err(storage)?;
1991 let result: Result<_, IngestError> = async {
1992 let lease_ids: Vec<Vec<u8>> = prepared
2001 .iter()
2002 .filter_map(|event| event.lease_id.clone())
2003 .collect::<std::collections::BTreeSet<_>>()
2004 .into_iter()
2005 .collect();
2006 struct LeaseRow {
2007 account_id: Vec<u8>,
2008 fence: i64,
2009 granted: i64,
2010 used: i64,
2011 used_delta: Option<NonZeroI64>,
2012 credited: i64,
2013 settled: bool,
2014 }
2015 let rows = sqlx::query(
2016 "SELECT lease_id, account_id, fencing_token, granted, used, credited, state
2017 FROM tollgate_leases
2018 WHERE lease_id = ANY($1)
2019 ORDER BY account_id, lease_id FOR UPDATE",
2020 )
2021 .bind(&lease_ids)
2022 .fetch_all(&mut *tx)
2023 .await
2024 .map_err(storage)?;
2025 let mut leases: std::collections::BTreeMap<Vec<u8>, LeaseRow> =
2026 std::collections::BTreeMap::new();
2027 for row in rows {
2028 let lease_id: Vec<u8> = row.get(0);
2029 let fence: i64 = row.get(2);
2030 let granted: i64 = row.get(3);
2031 let used: i64 = row.get(4);
2032 let credited: i64 = row.get(5);
2033 leases.insert(
2034 lease_id,
2035 LeaseRow {
2036 account_id: row.get(1),
2037 fence,
2038 granted,
2039 used,
2040 used_delta: None,
2041 credited,
2042 settled: row.get::<i16, _>(6) != STATE_ACTIVE,
2043 },
2044 );
2045 }
2046
2047 let request_ids: Vec<Vec<u8>> = prepared
2049 .iter()
2050 .map(|event| event.request_id.clone())
2051 .collect();
2052 let mut seen: std::collections::HashSet<Vec<u8>> = sqlx::query(
2053 "SELECT request_id FROM tollgate_usage_events WHERE request_id = ANY($1)",
2054 )
2055 .bind(&request_ids)
2056 .fetch_all(&mut *tx)
2057 .await
2058 .map_err(storage)?
2059 .into_iter()
2060 .map(|row| row.get::<Vec<u8>, _>(0))
2061 .collect();
2062
2063 let overage_account_ids: Vec<Vec<u8>> = prepared
2079 .iter()
2080 .filter(|event| event.lease_id.is_none())
2081 .map(|event| event.account_id.clone())
2082 .collect::<std::collections::BTreeSet<_>>()
2083 .into_iter()
2084 .collect();
2085 let known_overage_accounts: std::collections::HashSet<Vec<u8>> =
2086 if overage_account_ids.is_empty() {
2087 std::collections::HashSet::new()
2088 } else {
2089 sqlx::query(
2090 "SELECT account_id FROM tollgate_accounts WHERE account_id = ANY($1)",
2091 )
2092 .bind(&overage_account_ids)
2093 .fetch_all(&mut *tx)
2094 .await
2095 .map_err(storage)?
2096 .into_iter()
2097 .map(|row| row.get::<Vec<u8>, _>(0))
2098 .collect()
2099 };
2100
2101 struct Accepted {
2105 event_index: usize,
2106 settled: bool,
2107 fence: Option<i64>,
2112 overage: bool,
2115 units: i64,
2118 }
2119 let mut accepted: Vec<Accepted> = Vec::with_capacity(prepared.len());
2120 for (event_index, event) in prepared.iter().enumerate() {
2121 if seen.contains(event.request_id.as_slice()) {
2122 report.duplicate += 1;
2123 continue;
2124 }
2125 let Some(lease_key) = event.lease_id.as_deref() else {
2126 if !known_overage_accounts.contains(event.account_id.as_slice()) {
2135 report.rejected += 1;
2136 continue;
2137 }
2138 let Ok(units) = i64::try_from(event.event.units.get()) else {
2139 report.rejected += 1;
2142 continue;
2143 };
2144 seen.insert(event.request_id.clone());
2145 accepted.push(Accepted {
2146 event_index,
2147 settled: false,
2151 fence: None,
2152 overage: true,
2153 units,
2154 });
2155 report.accepted += 1;
2156 continue;
2157 };
2158 let Some(lease) = leases.get_mut(lease_key) else {
2159 report.rejected += 1;
2160 continue;
2161 };
2162 if Some(stored_fence(lease.fence)?) != event.event.source.fencing_token()
2163 || lease.account_id.as_slice() != event.account_id.as_slice()
2164 {
2165 report.rejected += 1;
2166 continue;
2167 }
2168 to_units(lease.granted, "lease granted")?;
2169 to_units(lease.used, "lease used")?;
2170 to_units(lease.credited, "lease credited")?;
2171 let Ok(units) = i64::try_from(event.event.units.get()) else {
2172 report.rejected += 1;
2173 continue;
2174 };
2175 let committed = lease
2176 .used
2177 .checked_add(lease.used_delta.map_or(0, NonZeroI64::get))
2178 .and_then(|used| used.checked_add(lease.credited))
2179 .ok_or_else(|| {
2180 StoreError(format!(
2181 "lease accounting overflow for {:#034x}",
2182 id_from(lease_key)
2183 ))
2184 })?;
2185 let remaining = lease.granted.checked_sub(committed).ok_or_else(|| {
2186 StoreError(format!(
2187 "lease accounting exceeds grant for {:#034x}: granted {}, committed {committed}",
2188 id_from(lease_key), lease.granted
2189 ))
2190 })?;
2191 if units > remaining {
2192 report.rejected += 1;
2193 continue;
2194 }
2195 let used_delta = lease
2196 .used_delta
2197 .map_or(0, NonZeroI64::get)
2198 .checked_add(units)
2199 .ok_or_else(|| {
2200 StoreError(format!(
2201 "lease usage delta overflow for {:#034x}",
2202 id_from(lease_key)
2203 ))
2204 })?;
2205 lease.used_delta = NonZeroI64::new(used_delta);
2206 seen.insert(event.request_id.clone());
2207 accepted.push(Accepted {
2208 event_index,
2209 settled: lease.settled,
2210 fence: Some(lease.fence),
2211 overage: false,
2212 units,
2213 });
2214 report.accepted += 1;
2215 }
2216
2217 if accepted.is_empty() {
2218 return Ok(report);
2219 }
2220
2221 #[derive(Default)]
2222 struct AccountDelta {
2223 usage: i64,
2224 loss: i64,
2225 overage: i64,
2229 }
2230
2231 let (mut rid, mut acct, mut lease, mut fence, mut units, mut at, mut revision, mut keys) = (
2233 Vec::with_capacity(accepted.len()),
2234 Vec::with_capacity(accepted.len()),
2235 Vec::with_capacity(accepted.len()),
2236 Vec::with_capacity(accepted.len()),
2237 Vec::with_capacity(accepted.len()),
2238 Vec::with_capacity(accepted.len()),
2239 Vec::with_capacity(accepted.len()),
2240 Vec::with_capacity(accepted.len()),
2241 );
2242 let mut account_deltas: std::collections::BTreeMap<Vec<u8>, AccountDelta> =
2243 std::collections::BTreeMap::new();
2244 for accepted_event in &accepted {
2245 let event = &prepared[accepted_event.event_index];
2246 rid.push(event.request_id.clone());
2247 acct.push(event.account_id.clone());
2248 keys.push(event.key_id.clone());
2249 lease.push(event.lease_id.clone());
2250 fence.push(accepted_event.fence);
2251 units.push(accepted_event.units);
2252 at.push(event.occurred_at_us);
2253 revision.push(event.event.policy_revision.as_bytes().to_vec());
2258 if accepted_event.overage {
2259 debug_assert!(
2260 event.lease_id.is_none() && accepted_event.fence.is_none(),
2261 "an overage row must carry neither half of a capability"
2262 );
2263 }
2264
2265 let entry = account_deltas.entry(event.account_id.clone()).or_default();
2266 entry.usage = entry.usage.checked_add(accepted_event.units).ok_or_else(|| {
2267 IngestError::Refused(StoreError(format!(
2268 "usage delta overflow for account {:#034x}",
2269 event.event.account_id.0
2270 )))
2271 })?;
2272 if accepted_event.settled {
2273 entry.loss = entry.loss.checked_add(accepted_event.units).ok_or_else(|| {
2274 IngestError::Refused(StoreError(format!(
2275 "settlement loss delta overflow for account {:#034x}",
2276 event.event.account_id.0
2277 )))
2278 })?;
2279 }
2280 if accepted_event.overage {
2281 entry.overage =
2282 entry.overage.checked_add(accepted_event.units).ok_or_else(|| {
2283 IngestError::Refused(StoreError(format!(
2284 "overage delta overflow for account {:#034x}",
2285 event.event.account_id.0
2286 )))
2287 })?;
2288 }
2289 }
2290 let inserted = sqlx::query(
2291 "INSERT INTO tollgate_usage_events
2292 (request_id, account_id, lease_id, fencing_token, units, occurred_at_us, policy_revision, key_id)
2293 SELECT * FROM UNNEST($1::bytea[], $2::bytea[], $3::bytea[], $4::bigint[], $5::bigint[], $6::bigint[], $7::bytea[], $8::bytea[])",
2294 )
2295 .bind(&rid)
2296 .bind(&acct)
2297 .bind(&lease)
2298 .bind(&fence)
2299 .bind(&units)
2300 .bind(&at)
2301 .bind(&revision)
2302 .bind(&keys)
2303 .execute(&mut *tx)
2304 .await
2305 .map_err(storage)?;
2306 let expected_event_rows = u64::try_from(accepted.len())
2307 .map_err(|_| StoreError("accepted event count exceeds u64 range".into()))?;
2308 if inserted.rows_affected() != expected_event_rows {
2309 return Err(IngestError::Unavailable(StoreError(format!(
2310 "ingest inserted {} of {} accepted usage rows",
2311 inserted.rows_affected(),
2312 accepted.len()
2313 ))));
2314 }
2315
2316 let (lease_update_ids, lease_used_deltas): (Vec<Vec<u8>>, Vec<i64>) = leases
2320 .iter()
2321 .filter_map(|(lease_id, row)| {
2322 row.used_delta
2323 .map(|used_delta| (lease_id.clone(), used_delta.get()))
2324 })
2325 .unzip();
2326 if !lease_update_ids.is_empty() {
2327 let updated_leases = sqlx::query(
2328 "UPDATE tollgate_leases AS lease
2329 SET used = lease.used + delta.used
2330 FROM UNNEST($1::bytea[], $2::bigint[]) AS delta(lease_id, used)
2331 WHERE lease.lease_id = delta.lease_id",
2332 )
2333 .bind(&lease_update_ids)
2334 .bind(&lease_used_deltas)
2335 .execute(&mut *tx)
2336 .await
2337 .map_err(storage)?;
2338 let expected_lease_rows = u64::try_from(lease_update_ids.len())
2339 .map_err(|_| StoreError("ingest lease row count exceeds u64 range".into()))?;
2340 if updated_leases.rows_affected() != expected_lease_rows {
2341 return Err(IngestError::Unavailable(StoreError(format!(
2342 "ingest updated {} of {} locked lease rows",
2343 updated_leases.rows_affected(),
2344 lease_update_ids.len()
2345 ))));
2346 }
2347 }
2348
2349 let mut account_ids = Vec::with_capacity(account_deltas.len());
2350 let mut account_usage_deltas = Vec::with_capacity(account_deltas.len());
2351 let mut account_loss_deltas = Vec::with_capacity(account_deltas.len());
2352 let mut account_overage_deltas = Vec::with_capacity(account_deltas.len());
2353 for (account_id, delta) in &account_deltas {
2354 account_ids.push(account_id.clone());
2355 account_usage_deltas.push(delta.usage);
2356 account_loss_deltas.push(delta.loss);
2357 account_overage_deltas.push(delta.overage);
2358 }
2359
2360 let locked_accounts = sqlx::query(
2366 "SELECT account_id, usage_recorded, settlement_loss, overage_recorded
2367 FROM tollgate_accounts
2368 WHERE account_id = ANY($1)
2369 ORDER BY account_id FOR UPDATE",
2370 )
2371 .bind(&account_ids)
2372 .fetch_all(&mut *tx)
2373 .await
2374 .map_err(storage)?;
2375 if locked_accounts.len() != account_ids.len() {
2376 return Err(IngestError::Unavailable(StoreError(format!(
2377 "ingest locked {} of {} referenced account rows",
2378 locked_accounts.len(),
2379 account_ids.len()
2380 ))));
2381 }
2382
2383 for row in locked_accounts {
2388 let account_id: Vec<u8> = row.get(0);
2389 let usage_recorded: i64 = row.get(1);
2390 let settlement_loss: i64 = row.get(2);
2391 let overage_recorded: i64 = row.get(3);
2392 to_units(usage_recorded, "account usage_recorded")?;
2393 to_units(settlement_loss, "account settlement_loss")?;
2394 to_units(overage_recorded, "account overage_recorded")?;
2395 let delta = account_deltas.get(&account_id).ok_or_else(|| {
2396 StoreError(format!(
2397 "ingest locked unexpected account {:#034x}",
2398 id_from(&account_id)
2399 ))
2400 })?;
2401 usage_recorded.checked_add(delta.usage).ok_or_else(|| {
2402 IngestError::Refused(StoreError(format!(
2403 "usage_recorded overflow for account {:#034x}",
2404 id_from(&account_id)
2405 )))
2406 })?;
2407 overage_recorded.checked_add(delta.overage).ok_or_else(|| {
2411 IngestError::Refused(StoreError(format!(
2412 "overage_recorded overflow for account {:#034x}",
2413 id_from(&account_id)
2414 )))
2415 })?;
2416 if settlement_loss < delta.loss {
2417 return Err(IngestError::Unavailable(StoreError(format!(
2418 "settlement_loss underflow for account {:#034x}: settled straggler \
2419 usage {} exceeds recorded loss",
2420 id_from(&account_id),
2421 delta.loss
2422 ))));
2423 }
2424 }
2425
2426 let updated_accounts = sqlx::query(
2427 "UPDATE tollgate_accounts AS account
2428 SET usage_recorded = account.usage_recorded + delta.usage,
2429 settlement_loss = account.settlement_loss - delta.loss,
2430 overage_recorded = account.overage_recorded + delta.overage
2431 FROM UNNEST($1::bytea[], $2::bigint[], $3::bigint[], $4::bigint[])
2432 AS delta(account_id, usage, loss, overage)
2433 WHERE account.account_id = delta.account_id
2434 AND account.settlement_loss >= delta.loss",
2435 )
2436 .bind(&account_ids)
2437 .bind(&account_usage_deltas)
2438 .bind(&account_loss_deltas)
2439 .bind(&account_overage_deltas)
2440 .execute(&mut *tx)
2441 .await
2442 .map_err(storage)?;
2443 let expected_account_rows = u64::try_from(account_ids.len())
2444 .map_err(|_| StoreError("ingest account row count exceeds u64 range".into()))?;
2445 if updated_accounts.rows_affected() != expected_account_rows {
2446 return Err(IngestError::Unavailable(StoreError(format!(
2447 "ingest updated {} of {} locked account rows",
2448 updated_accounts.rows_affected(),
2449 account_ids.len()
2450 ))));
2451 }
2452
2453 if keys.iter().all(Option::is_none) {
2457 report.unattributed = Some(expected_event_rows);
2458 return Ok(report);
2459 }
2460
2461 let attributed: i64 = sqlx::query_scalar(
2466 "WITH matched AS MATERIALIZED (
2467 SELECT k.key_id, b.occurred_at_us
2468 FROM UNNEST($1::bytea[], $2::bytea[], $3::bigint[])
2469 AS b(key_id, account_id, occurred_at_us)
2470 JOIN LATERAL (
2471 SELECT key_id FROM tollgate_credential_keys
2472 WHERE key_id = b.key_id AND account_id = b.account_id LIMIT 1
2473 ) k ON true
2474 ), updated AS (
2475 INSERT INTO tollgate_credential_activity AS activity (key_id, last_committed_at_us)
2476 SELECT key_id, MAX(occurred_at_us) FROM matched GROUP BY key_id ORDER BY key_id
2477 ON CONFLICT (key_id) DO UPDATE
2478 SET last_committed_at_us = EXCLUDED.last_committed_at_us
2479 WHERE activity.last_committed_at_us < EXCLUDED.last_committed_at_us
2480 RETURNING key_id
2481 ) SELECT COUNT(*) FROM matched"
2482 ).bind(&keys).bind(&acct).bind(&at).fetch_one(&mut *tx).await.map_err(storage)?;
2483 report.unattributed = Some(expected_event_rows.checked_sub(
2484 u64::try_from(attributed).map_err(|_| StoreError("negative attribution count".into()))?
2485 ).ok_or_else(|| StoreError("attribution count exceeds accepted events".into()))?);
2486
2487 Ok(report)
2488 }
2489 .await;
2490 finish_transaction(tx, result).await
2494 }
2495}
2496
2497fn decode_status(stored: String) -> Result<AccountStatus, StoreError> {
2505 match stored.as_str() {
2506 s if s == AccountStatus::Active.as_str() => Ok(AccountStatus::Active),
2507 s if s == AccountStatus::Suspended.as_str() => Ok(AccountStatus::Suspended),
2508 s if s == AccountStatus::Closed.as_str() => Ok(AccountStatus::Closed),
2509 other => Err(StoreError(format!("unrecognized account status {other:?}"))),
2510 }
2511}
2512
2513fn decode_capacity_class(stored: String) -> Result<CapacityClass, StoreError> {
2522 match stored.as_str() {
2523 s if s == CapacityClass::Assured.as_str() => Ok(CapacityClass::Assured),
2524 s if s == CapacityClass::BestEffort.as_str() => Ok(CapacityClass::BestEffort),
2525 other => Err(StoreError(format!("unrecognized capacity class {other:?}"))),
2526 }
2527}
2528
2529fn decode_authority(stored: String) -> Result<AdminAuthority, StoreError> {
2536 match stored.as_str() {
2537 s if s == AdminAuthority::Operator.as_str() => Ok(AdminAuthority::Operator),
2538 s if s == AdminAuthority::Provisioner.as_str() => Ok(AdminAuthority::Provisioner),
2539 other => Err(StoreError(format!(
2540 "unrecognized admin authority {other:?}"
2541 ))),
2542 }
2543}
2544
2545fn decode_schedule(
2554 allowance: Option<i64>,
2555 period: Option<String>,
2556 rollover: Option<String>,
2557) -> Result<Option<BudgetSchedule>, StoreError> {
2558 let populated = [allowance.is_some(), period.is_some(), rollover.is_some()];
2559 let (Some(allowance), Some(period), Some(rollover)) = (allowance, period, rollover) else {
2560 if populated.iter().any(|present| *present) {
2561 return Err(StoreError(
2562 "stored budget schedule is partially populated".into(),
2563 ));
2564 }
2565 return Ok(None);
2566 };
2567 let period = match period.as_str() {
2568 s if s == Period::UtcCalendarMonth.as_str() => Period::UtcCalendarMonth,
2569 other => return Err(StoreError(format!("unrecognized budget period {other:?}"))),
2570 };
2571 let rollover = match rollover.as_str() {
2572 s if s == Rollover::None.as_str() => Rollover::None,
2573 other => {
2574 return Err(StoreError(format!(
2575 "unrecognized budget rollover {other:?}"
2576 )));
2577 }
2578 };
2579 Ok(Some(BudgetSchedule {
2580 allowance: to_units(allowance, "budget allowance")?,
2581 period,
2582 rollover,
2583 }))
2584}
2585
2586const BUDGET_VIEW_SQL: &str = "SELECT account_id, deposited, overage_recorded, usage_recorded,
2589 settlement_loss, expired, budget_allowance, budget_period,
2590 budget_rollover, period_start_us
2591 FROM tollgate_accounts WHERE account_id = $1";
2592
2593const GRANT_DEBIT_SQL: &str = "UPDATE tollgate_accounts
2596 SET balance = balance - $2,
2597 allowance_balance = allowance_balance - $3,
2598 next_fence = next_fence + 1
2599 WHERE account_id = $1
2600 RETURNING account_id, deposited, overage_recorded, usage_recorded,
2601 settlement_loss, expired, budget_allowance, budget_period,
2602 budget_rollover, period_start_us";
2603
2604fn budget_view(row: &sqlx::postgres::PgRow) -> Result<BudgetView, StoreError> {
2619 let deposited = to_units(row.get::<i64, _>(1), "deposited")?;
2620 let overage = to_units(row.get::<i64, _>(2), "overage_recorded")?;
2621 let usage = to_units(row.get::<i64, _>(3), "usage_recorded")?;
2622 let loss = to_units(row.get::<i64, _>(4), "settlement_loss")?;
2623 let expired = to_units(row.get::<i64, _>(5), "expired")?;
2624 let schedule = decode_schedule(row.get(6), row.get(7), row.get(8))?;
2625 let period_start = micros_ts(row.get::<i64, _>(9), "period_start_us")?;
2626
2627 let funded = deposited
2628 .checked_add(overage)
2629 .ok_or_else(|| StoreError("account funding total overflows".into()))?;
2630 let consumed = usage
2631 .checked_add(loss)
2632 .and_then(|spent| spent.checked_add(expired))
2633 .ok_or_else(|| StoreError("account consumption total overflows".into()))?;
2634 Ok(BudgetView {
2635 balance_at_publish: funded.checked_sub(consumed).ok_or_else(|| {
2636 StoreError(format!(
2637 "consumption {} exceeds funding {}",
2638 consumed.get(),
2639 funded.get()
2640 ))
2641 })?,
2642 period_end: schedule.map(|schedule| schedule.period.end_after(period_start)),
2647 })
2648}
2649
2650fn decode_publishable(
2662 principal: Principal,
2663 generation: i64,
2664 value: serde_json::Value,
2665) -> Result<PublishableSnapshot, StoreError> {
2666 let generation = generation_from(generation)?;
2667 let snapshot: StoredSnapshot =
2668 serde_json::from_value(value).map_err(|e| StoreError(format!("snapshot decode: {e}")))?;
2669 let budget = snapshot.budget;
2674 let publishable = PublishableSnapshot::try_new(Arc::new(snapshot.into_snapshot(generation)))
2675 .map_err(|error| {
2676 StoreError(format!(
2677 "invalid stored snapshot for principal {:#034x}: {error}",
2678 principal.0
2679 ))
2680 })?;
2681 Ok(match budget {
2682 Some(budget) => publishable.with_budget(Some(budget)),
2683 None => publishable,
2684 })
2685}
2686
2687fn generation_from(column: i64) -> Result<Generation, StoreError> {
2699 u64::try_from(column)
2700 .map(Generation)
2701 .map_err(|_| StoreError("stored snapshot generation is negative".into()))
2702}
2703
2704#[async_trait]
2705impl SnapshotSource for PostgresStore {
2706 async fn snapshot(&self, principal: Principal) -> Result<SnapshotResolution, StoreError> {
2707 let row = sqlx::query(
2708 "SELECT generation, snapshot, deleted FROM tollgate_snapshots WHERE principal = $1",
2709 )
2710 .bind(id_bytes(principal.0))
2711 .fetch_optional(&self.pool)
2712 .await
2713 .map_err(storage)?;
2714 match row {
2715 Some(row) if row.get::<bool, _>(2) => Ok(SnapshotResolution::Revoked {
2721 generation: generation_from(row.get::<i64, _>(0))?,
2722 }),
2723 Some(row) => Ok(SnapshotResolution::Present(decode_publishable(
2724 principal,
2725 row.get::<i64, _>(0),
2726 row.get(1),
2727 )?)),
2728 None => Ok(SnapshotResolution::Unknown),
2729 }
2730 }
2731
2732 fn subscribe(&self) -> broadcast::Receiver<SnapshotPush> {
2733 self.push.subscribe()
2734 }
2735
2736 async fn principals(&self) -> Result<Option<Vec<Principal>>, StoreError> {
2744 let rows = sqlx::query("SELECT principal FROM tollgate_snapshots ORDER BY principal")
2745 .fetch_all(&self.pool)
2746 .await
2747 .map_err(storage)?;
2748 Ok(Some(
2749 rows.iter()
2750 .map(|row| Principal(id_from(row.get::<Vec<u8>, _>(0).as_slice())))
2751 .collect(),
2752 ))
2753 }
2754}
2755
2756#[async_trait]
2757impl StoreHealth for PostgresStore {
2758 async fn ping(&self) -> Result<(), StoreError> {
2759 sqlx::query("SELECT 1")
2760 .execute(&self.pool)
2761 .await
2762 .map_err(storage)?;
2763 Ok(())
2764 }
2765}
2766
2767async fn republish_patched_snapshots(
2797 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
2798 account: AccountId,
2799 json_path: &'static str,
2800 value: &str,
2801) -> Result<(Vec<(Principal, PublishableSnapshot)>, usize), SetStatusError> {
2802 let rows = sqlx::query(
2803 "UPDATE tollgate_snapshots
2804 SET generation = generation + 1,
2805 snapshot = jsonb_set(snapshot, $3::text[], to_jsonb($2::text))
2806 WHERE account_id = $1
2807 AND deleted = FALSE
2808 AND snapshot #>> $3::text[] IS DISTINCT FROM $2::text
2809 RETURNING principal, generation, snapshot",
2810 )
2811 .bind(id_bytes(account.0))
2812 .bind(value)
2813 .bind(json_path)
2814 .fetch_all(&mut **tx)
2815 .await
2816 .map_err(storage)?;
2817
2818 let mut republished = Vec::with_capacity(rows.len());
2819 let mut unreadable = 0usize;
2820 for row in rows {
2821 let principal = Principal(id_from(row.get::<Vec<u8>, _>(0).as_slice()));
2822 match decode_publishable(
2823 principal,
2824 row.get::<i64, _>(1),
2825 row.get::<serde_json::Value, _>(2),
2826 ) {
2827 Ok(snapshot) => republished.push((principal, snapshot)),
2828 Err(error) => {
2836 unreadable += 1;
2837 tracing::warn!(
2838 %principal,
2839 %error,
2840 "restamped snapshot could not be decoded for push"
2841 );
2842 }
2843 }
2844 }
2845 republished.sort_unstable_by_key(|(principal, _)| *principal);
2848 Ok((republished, unreadable))
2849}
2850#[async_trait]
2851impl AdminStore for PostgresStore {
2852 async fn create_account(
2853 &self,
2854 config: AccountConfig,
2855 ) -> Result<tollgate_store::AdminReceipt<()>, CreateAccountError> {
2856 self.create_account_as(config, AdminAuthority::Operator)
2857 .await
2858 }
2859
2860 async fn create_provisioned_account(
2861 &self,
2862 account: AccountId,
2863 ) -> Result<tollgate_store::AdminReceipt<()>, CreateAccountError> {
2864 self.create_account_as(
2865 AccountConfig {
2866 account_id: account,
2867 initial_balance: CostUnits::ZERO,
2868 status: AccountStatus::Suspended,
2869 capacity_class: CapacityClass::BestEffort,
2870 },
2871 AdminAuthority::Provisioner,
2872 )
2873 .await
2874 }
2875
2876 async fn deposit(
2877 &self,
2878 account: AccountId,
2879 units: CostUnits,
2880 ) -> Result<AdminReceipt<()>, AllocateError> {
2881 let row = sqlx::query(
2882 "UPDATE tollgate_accounts
2883 SET balance = balance + $2, deposited = deposited + $2
2884 WHERE account_id = $1
2885 RETURNING balance - $2 AS old_topup, deposited - $2 AS old_deposited,
2886 balance AS new_topup, deposited AS new_deposited",
2887 )
2888 .bind(id_bytes(account.0))
2889 .bind(i64::try_from(units.get()).map_err(|_| AllocateError::BalanceOverflow)?)
2890 .fetch_optional(&self.pool)
2891 .await
2892 .map_err(|error| {
2893 if error
2897 .as_database_error()
2898 .and_then(|db| db.code())
2899 .as_deref()
2900 == Some("22003")
2901 {
2902 AllocateError::BalanceOverflow
2903 } else {
2904 alloc_storage(error)
2905 }
2906 })?
2907 .ok_or(AllocateError::UnknownAccount)?;
2908 let state = |topup: &str, deposited: &str| -> Result<AdminState, StoreError> {
2909 Ok(AdminState::Funding {
2910 topup: to_units(row.get(topup), "audit topup")?,
2911 deposited: to_units(row.get(deposited), "audit deposited")?,
2912 })
2913 };
2914 Ok(AdminReceipt::new(
2915 (),
2916 state("old_topup", "old_deposited")?,
2917 state("new_topup", "new_deposited")?,
2918 ))
2919 }
2920
2921 async fn set_budget_schedule(
2922 &self,
2923 account: AccountId,
2924 schedule: Option<BudgetSchedule>,
2925 ) -> Result<AdminReceipt<()>, BudgetError> {
2926 let allowance = schedule
2936 .map(|s| to_i64(s.allowance, "budget allowance"))
2937 .transpose()
2938 .map_err(BudgetError::Storage)?;
2939 let mut tx = self.pool.begin().await.map_err(storage)?;
2940 let result = async {
2941 let row = sqlx::query(
2945 "SELECT budget_allowance, budget_period, budget_rollover
2946 FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE",
2947 )
2948 .bind(id_bytes(account.0))
2949 .fetch_optional(&mut *tx)
2950 .await
2951 .map_err(storage)?
2952 .ok_or(BudgetError::UnknownAccount)?;
2953 let before = decode_schedule(row.get(0), row.get(1), row.get(2))?;
2954 sqlx::query(
2955 "UPDATE tollgate_accounts
2956 SET budget_allowance = $2, budget_period = $3, budget_rollover = $4
2957 WHERE account_id = $1",
2958 )
2959 .bind(id_bytes(account.0))
2960 .bind(allowance)
2961 .bind(schedule.map(|s| s.period.as_str()))
2962 .bind(schedule.map(|s| s.rollover.as_str()))
2963 .execute(&mut *tx)
2964 .await
2965 .map_err(storage)?;
2966 Ok(AdminReceipt::new(
2967 (),
2968 AdminState::Budget { schedule: before },
2969 AdminState::Budget { schedule },
2970 ))
2971 }
2972 .await;
2973 finish_transaction(tx, result).await
2974 }
2975
2976 async fn account_view(&self, account: AccountId) -> Result<Option<AccountView>, StoreError> {
2977 let mut tx = self.pool.begin().await.map_err(storage)?;
2986 let result = async {
2987 sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY")
2988 .execute(&mut *tx)
2989 .await
2990 .map_err(storage)?;
2991 let account_row = sqlx::query(
2992 "SELECT deposited, balance, usage_recorded, settlement_loss, overage_recorded,
2993 expired, status, capacity_class, budget_allowance, budget_period,
2994 budget_rollover, period_start_us, origin, status_set_by
2995 FROM tollgate_accounts WHERE account_id = $1",
2996 )
2997 .bind(id_bytes(account.0))
2998 .fetch_optional(&mut *tx)
2999 .await
3000 .map_err(storage)?;
3001 let lease_row = sqlx::query(ACTIVE_LEASE_SUM_SQL)
3002 .bind(id_bytes(account.0))
3003 .fetch_one(&mut *tx)
3004 .await
3005 .map_err(storage)?;
3006 Ok::<_, StoreError>((account_row, lease_row))
3007 }
3008 .await;
3009 let (account_row, lease_row) = finish_transaction(tx, result).await?;
3010
3011 let Some(row) = account_row else {
3012 return Ok(None);
3013 };
3014 let active_grants = to_units(lease_row.get::<i64, _>(0), "active lease grants")?;
3015 let active_used = to_units(lease_row.get::<i64, _>(1), "active lease usage")?;
3016 let recorded = to_units(row.get::<i64, _>(2), "usage_recorded")?;
3017 Ok(Some(AccountView {
3018 account_id: account,
3019 status: decode_status(row.get::<String, _>(6))?,
3020 capacity_class: decode_capacity_class(row.get::<String, _>(7))?,
3021 origin: decode_authority(row.get::<String, _>(12))?,
3022 status_set_by: decode_authority(row.get::<String, _>(13))?,
3023 schedule: decode_schedule(
3024 row.get::<Option<i64>, _>(8),
3025 row.get::<Option<String>, _>(9),
3026 row.get::<Option<String>, _>(10),
3027 )?,
3028 period_start: StoredInstant {
3029 micros: row.get::<i64, _>(11),
3030 submicro_nanos: 0,
3031 }
3032 .timestamp()?,
3033 conservation: Conservation {
3034 deposited: to_units(row.get::<i64, _>(0), "deposited")?,
3035 overage_recorded: to_units(row.get::<i64, _>(4), "overage_recorded")?,
3036 balance: to_units(row.get::<i64, _>(1), "balance")?,
3037 active_lease_grants: active_grants,
3038 settled_usage: recorded.checked_sub(active_used).ok_or_else(|| {
3042 StoreError(format!(
3043 "active lease usage {} exceeds recorded usage {} for account {account}",
3044 active_used.get(),
3045 recorded.get()
3046 ))
3047 })?,
3048 settlement_loss: to_units(row.get::<i64, _>(3), "settlement_loss")?,
3049 expired: to_units(row.get::<i64, _>(5), "expired")?,
3050 },
3051 }))
3052 }
3053
3054 async fn roll_due_periods(
3055 &self,
3056 now: Timestamp,
3057 limit: NonZeroUsize,
3058 ) -> Result<RolloverBatch, StoreError> {
3059 let limit_i = i64::try_from(limit.get())
3060 .map_err(|_| StoreError(format!("rollover batch limit exceeds i64 range: {limit}")))?;
3061 let mut rolled = Vec::new();
3062 for period in Period::ALL {
3069 let boundary_us = ts_micros(period.start_of(now));
3070 let remaining = limit_i
3071 - i64::try_from(rolled.len())
3072 .map_err(|_| StoreError("rollover batch row count exceeds i64 range".into()))?;
3073 if remaining <= 0 {
3074 break;
3075 }
3076 let rows = sqlx::query(DUE_PERIODS_SQL)
3087 .bind(period.as_str())
3088 .bind(boundary_us)
3089 .bind(remaining)
3090 .fetch_all(&self.pool)
3091 .await
3092 .map_err(storage)?;
3093
3094 for row in rows {
3095 rolled.push((
3096 row.get::<i64, _>(3),
3097 RolledAccount {
3098 account_id: AccountId(id_from(&row.get::<Vec<u8>, _>(0))),
3099 deposited: to_units(row.get::<i64, _>(1), "budget allowance")?,
3100 expired: to_units(row.get::<i64, _>(2), "expiring allowance")?,
3101 },
3102 ));
3103 }
3104 }
3105 rolled
3112 .sort_unstable_by_key(|(crossed_from, account)| (*crossed_from, account.account_id.0));
3113 RolloverBatch::try_new(
3114 rolled.into_iter().map(|(_, account)| account).collect(),
3115 limit,
3116 )
3117 }
3118
3119 async fn set_account_status(
3120 &self,
3121 account: AccountId,
3122 status: AccountStatus,
3123 ) -> Result<tollgate_store::AdminReceipt<StatusChange>, SetStatusError> {
3124 self.set_status_as(account, status, AdminAuthority::Operator)
3125 .await
3126 }
3127
3128 async fn activate_provisioned(
3129 &self,
3130 account: AccountId,
3131 ) -> Result<tollgate_store::AdminReceipt<StatusChange>, SetStatusError> {
3132 self.set_status_as(account, AccountStatus::Active, AdminAuthority::Provisioner)
3133 .await
3134 }
3135
3136 async fn set_capacity_class(
3137 &self,
3138 account: AccountId,
3139 class: CapacityClass,
3140 ) -> Result<tollgate_store::AdminReceipt<StatusChange>, SetStatusError> {
3141 let mut tx = self.pool.begin().await.map_err(storage)?;
3142 let result = async {
3143 let row = sqlx::query(
3151 "SELECT status, capacity_class FROM tollgate_accounts WHERE account_id = $1 FOR UPDATE",
3152 )
3153 .bind(id_bytes(account.0))
3154 .fetch_optional(&mut *tx)
3155 .await
3156 .map_err(storage)?
3157 .ok_or(SetStatusError::UnknownAccount)?;
3158
3159 let before = decode_capacity_class(row.get::<String, _>(1))?;
3160 if decode_status(row.get::<String, _>(0))? == AccountStatus::Closed {
3165 return Err(SetStatusError::AccountClosed);
3166 }
3167
3168 sqlx::query("UPDATE tollgate_accounts SET capacity_class = $2 WHERE account_id = $1")
3169 .bind(id_bytes(account.0))
3170 .bind(class.as_str())
3171 .execute(&mut *tx)
3172 .await
3173 .map_err(storage)?;
3174
3175 let (republished, unreadable) =
3176 republish_patched_snapshots(&mut tx, account, "{capacity_class}", class.as_str()).await?;
3177 Ok((before, republished, unreadable))
3178 }
3179 .await;
3180
3181 let (before, republished, unreadable) = finish_transaction(tx, result).await?;
3182 if pushes_exceed_capacity(republished.len()) {
3183 tracing::warn!(
3184 %account,
3185 principals = republished.len(),
3186 capacity = PUSH_CHANNEL_CAPACITY,
3187 "capacity class change emitted more pushes than the channel holds; \
3188 subscribers will resync"
3189 );
3190 }
3191 let count = republished.len();
3192 for (principal, snapshot) in republished {
3193 self.push_to_subscribers(SnapshotPush {
3194 principal,
3195 resolution: SnapshotResolution::Present(snapshot),
3196 });
3197 }
3198 Ok(AdminReceipt::new(
3199 StatusChange {
3200 republished: count + unreadable,
3203 unreadable,
3204 },
3205 AdminState::CapacityClass {
3206 capacity_class: before,
3207 },
3208 AdminState::CapacityClass {
3209 capacity_class: class,
3210 },
3211 ))
3212 }
3213
3214 async fn publish_snapshot(
3215 &self,
3216 principal: Principal,
3217 snapshot: PublishableSnapshot,
3218 ) -> Result<tollgate_store::AdminReceipt<()>, PublishSnapshotError> {
3219 let generation = i64::try_from(snapshot.generation.0).map_err(|_| {
3220 StoreError("snapshot generation exceeds PostgreSQL BIGINT range".into())
3221 })?;
3222
3223 let mut tx = self.pool.begin().await.map_err(storage)?;
3224 let result = publish_in_tx(&mut tx, principal, Some(generation), snapshot).await;
3225
3226 let (written, published, before, after) = finish_transaction(tx, result).await?;
3227 if written {
3228 self.push_to_subscribers(SnapshotPush {
3232 principal,
3233 resolution: SnapshotResolution::Present(published),
3234 });
3235 }
3236 Ok(AdminReceipt::new((), before, after))
3237 }
3238
3239 async fn remove_snapshot(&self, principal: Principal) -> Result<AdminReceipt<()>, StoreError> {
3240 let mut tx = self.pool.begin().await.map_err(storage)?;
3241 let result = remove_in_tx(&mut tx, principal).await;
3242 let receipt = finish_transaction(tx, result).await?;
3243 self.announce_removal(principal, &receipt);
3244 Ok(receipt)
3245 }
3246}
3247
3248async fn lock_account_key(
3272 tx: &mut Transaction<'_, Postgres>,
3273 account: AccountId,
3274 key: KeyId,
3275) -> Result<(Principal, bool), KeySnapshotError> {
3276 let row = sqlx::query(
3277 "SELECT principal, revoked_at_us IS NOT NULL
3278 FROM tollgate_credential_keys WHERE key_id = $1 AND account_id = $2 FOR SHARE",
3279 )
3280 .bind(id_bytes(key.0))
3281 .bind(id_bytes(account.0))
3282 .fetch_optional(&mut **tx)
3283 .await
3284 .map_err(storage)?
3285 .ok_or(KeySnapshotError::UnknownCredential)?;
3286 let principal: Vec<u8> = row.get(0);
3287 let principal: [u8; 16] = principal
3288 .try_into()
3289 .map_err(|_| StoreError("credential principal is not 16 bytes".into()))?;
3290 Ok((Principal(u128::from_be_bytes(principal)), row.get(1)))
3291}
3292
3293async fn publish_in_tx(
3301 tx: &mut Transaction<'_, Postgres>,
3302 principal: Principal,
3303 generation: Option<i64>,
3304 snapshot: PublishableSnapshot,
3305) -> Result<(bool, PublishableSnapshot, AdminState, AdminState), PublishSnapshotError> {
3306 if let Some(key_id) = snapshot.key_id {
3307 let matches: bool = sqlx::query_scalar(
3308 "SELECT EXISTS(SELECT 1 FROM tollgate_credential_keys
3309 WHERE key_id = $1 AND principal = $2 AND account_id = $3)",
3310 )
3311 .bind(id_bytes(key_id.0))
3312 .bind(id_bytes(principal.0))
3313 .bind(id_bytes(snapshot.account_id.0))
3314 .fetch_one(&mut **tx)
3315 .await
3316 .map_err(storage)?;
3317 if !matches {
3318 return Err(PublishSnapshotError::CredentialMismatch { key_id });
3319 }
3320 }
3321 let ledger = sqlx::query(
3333 "SELECT status, deposited, overage_recorded, usage_recorded, settlement_loss,
3334 expired, budget_allowance, budget_period, budget_rollover, period_start_us,
3335 capacity_class
3336 FROM tollgate_accounts WHERE account_id = $1 FOR SHARE",
3337 )
3338 .bind(id_bytes(snapshot.account_id.0))
3339 .fetch_optional(&mut **tx)
3340 .await
3341 .map_err(storage)?;
3342
3343 let view = match &ledger {
3351 Some(row) => {
3352 let ledger = decode_status(row.get::<String, _>(0))?;
3353 if ledger != snapshot.status {
3354 return Err(PublishSnapshotError::StatusMismatch {
3355 ledger,
3356 submitted: snapshot.status,
3357 });
3358 }
3359 let ledger_class = decode_capacity_class(row.get::<String, _>(10))?;
3363 if ledger_class != snapshot.capacity_class {
3364 return Err(PublishSnapshotError::CapacityClassMismatch {
3365 ledger: ledger_class,
3366 submitted: snapshot.capacity_class,
3367 });
3368 }
3369 Some(budget_view(row)?)
3370 }
3371 None => None,
3372 };
3373 let published = snapshot.with_budget(view);
3374 let value = serde_json::to_value(StoredSnapshotRef::from(published.as_snapshot()))
3375 .map_err(|e| StoreError(format!("snapshot encode: {e}")))?;
3376
3377 let (written, before, after) = write_snapshot_audited(tx, principal, generation, value).await?;
3378 let published = if generation.is_none() {
3379 let AdminState::Snapshot { generation, .. } = after else {
3380 return Err(StoreError("published snapshot has no generation".into()).into());
3381 };
3382 published.restamped(published.status, generation)
3383 } else {
3384 published
3385 };
3386 Ok((written, published, before, after))
3387}
3388
3389async fn remove_in_tx(
3392 tx: &mut Transaction<'_, Postgres>,
3393 principal: Principal,
3394) -> Result<AdminReceipt<()>, StoreError> {
3395 let before = snapshot_audit_row(tx, principal).await?;
3396 let after = match before {
3397 AdminState::Snapshot {
3398 generation,
3399 revoked: false,
3400 } => {
3401 sqlx::query("UPDATE tollgate_snapshots SET deleted = TRUE WHERE principal = $1")
3402 .bind(id_bytes(principal.0))
3403 .execute(&mut **tx)
3404 .await
3405 .map_err(storage)?;
3406 AdminState::Snapshot {
3407 generation,
3408 revoked: true,
3409 }
3410 }
3411 state => state,
3412 };
3413 Ok(AdminReceipt::new((), before, after))
3414}
3415
3416async fn snapshot_audit_row(
3419 tx: &mut Transaction<'_, Postgres>,
3420 principal: Principal,
3421) -> Result<AdminState, StoreError> {
3422 let row = sqlx::query(
3423 "SELECT generation, deleted FROM tollgate_snapshots WHERE principal = $1 FOR UPDATE",
3424 )
3425 .bind(id_bytes(principal.0))
3426 .fetch_optional(&mut **tx)
3427 .await
3428 .map_err(storage)?;
3429 row.map(|row| {
3430 Ok(AdminState::Snapshot {
3431 generation: generation_from(row.get(0))?,
3432 revoked: row.get(1),
3433 })
3434 })
3435 .unwrap_or(Ok(AdminState::Absent))
3436}
3437
3438async fn write_snapshot_audited(
3439 tx: &mut Transaction<'_, Postgres>,
3440 principal: Principal,
3441 generation: Option<i64>,
3442 value: serde_json::Value,
3443) -> Result<(bool, AdminState, AdminState), StoreError> {
3444 let mut before = snapshot_audit_row(tx, principal).await?;
3445 if before == AdminState::Absent {
3446 let initial = generation.unwrap_or(1);
3447 let after = AdminState::Snapshot {
3448 generation: generation_from(initial)?,
3449 revoked: false,
3450 };
3451 let inserted = sqlx::query(
3452 "INSERT INTO tollgate_snapshots (principal, generation, snapshot, deleted)
3453 VALUES ($1, $2, $3, FALSE) ON CONFLICT (principal) DO NOTHING",
3454 )
3455 .bind(id_bytes(principal.0))
3456 .bind(initial)
3457 .bind(&value)
3458 .execute(&mut **tx)
3459 .await
3460 .map_err(storage)?;
3461 if inserted.rows_affected() == 1 {
3462 return Ok((true, before, after));
3463 }
3464 before = snapshot_audit_row(tx, principal).await?;
3467 }
3468 let AdminState::Snapshot {
3469 generation: previous,
3470 ..
3471 } = before
3472 else {
3473 return Err(StoreError("snapshot disappeared during publication".into()));
3474 };
3475 let generation = match generation {
3478 Some(stated) => stated,
3479 None => i64::try_from(previous.0)
3480 .ok()
3481 .and_then(|previous| previous.checked_add(1))
3482 .ok_or_else(|| StoreError("snapshot generation overflow".into()))?,
3483 };
3484 let after = AdminState::Snapshot {
3485 generation: generation_from(generation)?,
3486 revoked: false,
3487 };
3488 if previous >= generation_from(generation)? {
3489 return Ok((false, before, before));
3490 }
3491 sqlx::query("UPDATE tollgate_snapshots SET generation = $2, snapshot = $3, deleted = FALSE WHERE principal = $1")
3492 .bind(id_bytes(principal.0)).bind(generation).bind(value)
3493 .execute(&mut **tx).await.map_err(storage)?;
3494 Ok((true, before, after))
3495}
3496
3497#[cfg(test)]
3498mod tests {
3499 use super::*;
3500
3501 #[test]
3502 fn stored_authorities_reject_unknown_vocabulary() {
3503 assert_eq!(
3504 decode_authority("Operator".into()).unwrap(),
3505 AdminAuthority::Operator
3506 );
3507 assert_eq!(
3508 decode_authority("Provisioner".into()).unwrap(),
3509 AdminAuthority::Provisioner
3510 );
3511 for invalid in [
3512 "",
3513 "operator",
3514 "provisioner",
3515 "Administrator",
3516 "Provisioner ",
3517 ] {
3518 assert!(decode_authority(invalid.into()).is_err(), "{invalid:?}");
3519 }
3520 }
3521
3522 #[test]
3523 fn stored_fences_use_the_exact_positive_bigint_domain() {
3524 for invalid in [i64::MIN, -1, 0] {
3525 assert!(stored_fence(invalid).is_err());
3526 }
3527 for valid in [1, 2, i64::MAX] {
3528 assert_eq!(stored_fence(valid).unwrap(), FencingToken(valid as u64));
3529 }
3530 }
3531
3532 #[tokio::test]
3535 async fn ping_surfaces_a_closed_pool() {
3536 let pool = PgPoolOptions::new()
3537 .connect_lazy("postgres://localhost/tollgate")
3538 .unwrap();
3539 pool.close().await;
3540 let (push, _) = broadcast::channel(1);
3541 let store = PostgresStore {
3542 pool,
3543 policy: GrantPolicy::default(),
3544 push,
3545 };
3546
3547 assert!(store.ping().await.is_err());
3548 assert!(matches!(
3549 AdminStore::deposit(&store, AccountId(1), CostUnits(1)).await,
3550 Err(AllocateError::Storage(_))
3551 ));
3552 }
3553}