70 lines
2.6 KiB
Python
70 lines
2.6 KiB
Python
"""Durable atomic replay protection owned by the Brain connector."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from pathlib import Path
|
|
|
|
|
|
class SQLiteReplayCache:
|
|
"""SQLite implementation of the machine-identity ``ReplayCache`` protocol.
|
|
|
|
A primary key makes consumption atomic across threads and processes. Entries
|
|
remain durable across connector restarts until their verifier expiry passes.
|
|
"""
|
|
|
|
def __init__(self, database: str | Path, *, timeout_seconds: float = 5.0) -> None:
|
|
self._database = str(Path(database))
|
|
self._timeout_seconds = timeout_seconds
|
|
if timeout_seconds <= 0:
|
|
raise ValueError("SQLite replay timeout must be positive")
|
|
with self._connect() as connection:
|
|
connection.execute("PRAGMA journal_mode=WAL")
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS machine_token_replay (
|
|
principal TEXT NOT NULL,
|
|
token_id TEXT NOT NULL,
|
|
expires_at INTEGER NOT NULL,
|
|
PRIMARY KEY (principal, token_id)
|
|
) WITHOUT ROWID
|
|
"""
|
|
)
|
|
|
|
def consume(self, principal: str, token_id: str, expires_at: int, now: int) -> bool:
|
|
if not principal or not token_id:
|
|
raise ValueError("replay identity must be non-empty")
|
|
if any(isinstance(value, bool) or not isinstance(value, int) for value in (expires_at, now)):
|
|
raise ValueError("replay timestamps must be integers")
|
|
if expires_at <= now:
|
|
return False
|
|
|
|
connection = self._connect()
|
|
try:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
connection.execute("DELETE FROM machine_token_replay WHERE expires_at <= ?", (now,))
|
|
try:
|
|
connection.execute(
|
|
"INSERT INTO machine_token_replay (principal, token_id, expires_at) VALUES (?, ?, ?)",
|
|
(principal, token_id, expires_at),
|
|
)
|
|
except sqlite3.IntegrityError:
|
|
connection.rollback()
|
|
return False
|
|
connection.commit()
|
|
return True
|
|
except BaseException:
|
|
connection.rollback()
|
|
raise
|
|
finally:
|
|
connection.close()
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
connection = sqlite3.connect(
|
|
self._database,
|
|
timeout=self._timeout_seconds,
|
|
isolation_level=None,
|
|
)
|
|
connection.execute(f"PRAGMA busy_timeout={int(self._timeout_seconds * 1000)}")
|
|
return connection
|