use std::collections::HashMap;
use std::fmt;
pub const SESSION_TOKEN_HALF_BYTES: usize = 16;
pub const SESSION_TOKEN_TOTAL_BYTES: usize = SESSION_TOKEN_HALF_BYTES * 2;
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct TokenHalf([u8; SESSION_TOKEN_HALF_BYTES]);
impl TokenHalf {
pub fn generate() -> Result<Self, SessionTokenError> {
let mut bytes = [0_u8; SESSION_TOKEN_HALF_BYTES];
getrandom::fill(&mut bytes)?;
Ok(Self(bytes))
}
pub fn from_bytes(bytes: [u8; SESSION_TOKEN_HALF_BYTES]) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8; SESSION_TOKEN_HALF_BYTES] {
&self.0
}
}
impl fmt::Debug for TokenHalf {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("TokenHalf(<redacted>)")
}
}
pub type DaemonId = String;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SessionTokenRejection {
MalformedLength,
BrokerHalfMismatch,
DaemonUnknown,
DaemonHalfMismatch,
}
#[derive(Debug)]
pub struct SessionTokenAuthority {
broker_token: TokenHalf,
daemon_tokens: HashMap<DaemonId, TokenHalf>,
}
impl SessionTokenAuthority {
pub fn new() -> Result<Self, SessionTokenError> {
Ok(Self {
broker_token: TokenHalf::generate()?,
daemon_tokens: HashMap::new(),
})
}
pub fn with_broker_token(broker_token: TokenHalf) -> Self {
Self {
broker_token,
daemon_tokens: HashMap::new(),
}
}
pub fn broker_token(&self) -> &TokenHalf {
&self.broker_token
}
pub fn rotate_broker_token(&mut self) -> Result<TokenHalf, SessionTokenError> {
let fresh = TokenHalf::generate()?;
self.broker_token = fresh;
Ok(fresh)
}
pub fn register_daemon(&mut self, daemon_id: DaemonId) -> Result<TokenHalf, SessionTokenError> {
let token = TokenHalf::generate()?;
self.daemon_tokens.insert(daemon_id, token);
Ok(token)
}
pub fn invalidate_daemon(&mut self, daemon_id: &str) -> bool {
self.daemon_tokens.remove(daemon_id).is_some()
}
pub fn daemon_count(&self) -> usize {
self.daemon_tokens.len()
}
pub fn composed_token_for(&self, daemon_id: &str) -> Option<Vec<u8>> {
let daemon_half = self.daemon_tokens.get(daemon_id)?;
Some(compose_presented_token(&self.broker_token, daemon_half))
}
pub fn validate(&self, presented: &[u8], daemon_id: &str) -> Result<(), SessionTokenRejection> {
if presented.len() != SESSION_TOKEN_TOTAL_BYTES {
return Err(SessionTokenRejection::MalformedLength);
}
let (broker_half, daemon_half) = presented.split_at(SESSION_TOKEN_HALF_BYTES);
if broker_half != self.broker_token.as_bytes() {
return Err(SessionTokenRejection::BrokerHalfMismatch);
}
let Some(expected_daemon_token) = self.daemon_tokens.get(daemon_id) else {
return Err(SessionTokenRejection::DaemonUnknown);
};
if daemon_half != expected_daemon_token.as_bytes() {
return Err(SessionTokenRejection::DaemonHalfMismatch);
}
Ok(())
}
}
pub fn compose_presented_token(broker_half: &TokenHalf, daemon_half: &TokenHalf) -> Vec<u8> {
let mut out = Vec::with_capacity(SESSION_TOKEN_TOTAL_BYTES);
out.extend_from_slice(broker_half.as_bytes());
out.extend_from_slice(daemon_half.as_bytes());
out
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum SessionTokenError {
#[error("session token random generation failed: {0}")]
Random(String),
}
impl From<getrandom::Error> for SessionTokenError {
fn from(value: getrandom::Error) -> Self {
Self::Random(value.to_string())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn half(byte: u8) -> TokenHalf {
TokenHalf::from_bytes([byte; SESSION_TOKEN_HALF_BYTES])
}
#[test]
fn generate_produces_distinct_halves() {
let a = TokenHalf::generate().expect("random");
let b = TokenHalf::generate().expect("random");
assert_ne!(a.as_bytes(), b.as_bytes(), "two mints must not collide");
}
#[test]
fn valid_composite_token_validates() {
let mut authority = SessionTokenAuthority::with_broker_token(half(0xAA));
let daemon_token = authority.register_daemon("daemon-1".into()).expect("mint");
let presented = compose_presented_token(authority.broker_token(), &daemon_token);
assert_eq!(authority.validate(&presented, "daemon-1"), Ok(()));
}
#[test]
fn wrong_broker_half_is_rejected_even_for_a_valid_daemon() {
let mut authority = SessionTokenAuthority::with_broker_token(half(0xAA));
let daemon_token = authority.register_daemon("daemon-1".into()).expect("mint");
let wrong_broker_half = half(0xFF);
let presented = compose_presented_token(&wrong_broker_half, &daemon_token);
assert_eq!(
authority.validate(&presented, "daemon-1"),
Err(SessionTokenRejection::BrokerHalfMismatch)
);
}
#[test]
fn broker_half_is_checked_before_daemon_lookup_for_an_unknown_daemon() {
let authority = SessionTokenAuthority::with_broker_token(half(0xAA));
let wrong_broker_half = half(0xFF);
let daemon_half = half(0x11);
let presented = compose_presented_token(&wrong_broker_half, &daemon_half);
assert_eq!(
authority.validate(&presented, "never-registered"),
Err(SessionTokenRejection::BrokerHalfMismatch)
);
}
#[test]
fn correct_broker_half_but_unregistered_daemon_is_rejected() {
let authority = SessionTokenAuthority::with_broker_token(half(0xAA));
let daemon_half = half(0x11);
let presented = compose_presented_token(&half(0xAA), &daemon_half);
assert_eq!(
authority.validate(&presented, "never-registered"),
Err(SessionTokenRejection::DaemonUnknown)
);
}
#[test]
fn wrong_daemon_half_is_rejected() {
let mut authority = SessionTokenAuthority::with_broker_token(half(0xAA));
authority.register_daemon("daemon-1".into()).expect("mint");
let wrong_daemon_half = half(0xEE);
let presented = compose_presented_token(&half(0xAA), &wrong_daemon_half);
assert_eq!(
authority.validate(&presented, "daemon-1"),
Err(SessionTokenRejection::DaemonHalfMismatch)
);
}
#[test]
fn malformed_length_is_rejected_before_any_comparison() {
let authority = SessionTokenAuthority::with_broker_token(half(0xAA));
assert_eq!(
authority.validate(&[0xAA; 5], "daemon-1"),
Err(SessionTokenRejection::MalformedLength)
);
assert_eq!(
authority.validate(&[], "daemon-1"),
Err(SessionTokenRejection::MalformedLength)
);
assert_eq!(
authority.validate(&[0xAA; SESSION_TOKEN_TOTAL_BYTES + 1], "daemon-1"),
Err(SessionTokenRejection::MalformedLength)
);
}
#[test]
fn broker_rotation_invalidates_every_daemons_sessions() {
let mut authority = SessionTokenAuthority::with_broker_token(half(0xAA));
let daemon_a = authority.register_daemon("daemon-a".into()).expect("mint");
let daemon_b = authority.register_daemon("daemon-b".into()).expect("mint");
let old_broker_half = *authority.broker_token();
let presented_a = compose_presented_token(&old_broker_half, &daemon_a);
let presented_b = compose_presented_token(&old_broker_half, &daemon_b);
assert_eq!(authority.validate(&presented_a, "daemon-a"), Ok(()));
assert_eq!(authority.validate(&presented_b, "daemon-b"), Ok(()));
authority.rotate_broker_token().expect("rotate");
assert_eq!(
authority.validate(&presented_a, "daemon-a"),
Err(SessionTokenRejection::BrokerHalfMismatch)
);
assert_eq!(
authority.validate(&presented_b, "daemon-b"),
Err(SessionTokenRejection::BrokerHalfMismatch)
);
}
#[test]
fn invalidating_one_daemon_does_not_disrupt_another() {
let mut authority = SessionTokenAuthority::with_broker_token(half(0xAA));
let daemon_a = authority.register_daemon("daemon-a".into()).expect("mint");
let daemon_b = authority.register_daemon("daemon-b".into()).expect("mint");
let broker_half = *authority.broker_token();
let presented_a = compose_presented_token(&broker_half, &daemon_a);
let presented_b = compose_presented_token(&broker_half, &daemon_b);
assert!(authority.invalidate_daemon("daemon-a"));
assert_eq!(
authority.validate(&presented_a, "daemon-a"),
Err(SessionTokenRejection::DaemonUnknown),
"daemon-a's session must be gone"
);
assert_eq!(
authority.validate(&presented_b, "daemon-b"),
Ok(()),
"daemon-b's session must be completely unaffected"
);
}
#[test]
fn invalidate_daemon_reports_whether_a_token_was_present() {
let mut authority = SessionTokenAuthority::with_broker_token(half(0xAA));
assert!(!authority.invalidate_daemon("never-registered"));
authority.register_daemon("daemon-1".into()).expect("mint");
assert!(authority.invalidate_daemon("daemon-1"));
assert!(!authority.invalidate_daemon("daemon-1"));
}
#[test]
fn re_registering_a_daemon_mints_a_new_token_and_invalidates_the_old_one() {
let mut authority = SessionTokenAuthority::with_broker_token(half(0xAA));
let broker_half = *authority.broker_token();
let first_token = authority.register_daemon("daemon-1".into()).expect("mint");
let stale_presented = compose_presented_token(&broker_half, &first_token);
let second_token = authority.register_daemon("daemon-1".into()).expect("mint");
assert_ne!(first_token.as_bytes(), second_token.as_bytes());
assert_eq!(
authority.validate(&stale_presented, "daemon-1"),
Err(SessionTokenRejection::DaemonHalfMismatch),
"the pre-restart token must no longer validate"
);
}
#[test]
fn daemon_count_tracks_registration_and_invalidation() {
let mut authority = SessionTokenAuthority::with_broker_token(half(0xAA));
assert_eq!(authority.daemon_count(), 0);
authority.register_daemon("daemon-1".into()).expect("mint");
authority.register_daemon("daemon-2".into()).expect("mint");
assert_eq!(authority.daemon_count(), 2);
authority.invalidate_daemon("daemon-1");
assert_eq!(authority.daemon_count(), 1);
}
#[test]
fn debug_impl_redacts_token_bytes() {
let token = TokenHalf::from_bytes([0x42; SESSION_TOKEN_HALF_BYTES]);
let rendered = format!("{token:?}");
assert!(
!rendered.contains("42"),
"token bytes must not leak into Debug output: {rendered}"
);
assert!(rendered.contains("redacted"));
}
}