use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use cheers_core::{Scope, StoreError, UserId};
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct UserTokenRecord {
pub jti: String,
pub user_id: UserId,
pub name: String,
pub scopes: Vec<Scope>,
pub aud: String,
pub created_at: i64,
pub last_used_at: Option<i64>,
pub expires_at: i64,
pub revoked: bool,
}
impl UserTokenRecord {
#[allow(clippy::too_many_arguments)]
pub fn new(
jti: impl Into<String>,
user_id: UserId,
name: impl Into<String>,
scopes: Vec<Scope>,
aud: impl Into<String>,
created_at: i64,
expires_at: i64,
) -> Self {
Self {
jti: jti.into(),
user_id,
name: name.into(),
scopes,
aud: aud.into(),
created_at,
last_used_at: None,
expires_at,
revoked: false,
}
}
pub fn is_live_at(&self, now: i64) -> bool {
!self.revoked && self.expires_at > now
}
}
pub fn encode_scopes(scopes: &[Scope]) -> String {
scopes
.iter()
.map(|s| s.as_wire())
.collect::<Vec<_>>()
.join(" ")
}
pub fn decode_scopes(raw: &str) -> Result<Vec<Scope>, StoreError> {
raw.split_whitespace()
.map(|s| {
s.parse::<Scope>()
.map_err(|e| StoreError::Backend(format!("user_tokens.scopes: {e}")))
})
.collect()
}
#[async_trait]
pub trait UserTokenStore: Send + Sync {
async fn insert(&self, record: &UserTokenRecord) -> Result<(), StoreError>;
async fn list_live_for_user(
&self,
user_id: &UserId,
now: i64,
) -> Result<Vec<UserTokenRecord>, StoreError>;
async fn get(&self, jti: &str) -> Result<Option<UserTokenRecord>, StoreError>;
async fn mark_revoked(&self, jti: &str) -> Result<(), StoreError>;
async fn touch_last_used(&self, jti: &str, now: i64) -> Result<(), StoreError>;
}
#[derive(Default, Clone)]
pub struct MemoryUserTokenStore {
inner: Arc<Mutex<HashMap<String, UserTokenRecord>>>,
}
impl MemoryUserTokenStore {
pub fn new() -> Self {
Self::default()
}
pub fn len(&self) -> usize {
self.inner.lock().expect("user token store poisoned").len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[async_trait]
impl UserTokenStore for MemoryUserTokenStore {
async fn insert(&self, record: &UserTokenRecord) -> Result<(), StoreError> {
self.inner
.lock()
.expect("user token store poisoned")
.insert(record.jti.clone(), record.clone());
Ok(())
}
async fn list_live_for_user(
&self,
user_id: &UserId,
now: i64,
) -> Result<Vec<UserTokenRecord>, StoreError> {
Ok(self
.inner
.lock()
.expect("user token store poisoned")
.values()
.filter(|r| &r.user_id == user_id && r.is_live_at(now))
.cloned()
.collect())
}
async fn get(&self, jti: &str) -> Result<Option<UserTokenRecord>, StoreError> {
Ok(self
.inner
.lock()
.expect("user token store poisoned")
.get(jti)
.cloned())
}
async fn mark_revoked(&self, jti: &str) -> Result<(), StoreError> {
let mut guard = self.inner.lock().expect("user token store poisoned");
match guard.get_mut(jti) {
Some(row) => {
row.revoked = true;
Ok(())
}
None => Err(StoreError::NotFound),
}
}
async fn touch_last_used(&self, jti: &str, now: i64) -> Result<(), StoreError> {
let mut guard = self.inner.lock().expect("user token store poisoned");
match guard.get_mut(jti) {
Some(row) => {
row.last_used_at = Some(now);
Ok(())
}
None => Err(StoreError::NotFound),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use pollster::block_on;
fn rec(jti: &str, user: &str, expires_at: i64) -> UserTokenRecord {
UserTokenRecord::new(
jti,
UserId::new(user),
"ci",
vec![Scope::CloudRead],
"https://kamaji.example",
1_000,
expires_at,
)
}
#[test]
fn insert_then_list_returns_only_the_owner_s_live_rows() {
let store = MemoryUserTokenStore::new();
block_on(store.insert(&rec("j1", "alice", 9_000))).unwrap();
block_on(store.insert(&rec("j2", "bob", 9_000))).unwrap();
let alice = block_on(store.list_live_for_user(&UserId::new("alice"), 2_000)).unwrap();
assert_eq!(alice.len(), 1);
assert_eq!(alice[0].jti, "j1");
}
#[test]
fn expired_and_revoked_rows_drop_out_of_the_live_list_but_stay_gettable() {
let store = MemoryUserTokenStore::new();
block_on(store.insert(&rec("live", "alice", 9_000))).unwrap();
block_on(store.insert(&rec("expired", "alice", 1_500))).unwrap();
block_on(store.insert(&rec("revoked", "alice", 9_000))).unwrap();
block_on(store.mark_revoked("revoked")).unwrap();
let live = block_on(store.list_live_for_user(&UserId::new("alice"), 2_000)).unwrap();
assert_eq!(live.len(), 1);
assert_eq!(live[0].jti, "live");
assert!(block_on(store.get("revoked")).unwrap().unwrap().revoked);
assert!(block_on(store.get("expired")).unwrap().is_some());
}
#[test]
fn mark_revoked_is_idempotent_and_unknown_jti_is_not_found() {
let store = MemoryUserTokenStore::new();
block_on(store.insert(&rec("j1", "alice", 9_000))).unwrap();
block_on(store.mark_revoked("j1")).unwrap();
block_on(store.mark_revoked("j1")).unwrap();
match block_on(store.mark_revoked("nope")) {
Err(StoreError::NotFound) => {}
other => panic!("expected NotFound, got {other:?}"),
}
}
#[test]
fn scopes_round_trip_through_the_column_encoding() {
let scopes = vec![Scope::CloudRead, Scope::BoardWrite, Scope::CampAdmin];
let raw = encode_scopes(&scopes);
assert_eq!(raw, "cloud:read board:write camp:admin");
assert_eq!(decode_scopes(&raw).unwrap(), scopes);
assert_eq!(encode_scopes(&[]), "");
assert_eq!(decode_scopes("").unwrap(), Vec::<Scope>::new());
}
#[test]
fn decode_scopes_refuses_an_unknown_token() {
match decode_scopes("cloud:read cloud:teleport") {
Err(StoreError::Backend(msg)) => assert!(msg.contains("cloud:teleport"), "{msg}"),
other => panic!("expected Backend, got {other:?}"),
}
}
#[test]
fn touch_last_used_stamps_the_row() {
let store = MemoryUserTokenStore::new();
block_on(store.insert(&rec("j1", "alice", 9_000))).unwrap();
assert!(block_on(store.get("j1")).unwrap().unwrap().last_used_at.is_none());
block_on(store.touch_last_used("j1", 4_242)).unwrap();
assert_eq!(
block_on(store.get("j1")).unwrap().unwrap().last_used_at,
Some(4_242)
);
match block_on(store.touch_last_used("nope", 1)) {
Err(StoreError::NotFound) => {}
other => panic!("expected NotFound, got {other:?}"),
}
}
}