use crate::pb::session_service_client::SessionServiceClient;
use crate::pb::{resolve_response, ResolveRequest};
use crate::UserId;
use chrono::{DateTime, Utc};
use futures::future::{BoxFuture, FutureExt, Shared};
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, HashMap};
use std::future::Future;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tonic::transport::Channel;
const RESOLVE_TIMEOUT: Duration = Duration::from_secs(5);
pub fn token_hash(raw_token: &str) -> String {
let digest = Sha256::digest(raw_token.as_bytes());
hex::encode(digest)
}
pub struct TokenHash(pub String);
impl TokenHash {
pub fn from_raw(raw_token: &str) -> Self {
TokenHash(token_hash(raw_token))
}
pub fn as_str(&self) -> &str {
&self.0
}
}
#[derive(Clone, Debug)]
pub struct ResolvedSession {
pub principal: UserId,
pub tenant_id: String,
pub expires_at: DateTime<Utc>,
}
#[derive(Clone, Debug, thiserror::Error)]
pub enum ResolveError {
#[error("session resolve transport error: {0}")]
Transport(String),
#[error("session resolve backend error: {0}")]
Backend(String),
}
impl ResolveError {
fn is_transport(&self) -> bool {
matches!(self, ResolveError::Transport(_))
}
}
pub type ResolveFuture<'a> = BoxFuture<'a, Result<Option<ResolvedSession>, ResolveError>>;
pub trait SessionFetcher: Send + Sync + 'static {
fn fetch<'a>(&'a self, token_hash: &'a str) -> ResolveFuture<'a>;
}
pub trait SessionResolver: Send + Sync + 'static {
fn resolve<'a>(&'a self, token_hash: &'a str) -> ResolveFuture<'a>;
fn evict(&self, token_hash: &str);
}
#[derive(Clone, Debug)]
pub struct ResolverConfig {
pub capacity: usize,
pub l1_ttl: Duration,
pub neg_ttl: Duration,
pub stale_if_error: Duration,
}
impl Default for ResolverConfig {
fn default() -> Self {
ResolverConfig {
capacity: 10_000,
l1_ttl: Duration::from_secs(30),
neg_ttl: Duration::from_secs(2),
stale_if_error: Duration::ZERO,
}
}
}
#[derive(Clone)]
struct Entry {
outcome: Option<ResolvedSession>,
fetched_at: Instant,
fresh_until: Instant,
stale_until: Instant,
effective_ttl: Duration,
}
struct Lru {
map: HashMap<String, (Entry, u64)>,
recency: BTreeMap<u64, String>,
next_gen: u64,
capacity: usize,
}
impl Lru {
fn new(capacity: usize) -> Self {
Lru {
map: HashMap::new(),
recency: BTreeMap::new(),
next_gen: 0,
capacity,
}
}
fn get(&mut self, key: &str) -> Option<Entry> {
if self.capacity == 0 {
return None;
}
let (entry, old_gen) = self.map.get(key)?;
let entry = entry.clone();
let old_gen = *old_gen;
self.recency.remove(&old_gen);
let gen = self.next_gen;
self.next_gen += 1;
self.recency.insert(gen, key.to_string());
if let Some(e) = self.map.get_mut(key) {
e.1 = gen;
}
Some(entry)
}
fn peek(&self, key: &str) -> Option<Entry> {
self.map.get(key).map(|(e, _)| e.clone())
}
fn put(&mut self, key: String, entry: Entry) {
if self.capacity == 0 {
return;
}
if let Some((_, old_gen)) = self.map.get(&key) {
let old_gen = *old_gen;
self.recency.remove(&old_gen);
}
let gen = self.next_gen;
self.next_gen += 1;
self.recency.insert(gen, key.clone());
self.map.insert(key, (entry, gen));
while self.map.len() > self.capacity {
let Some((&lru_gen, lru_key)) = self.recency.iter().next() else {
break;
};
let lru_key = lru_key.clone();
self.recency.remove(&lru_gen);
self.map.remove(&lru_key);
}
}
fn remove(&mut self, key: &str) {
if let Some((_, gen)) = self.map.remove(key) {
self.recency.remove(&gen);
}
}
}
type FillResult = Result<Option<ResolvedSession>, Arc<ResolveError>>;
type SharedFill = Shared<BoxFuture<'static, FillResult>>;
#[derive(Clone, Default)]
struct SingleFlight {
inflight: Arc<Mutex<HashMap<String, (u64, SharedFill)>>>,
next_id: Arc<AtomicU64>,
}
impl SingleFlight {
async fn run<F, Fut>(&self, key: &str, make: F) -> Result<Option<ResolvedSession>, ResolveError>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<Option<ResolvedSession>, ResolveError>> + Send + 'static,
{
let shared = {
let mut map = self.inflight.lock().expect("singleflight mutex poisoned");
if let Some((_, existing)) = map.get(key) {
existing.clone()
} else {
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
let inflight = Arc::clone(&self.inflight);
let owned_key = key.to_string();
let inner = make().map(|r| r.map_err(Arc::new));
let fut: SharedFill = async move {
let result = inner.await;
let mut map = inflight.lock().expect("singleflight mutex poisoned");
if map
.get(&owned_key)
.map(|(gid, _)| *gid == id)
.unwrap_or(false)
{
map.remove(&owned_key);
}
result
}
.boxed()
.shared();
map.insert(key.to_string(), (id, fut.clone()));
tokio::spawn(fut.clone());
fut
}
};
shared.await.map_err(|e| (*e).clone())
}
}
struct ResolverInner {
fetcher: Arc<dyn SessionFetcher>,
cache: Mutex<Lru>,
flight: SingleFlight,
cfg: ResolverConfig,
}
impl ResolverInner {
fn effective_ttl(&self) -> Duration {
let jitter = 0.8 + 0.2 * rand::random::<f64>();
Duration::from_secs_f64(self.cfg.l1_ttl.as_secs_f64() * jitter)
}
async fn resolve(
self: Arc<Self>,
hash: String,
) -> Result<Option<ResolvedSession>, ResolveError> {
let now = Instant::now();
let now_wall = Utc::now();
let cached = {
let mut guard = self.cache.lock().expect("session cache mutex poisoned");
guard.get(&hash)
};
if let Some(entry) = cached {
let wall_valid = match &entry.outcome {
Some(s) => s.expires_at > now_wall,
None => true,
};
if !wall_valid {
self.cache
.lock()
.expect("session cache mutex poisoned")
.remove(&hash);
} else if now < entry.fresh_until {
if entry.outcome.is_some() {
let remaining = entry.fresh_until.saturating_duration_since(now);
if remaining.as_secs_f64() < 0.10 * entry.effective_ttl.as_secs_f64() {
self.clone().spawn_refresh(hash.clone());
}
}
return Ok(entry.outcome);
}
}
let stale = self.stale_candidate(&hash, now, now_wall);
let this = self.clone();
let key = hash.clone();
let outcome = self
.flight
.run(&hash, move || {
let this = this.clone();
async move { this.fill(key).await }
})
.await;
match outcome {
Ok(v) => Ok(v),
Err(e) => {
if e.is_transport() {
if let Some((s, _fetched_at)) = stale {
log::warn!("session resolver: serving stale entry on transport error: {e}");
return Ok(Some(s));
}
}
Err(e)
}
}
}
async fn fill(self: Arc<Self>, hash: String) -> Result<Option<ResolvedSession>, ResolveError> {
let fetched = tokio::time::timeout(RESOLVE_TIMEOUT, self.fetcher.fetch(&hash))
.await
.map_err(|_| {
ResolveError::Transport("session resolve timed out after 5s".to_string())
})??;
let now = Instant::now();
let entry = match &fetched {
Some(s) => {
let eff = self.effective_ttl();
let wall_remaining = (s.expires_at - Utc::now())
.to_std()
.unwrap_or(Duration::ZERO);
let ttl = eff.min(wall_remaining);
Entry {
outcome: Some(s.clone()),
fetched_at: now,
fresh_until: now + ttl,
stale_until: now + ttl + self.cfg.stale_if_error,
effective_ttl: eff,
}
}
None => Entry {
outcome: None,
fetched_at: now,
fresh_until: now + self.cfg.neg_ttl,
stale_until: now + self.cfg.neg_ttl,
effective_ttl: self.cfg.neg_ttl,
},
};
{
let mut guard = self.cache.lock().expect("session cache mutex poisoned");
guard.put(hash, entry);
}
Ok(fetched)
}
fn stale_candidate(
&self,
hash: &str,
now: Instant,
now_wall: DateTime<Utc>,
) -> Option<(ResolvedSession, Instant)> {
if self.cfg.stale_if_error.is_zero() {
return None;
}
let guard = self.cache.lock().expect("session cache mutex poisoned");
let entry = guard.peek(hash)?;
match entry.outcome {
Some(s) if s.expires_at > now_wall && now < entry.stale_until => {
Some((s, entry.fetched_at))
}
_ => None,
}
}
fn spawn_refresh(self: Arc<Self>, hash: String) {
tokio::spawn(async move {
let this = self.clone();
let key = hash.clone();
let _ = self
.flight
.run(&hash, move || {
let this = this.clone();
async move { this.fill(key).await }
})
.await;
});
}
}
#[derive(Clone)]
pub struct CachedResolver {
shared: Arc<ResolverInner>,
}
impl CachedResolver {
pub fn new(fetcher: Arc<dyn SessionFetcher>, cfg: ResolverConfig) -> Self {
CachedResolver {
shared: Arc::new(ResolverInner {
cache: Mutex::new(Lru::new(cfg.capacity)),
flight: SingleFlight::default(),
fetcher,
cfg,
}),
}
}
pub fn into_dyn(self) -> Arc<dyn SessionResolver> {
Arc::new(self)
}
}
impl SessionResolver for CachedResolver {
fn resolve<'a>(&'a self, token_hash: &'a str) -> ResolveFuture<'a> {
let shared = self.shared.clone();
let hash = token_hash.to_string();
Box::pin(async move { shared.resolve(hash).await })
}
fn evict(&self, token_hash: &str) {
let mut guard = self
.shared
.cache
.lock()
.expect("session cache mutex poisoned");
guard.remove(token_hash);
}
}
pub struct GrpcSessionResolver;
impl GrpcSessionResolver {
#[allow(clippy::new_ret_no_self)]
pub fn new(channel: Channel, cfg: ResolverConfig) -> Arc<dyn SessionResolver> {
let client = SessionServiceClient::new(channel);
let fetcher: Arc<dyn SessionFetcher> = Arc::new(GrpcFetcher { client });
CachedResolver::new(fetcher, cfg).into_dyn()
}
}
struct GrpcFetcher {
client: SessionServiceClient<Channel>,
}
impl SessionFetcher for GrpcFetcher {
fn fetch<'a>(&'a self, token_hash: &'a str) -> ResolveFuture<'a> {
let mut client = self.client.clone();
let token_hash = token_hash.to_string();
Box::pin(async move {
let resp = client
.resolve(ResolveRequest { token_hash })
.await
.map_err(classify_status)?;
match resp.into_inner().outcome {
Some(resolve_response::Outcome::Session(s)) => Ok(Some(ResolvedSession {
principal: UserId::try_from(s.principal)
.map_err(|e| ResolveError::Backend(e.to_string()))?,
tenant_id: s.tenant_id,
expires_at: DateTime::from_timestamp(s.expires_at_unix_seconds, 0)
.unwrap_or_else(Utc::now),
})),
Some(resolve_response::Outcome::NotFound(_)) | None => Ok(None),
}
})
}
}
fn classify_status(status: tonic::Status) -> ResolveError {
match status.code() {
tonic::Code::Unavailable | tonic::Code::DeadlineExceeded => {
ResolveError::Transport(status.message().to_string())
}
_ => ResolveError::Backend(status.message().to_string()),
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::AtomicUsize;
struct CountingFetcher {
calls: Arc<AtomicUsize>,
session: Option<ResolvedSession>,
}
impl SessionFetcher for CountingFetcher {
fn fetch<'a>(&'a self, _token_hash: &'a str) -> ResolveFuture<'a> {
self.calls.fetch_add(1, Ordering::Relaxed);
let session = self.session.clone();
Box::pin(async move { Ok(session) })
}
}
struct FailingFetcher {
calls: Arc<AtomicUsize>,
}
impl SessionFetcher for FailingFetcher {
fn fetch<'a>(&'a self, _token_hash: &'a str) -> ResolveFuture<'a> {
self.calls.fetch_add(1, Ordering::Relaxed);
Box::pin(async move { Err(ResolveError::Transport("down".into())) })
}
}
fn session_valid_for(mins: i64) -> ResolvedSession {
ResolvedSession {
principal: UserId::try_from(42).unwrap(),
tenant_id: String::new(),
expires_at: Utc::now() + chrono::TimeDelta::minutes(mins),
}
}
fn cfg() -> ResolverConfig {
ResolverConfig {
capacity: 100,
l1_ttl: Duration::from_secs(30),
neg_ttl: Duration::from_secs(2),
stale_if_error: Duration::ZERO,
}
}
#[test]
fn token_hash_is_hex_sha256() {
assert_eq!(
token_hash(""),
"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
);
assert_eq!(token_hash("abc").len(), 64);
}
#[test]
fn default_config_matches_243() {
let cfg = ResolverConfig::default();
assert_eq!(cfg.capacity, 10_000);
assert_eq!(cfg.l1_ttl, Duration::from_secs(30));
assert_eq!(cfg.neg_ttl, Duration::from_secs(2));
assert_eq!(cfg.stale_if_error, Duration::ZERO);
}
#[tokio::test]
async fn hit_does_not_refetch() {
let calls = Arc::new(AtomicUsize::new(0));
let fetcher = Arc::new(CountingFetcher {
calls: calls.clone(),
session: Some(session_valid_for(120)),
});
let r = CachedResolver::new(fetcher, cfg());
let first = r.resolve("deadbeef").await.unwrap();
assert!(first.is_some());
let _ = r.resolve("deadbeef").await.unwrap();
let _ = r.resolve("deadbeef").await.unwrap();
assert_eq!(calls.load(Ordering::Relaxed), 1, "L1 hit must not refetch");
}
#[tokio::test]
async fn unknown_token_is_tombstoned() {
let calls = Arc::new(AtomicUsize::new(0));
let fetcher = Arc::new(CountingFetcher {
calls: calls.clone(),
session: None,
});
let r = CachedResolver::new(fetcher, cfg());
assert!(r.resolve("nope").await.unwrap().is_none());
assert!(r.resolve("nope").await.unwrap().is_none());
assert_eq!(
calls.load(Ordering::Relaxed),
1,
"tombstone must not refetch"
);
}
#[tokio::test]
async fn evict_forces_refetch() {
let calls = Arc::new(AtomicUsize::new(0));
let fetcher = Arc::new(CountingFetcher {
calls: calls.clone(),
session: Some(session_valid_for(120)),
});
let r = CachedResolver::new(fetcher, cfg());
let _ = r.resolve("k").await.unwrap();
r.evict("k");
let _ = r.resolve("k").await.unwrap();
assert_eq!(
calls.load(Ordering::Relaxed),
2,
"evict must force a refetch"
);
}
#[tokio::test]
async fn expired_session_is_forced_miss() {
let calls = Arc::new(AtomicUsize::new(0));
let fetcher = Arc::new(CountingFetcher {
calls: calls.clone(),
session: Some(session_valid_for(-1)),
});
let r = CachedResolver::new(fetcher, cfg());
let _ = r.resolve("k").await.unwrap();
let _ = r.resolve("k").await.unwrap();
assert_eq!(
calls.load(Ordering::Relaxed),
2,
"expired entry must not serve from L1"
);
}
#[tokio::test]
async fn transport_error_propagates_without_stale_window() {
let calls = Arc::new(AtomicUsize::new(0));
let fetcher = Arc::new(FailingFetcher {
calls: calls.clone(),
});
let r = CachedResolver::new(fetcher, cfg());
let err = r.resolve("k").await.expect_err("must fail");
assert!(matches!(err, ResolveError::Transport(_)));
}
#[tokio::test]
async fn zero_capacity_disables_cache() {
let calls = Arc::new(AtomicUsize::new(0));
let fetcher = Arc::new(CountingFetcher {
calls: calls.clone(),
session: Some(session_valid_for(120)),
});
let mut c = cfg();
c.capacity = 0;
let r = CachedResolver::new(fetcher, c);
let _ = r.resolve("k").await.unwrap();
let _ = r.resolve("k").await.unwrap();
assert_eq!(
calls.load(Ordering::Relaxed),
2,
"capacity 0 must not cache"
);
}
#[tokio::test]
async fn lru_evicts_oldest() {
let calls = Arc::new(AtomicUsize::new(0));
let fetcher = Arc::new(CountingFetcher {
calls: calls.clone(),
session: Some(session_valid_for(120)),
});
let mut c = cfg();
c.capacity = 2;
let r = CachedResolver::new(fetcher, c);
let _ = r.resolve("a").await.unwrap();
let _ = r.resolve("b").await.unwrap();
let _ = r.resolve("c").await.unwrap(); let _ = r.resolve("a").await.unwrap(); assert_eq!(calls.load(Ordering::Relaxed), 4, "lru must evict oldest");
}
struct GatedFirstFetcher {
calls: Arc<AtomicUsize>,
gate: Arc<tokio::sync::Semaphore>,
session: Option<ResolvedSession>,
}
impl SessionFetcher for GatedFirstFetcher {
fn fetch<'a>(&'a self, _token_hash: &'a str) -> ResolveFuture<'a> {
let n = self.calls.fetch_add(1, Ordering::Relaxed);
let gate = self.gate.clone();
let session = self.session.clone();
Box::pin(async move {
if n == 0 {
let _permit = gate.acquire_owned().await.expect("gate closed");
}
Ok(session)
})
}
}
#[tokio::test]
async fn cancelled_leader_does_not_poison_future_resolves() {
let calls = Arc::new(AtomicUsize::new(0));
let gate = Arc::new(tokio::sync::Semaphore::new(0));
let fetcher = Arc::new(GatedFirstFetcher {
calls: calls.clone(),
gate: gate.clone(),
session: Some(session_valid_for(120)),
});
let r = CachedResolver::new(fetcher, cfg());
let leader = {
let r = r.clone();
tokio::spawn(async move { r.resolve("k").await })
};
while calls.load(Ordering::Relaxed) == 0 {
tokio::task::yield_now().await;
}
leader.abort();
let _ = leader.await;
gate.add_permits(1);
let got = r.resolve("k").await.unwrap();
assert!(got.is_some());
assert_eq!(calls.load(Ordering::Relaxed), 1);
r.evict("k");
let got = r.resolve("k").await.unwrap();
assert!(got.is_some());
assert_eq!(
calls.load(Ordering::Relaxed),
2,
"resolve after evict must refetch even after a cancelled leader"
);
}
struct HangingFetcher;
impl SessionFetcher for HangingFetcher {
fn fetch<'a>(&'a self, _token_hash: &'a str) -> ResolveFuture<'a> {
Box::pin(futures::future::pending())
}
}
#[tokio::test(start_paused = true)]
async fn fill_timeout_is_transport_error() {
let r = CachedResolver::new(Arc::new(HangingFetcher), cfg());
let err = r
.resolve("k")
.await
.expect_err("hung backend must time out");
assert!(matches!(err, ResolveError::Transport(_)));
}
struct SwitchableFetcher {
session: ResolvedSession,
err: Mutex<Option<ResolveError>>,
}
impl SessionFetcher for SwitchableFetcher {
fn fetch<'a>(&'a self, _token_hash: &'a str) -> ResolveFuture<'a> {
let err = self.err.lock().unwrap().clone();
let session = self.session.clone();
Box::pin(async move {
match err {
Some(e) => Err(e),
None => Ok(Some(session)),
}
})
}
}
fn stale_cfg() -> ResolverConfig {
ResolverConfig {
capacity: 100,
l1_ttl: Duration::from_millis(1),
neg_ttl: Duration::from_secs(2),
stale_if_error: Duration::from_secs(60),
}
}
#[tokio::test]
async fn transport_error_serves_stale_within_window() {
let fetcher = Arc::new(SwitchableFetcher {
session: session_valid_for(120),
err: Mutex::new(None),
});
let r = CachedResolver::new(fetcher.clone(), stale_cfg());
assert!(r.resolve("k").await.unwrap().is_some());
tokio::time::sleep(Duration::from_millis(20)).await; *fetcher.err.lock().unwrap() = Some(ResolveError::Transport("unavailable".into()));
let got = r
.resolve("k")
.await
.expect("transport error must serve stale");
assert!(got.is_some());
}
#[tokio::test]
async fn backend_error_propagates_despite_stale_window() {
let fetcher = Arc::new(SwitchableFetcher {
session: session_valid_for(120),
err: Mutex::new(None),
});
let r = CachedResolver::new(fetcher.clone(), stale_cfg());
assert!(r.resolve("k").await.unwrap().is_some());
tokio::time::sleep(Duration::from_millis(20)).await;
*fetcher.err.lock().unwrap() = Some(ResolveError::Backend("boom".into()));
r.resolve("k")
.await
.expect_err("backend error must propagate, not serve stale");
}
#[test]
fn classify_status_transport_vs_backend() {
assert!(matches!(
classify_status(tonic::Status::unavailable("x")),
ResolveError::Transport(_)
));
assert!(matches!(
classify_status(tonic::Status::deadline_exceeded("x")),
ResolveError::Transport(_)
));
assert!(matches!(
classify_status(tonic::Status::internal("x")),
ResolveError::Backend(_)
));
assert!(matches!(
classify_status(tonic::Status::permission_denied("x")),
ResolveError::Backend(_)
));
}
#[test]
fn token_hash_newtype_matches_fn() {
assert_eq!(TokenHash::from_raw("tok").as_str(), token_hash("tok"));
}
}