use chrono::{DateTime, Utc};
use std::collections::{HashMap, HashSet};
use std::error::Error;
use std::fmt::{Display, Formatter};
use std::sync::RwLock;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TrustStateError(String);
impl TrustStateError {
pub fn new(message: impl Into<String>) -> Self {
Self(message.into())
}
}
impl Display for TrustStateError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
impl Error for TrustStateError {}
pub trait TrustStateStore: Send + Sync {
fn is_token_revoked(&self, issuer: &str, token_id: &str) -> Result<bool, TrustStateError>;
fn is_agent_denied(&self, issuer: &str, agent_id: &str) -> Result<bool, TrustStateError>;
fn consume_nonce(
&self,
issuer: &str,
nonce: &str,
observed_at: DateTime<Utc>,
valid_until: DateTime<Utc>,
) -> Result<bool, TrustStateError>;
}
pub trait TrustStateAdmin: Send + Sync {
fn revoke_token(&self, issuer: &str, token_id: &str) -> Result<(), TrustStateError>;
fn deny_agent(&self, issuer: &str, agent_id: &str) -> Result<(), TrustStateError>;
}
#[derive(Default)]
pub struct InMemoryTrustState {
inner: RwLock<State>,
}
#[derive(Default)]
struct State {
revoked_tokens: HashSet<(String, String)>,
denied_agents: HashSet<(String, String)>,
nonces: HashMap<(String, String), DateTime<Utc>>,
}
impl InMemoryTrustState {
pub fn new() -> Self {
Self::default()
}
}
impl TrustStateStore for InMemoryTrustState {
fn is_token_revoked(&self, issuer: &str, token_id: &str) -> Result<bool, TrustStateError> {
Ok(self
.inner
.read()
.map_err(lock_error)?
.revoked_tokens
.contains(&(issuer.to_owned(), token_id.to_owned())))
}
fn is_agent_denied(&self, issuer: &str, agent_id: &str) -> Result<bool, TrustStateError> {
Ok(self
.inner
.read()
.map_err(lock_error)?
.denied_agents
.contains(&(issuer.to_owned(), agent_id.to_owned())))
}
fn consume_nonce(
&self,
issuer: &str,
nonce: &str,
observed_at: DateTime<Utc>,
valid_until: DateTime<Utc>,
) -> Result<bool, TrustStateError> {
let mut state = self.inner.write().map_err(lock_error)?;
state.nonces.retain(|_, expiry| *expiry > observed_at);
let key = (issuer.to_owned(), nonce.to_owned());
if state.nonces.contains_key(&key) {
return Ok(false);
}
state.nonces.insert(key, valid_until);
Ok(true)
}
}
impl TrustStateAdmin for InMemoryTrustState {
fn revoke_token(&self, issuer: &str, token_id: &str) -> Result<(), TrustStateError> {
self.inner
.write()
.map_err(lock_error)?
.revoked_tokens
.insert((issuer.to_owned(), token_id.to_owned()));
Ok(())
}
fn deny_agent(&self, issuer: &str, agent_id: &str) -> Result<(), TrustStateError> {
self.inner
.write()
.map_err(lock_error)?
.denied_agents
.insert((issuer.to_owned(), agent_id.to_owned()));
Ok(())
}
}
fn lock_error<T>(_: std::sync::PoisonError<T>) -> TrustStateError {
TrustStateError::new("trust state lock poisoned")
}