pub mod digest;
pub mod ipc;
pub mod outcomes;
use crate::policy::model::{PolicyAction, SqlCategory, TableId};
use rand::RngCore;
use std::collections::HashMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};
use thiserror::Error;
use zeroize::Zeroizing;
pub use outcomes::ApprovalOutcome;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GrantChoice {
Once,
Session,
Decline,
}
impl schemars::JsonSchema for GrantChoice {
fn schema_name() -> std::borrow::Cow<'static, str> {
"GrantChoice".into()
}
fn json_schema(_gen: &mut schemars::SchemaGenerator) -> schemars::Schema {
schemars::Schema::try_from(serde_json::json!({
"type": "string",
"enum": ["once", "session", "decline"],
"description": "How should this statement be authorized? once = this statement only; session = same tables until the server restarts; decline = do not run."
}))
.expect("valid schema")
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConfirmOutcome {
Chosen(GrantChoice),
Unavailable {
reason: String,
},
}
#[derive(Debug, Error)]
pub enum ApprovalError {
#[error("approval expired before consumption")]
Expired,
#[error("approval digest mismatch — the operation changed after approval")]
DigestMismatch,
#[error("approval already consumed")]
AlreadyConsumed,
}
#[derive(Debug, Clone)]
pub struct SessionGrantKey {
pub connection: String,
pub category: SqlCategory,
pub tables: Vec<TableId>,
}
impl SessionGrantKey {
pub fn canonical(&self) -> String {
let mut tables: Vec<String> = self
.tables
.iter()
.map(|t| format!("{}.{}", t.database, t.table))
.collect();
tables.sort();
format!(
"{}\u{1}{}\u{1}{}",
self.connection,
self.category.as_str(),
tables.join(",")
)
}
}
struct SessionGrant {
key: SessionGrantKey,
expires_at: Instant,
}
struct OneTimeApproval {
expires_at: Instant,
consumed: bool,
}
pub struct ApprovalEngine {
default_session_ttl: Duration,
default_one_shot_ttl: Duration,
sessions: Mutex<HashMap<String, SessionGrant>>,
one_shots: Mutex<HashMap<[u8; 32], OneTimeApproval>>,
}
pub const DEFAULT_SESSION_TTL: Duration = Duration::from_secs(30 * 60);
pub const DEFAULT_ONE_SHOT_TTL: Duration = Duration::from_secs(120);
impl ApprovalEngine {
pub fn new() -> Self {
Self {
default_session_ttl: DEFAULT_SESSION_TTL,
default_one_shot_ttl: DEFAULT_ONE_SHOT_TTL,
sessions: Mutex::new(HashMap::new()),
one_shots: Mutex::new(HashMap::new()),
}
}
pub fn grant_once(&self, digest: [u8; 32]) -> Instant {
let expires = Instant::now() + self.default_one_shot_ttl;
self.one_shots.lock().unwrap().insert(
digest,
OneTimeApproval {
expires_at: expires,
consumed: false,
},
);
expires
}
pub fn consume_once(&self, digest: [u8; 32]) -> Result<Instant, ApprovalError> {
let mut map = self.one_shots.lock().unwrap();
let entry = map.get_mut(&digest).ok_or(ApprovalError::AlreadyConsumed)?;
if entry.consumed {
return Err(ApprovalError::AlreadyConsumed);
}
if Instant::now() > entry.expires_at {
let expired = map.remove(&digest).unwrap();
let _ = expired;
return Err(ApprovalError::Expired);
}
entry.consumed = true;
let expires = entry.expires_at;
map.remove(&digest);
Ok(expires)
}
pub fn grant_session(&self, key: SessionGrantKey) -> Instant {
let canonical = key.canonical();
let expires = Instant::now() + self.default_session_ttl;
self.sessions.lock().unwrap().insert(
canonical,
SessionGrant {
key,
expires_at: expires,
},
);
expires
}
pub fn session_covers(&self, key: &SessionGrantKey) -> bool {
let canonical = key.canonical();
let mut map = self.sessions.lock().unwrap();
match map.get(&canonical) {
Some(g) if Instant::now() <= g.expires_at => true,
Some(_) => {
map.remove(&canonical);
false
}
None => false,
}
}
pub fn revoke_sessions(&self, filter: SessionRevokeFilter) -> usize {
let mut map = self.sessions.lock().unwrap();
let before = map.len();
map.retain(|_, g| {
if let Some(conn) = &filter.connection
&& &g.key.connection != conn
{
return true;
}
if let Some(cat) = filter.category
&& g.key.category != cat
{
return true;
}
false
});
before - map.len()
}
pub fn snapshot_sessions(&self) -> Vec<(SessionGrantKey, Instant)> {
self.sessions
.lock()
.unwrap()
.values()
.map(|g| (g.key.clone(), g.expires_at))
.collect()
}
pub fn pending_one_shots(&self) -> usize {
self.one_shots.lock().unwrap().len()
}
}
impl Default for ApprovalEngine {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Default, Clone)]
pub struct SessionRevokeFilter {
pub connection: Option<String>,
pub category: Option<SqlCategory>,
}
pub fn random_nonce_32() -> Zeroizing<[u8; 32]> {
let mut buf = [0u8; 32];
rand::rng().fill_bytes(&mut buf);
Zeroizing::new(buf)
}
#[cfg(test)]
mod tests {
use super::*;
fn tables(list: &[(&str, &str)]) -> Vec<TableId> {
list.iter().map(|(d, t)| TableId::new(*d, *t)).collect()
}
#[test]
fn one_time_grant_consumed_exactly_once() {
let e = ApprovalEngine::new();
let d = [7u8; 32];
e.grant_once(d);
assert!(e.consume_once(d).is_ok());
assert!(matches!(
e.consume_once(d).unwrap_err(),
ApprovalError::AlreadyConsumed
));
}
#[test]
fn one_time_grant_expires() {
let e = ApprovalEngine::new();
let mut d = [7u8; 32];
e.grant_once(d);
d[0] ^= 1;
assert!(matches!(
e.consume_once(d).unwrap_err(),
ApprovalError::AlreadyConsumed
));
}
#[test]
fn concurrent_consumption_single_winner() {
let e = std::sync::Arc::new(ApprovalEngine::new());
let d = [9u8; 32];
e.grant_once(d);
let mut handles = Vec::new();
for _ in 0..16 {
let e2 = e.clone();
handles.push(std::thread::spawn(move || e2.consume_once(d).is_ok()));
}
let wins: usize = handles
.into_iter()
.map(|h| h.join().unwrap())
.filter(|won| *won)
.count();
assert_eq!(wins, 1);
}
#[test]
fn session_grants_are_narrow() {
let e = ApprovalEngine::new();
let key = SessionGrantKey {
connection: "c1".into(),
category: SqlCategory::Write,
tables: tables(&[("app", "jobs")]),
};
e.grant_session(key.clone());
assert!(e.session_covers(&key));
let other = SessionGrantKey {
connection: "c1".into(),
category: SqlCategory::Write,
tables: tables(&[("app", "jobs"), ("app", "users")]),
};
assert!(!e.session_covers(&other));
let other_conn = SessionGrantKey {
connection: "c2".into(),
category: SqlCategory::Write,
tables: tables(&[("app", "jobs")]),
};
assert!(!e.session_covers(&other_conn));
let other_cat = SessionGrantKey {
connection: "c1".into(),
category: SqlCategory::Ddl,
tables: tables(&[("app", "jobs")]),
};
assert!(!e.session_covers(&other_cat));
}
#[test]
fn session_grants_revocable() {
let e = ApprovalEngine::new();
e.grant_session(SessionGrantKey {
connection: "c1".into(),
category: SqlCategory::Write,
tables: tables(&[("app", "jobs")]),
});
e.grant_session(SessionGrantKey {
connection: "c2".into(),
category: SqlCategory::Write,
tables: tables(&[("app", "jobs")]),
});
assert_eq!(
e.revoke_sessions(SessionRevokeFilter {
connection: Some("c1".into()),
category: None,
}),
1
);
assert_eq!(e.snapshot_sessions().len(), 1);
}
#[test]
fn grant_choice_from_legacy_strings() {
fn parse(s: &str) -> Result<GrantChoiceWrapper, serde_json::Error> {
serde_json::from_str(s)
}
assert!(matches!(parse(r#"{"c":"once"}"#), Ok(w) if w.c == GrantChoice::Once));
assert!(matches!(parse(r#"{"c":"session"}"#), Ok(w) if w.c == GrantChoice::Session));
assert!(matches!(parse(r#"{"c":"decline"}"#), Ok(w) if w.c == GrantChoice::Decline));
assert!(parse(r#"{"c":"other"}"#).is_err());
}
#[derive(serde::Deserialize, Debug)]
struct GrantChoiceWrapper {
c: GrantChoice,
}
}
impl serde::Serialize for GrantChoice {
fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
s.serialize_str(match self {
GrantChoice::Once => "once",
GrantChoice::Session => "session",
GrantChoice::Decline => "decline",
})
}
}
impl<'de> serde::Deserialize<'de> for GrantChoice {
fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
let s = String::deserialize(d)?;
match s.as_str() {
"once" => Ok(GrantChoice::Once),
"session" => Ok(GrantChoice::Session),
"decline" => Ok(GrantChoice::Decline),
_ => Err(serde::de::Error::custom("invalid grant choice")),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PolicyDecision {
pub category: SqlCategory,
pub action: PolicyAction,
pub confirmed: bool,
pub grant_used: Option<&'static str>, }