use std::collections::{HashMap, HashSet};
use std::fmt;
use tensor_wasm_core::types::TenantId;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TenantScope {
All,
Set(HashSet<TenantId>),
}
impl TenantScope {
pub fn allows(&self, tenant: TenantId) -> bool {
match self {
TenantScope::All => true,
TenantScope::Set(s) => s.contains(&tenant),
}
}
pub fn is_all(&self) -> bool {
matches!(self, TenantScope::All)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TokenScope {
pub tenants: TenantScope,
}
impl TokenScope {
pub fn all() -> Self {
Self {
tenants: TenantScope::All,
}
}
pub fn from_tenants<I>(iter: I) -> Self
where
I: IntoIterator<Item = TenantId>,
{
Self {
tenants: TenantScope::Set(iter.into_iter().collect()),
}
}
pub fn allows(&self, tenant: TenantId) -> bool {
self.tenants.allows(tenant)
}
}
pub type BearerString = String;
#[derive(Debug, Default, Clone)]
pub struct ParsedTokens {
pub token_scopes: HashMap<BearerString, TokenScope>,
pub deprecated_count: usize,
}
#[derive(Debug, PartialEq, Eq)]
pub enum ScopeParseError {
EmptyEntry,
MissingBearer,
EmptyTenantList,
InvalidTenant(String),
WildcardMixedWithIds,
UnknownKey(String),
}
impl fmt::Display for ScopeParseError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
ScopeParseError::EmptyEntry => f.write_str("empty token entry"),
ScopeParseError::MissingBearer => {
f.write_str("token entry has scope clause but no bearer token")
}
ScopeParseError::EmptyTenantList => f.write_str("tenant= clause has no value"),
ScopeParseError::InvalidTenant(s) => {
write!(f, "tenant id {s:?} is neither '*' nor a u64")
}
ScopeParseError::WildcardMixedWithIds => {
f.write_str("tenant=* cannot be mixed with explicit tenant ids")
}
ScopeParseError::UnknownKey(k) => {
write!(f, "unknown scope key {k:?} (only 'tenant=' is supported)")
}
}
}
}
impl std::error::Error for ScopeParseError {}
pub fn parse_token_entry(s: &str) -> Result<(BearerString, TokenScope), ScopeParseError> {
let trimmed = s.trim();
if trimmed.is_empty() {
return Err(ScopeParseError::EmptyEntry);
}
match trimmed.split_once(':') {
None => {
Ok((trimmed.to_owned(), TokenScope::all()))
}
Some((bearer, rest)) => {
let bearer = bearer.trim();
if bearer.is_empty() {
return Err(ScopeParseError::MissingBearer);
}
let rest = rest.trim();
let value = match rest.split_once('=') {
Some(("tenant", v)) => v.trim(),
Some((other, _)) => {
return Err(ScopeParseError::UnknownKey(other.to_owned()));
}
None => {
return Err(ScopeParseError::UnknownKey(rest.to_owned()));
}
};
if value.is_empty() {
return Err(ScopeParseError::EmptyTenantList);
}
let mut saw_wildcard = false;
let mut ids: HashSet<TenantId> = HashSet::new();
for raw in value.split(',') {
let frag = raw.trim();
if frag.is_empty() {
return Err(ScopeParseError::InvalidTenant(frag.to_owned()));
}
if frag == "*" {
saw_wildcard = true;
} else {
match frag.parse::<u64>() {
Ok(v) => {
ids.insert(TenantId(v));
}
Err(_) => return Err(ScopeParseError::InvalidTenant(frag.to_owned())),
}
}
}
if saw_wildcard && !ids.is_empty() {
return Err(ScopeParseError::WildcardMixedWithIds);
}
let scope = if saw_wildcard {
TokenScope::all()
} else {
TokenScope {
tenants: TenantScope::Set(ids),
}
};
Ok((bearer.to_owned(), scope))
}
}
}
pub fn parse_tokens_env(env_value: &str) -> ParsedTokens {
let mut out = ParsedTokens::default();
if env_value.trim().is_empty() {
return out;
}
let raw_fragments: Vec<&str> = env_value.split(',').collect();
let mut entries: Vec<String> = Vec::with_capacity(raw_fragments.len());
let mut current = String::new();
let mut in_tenant_list = false;
for frag in raw_fragments {
let trimmed = frag.trim();
if trimmed.is_empty() {
if !current.is_empty() {
entries.push(std::mem::take(&mut current));
}
in_tenant_list = false;
continue;
}
let has_colon = trimmed.contains(':');
let looks_like_tenant_continuation = trimmed == "*" || trimmed.parse::<u64>().is_ok();
if in_tenant_list && !has_colon && looks_like_tenant_continuation {
current.push(',');
current.push_str(trimmed);
} else {
if !current.is_empty() {
entries.push(std::mem::take(&mut current));
}
current.push_str(trimmed);
in_tenant_list = has_colon
&& trimmed
.split_once(':')
.is_some_and(|(_, rest)| rest.trim().starts_with("tenant="));
}
}
if !current.is_empty() {
entries.push(current);
}
for entry in entries {
let is_bare = !entry.contains(':');
match parse_token_entry(&entry) {
Ok((bearer, scope)) => {
if is_bare {
out.deprecated_count += 1;
}
out.token_scopes.insert(bearer, scope);
}
Err(e) => {
tracing::warn!(
target: "tensor_wasm_api::token_scope",
error = %e,
entry = %entry,
"ignoring malformed TENSOR_WASM_API_TOKENS entry",
);
}
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
fn ids(it: impl IntoIterator<Item = u64>) -> HashSet<TenantId> {
it.into_iter().map(TenantId).collect()
}
#[test]
fn parse_bare_entry_yields_wildcard() {
let (bearer, scope) = parse_token_entry("alpha").expect("parses");
assert_eq!(bearer, "alpha");
assert!(scope.tenants.is_all());
}
#[test]
fn parse_bare_entry_trims_whitespace() {
let (bearer, _) = parse_token_entry(" alpha ").expect("parses");
assert_eq!(bearer, "alpha");
}
#[test]
fn parse_wildcard_entry_yields_wildcard() {
let (bearer, scope) = parse_token_entry("alpha:tenant=*").expect("parses");
assert_eq!(bearer, "alpha");
assert!(scope.tenants.is_all());
}
#[test]
fn parse_single_tenant_entry() {
let (bearer, scope) = parse_token_entry("alpha:tenant=7").expect("parses");
assert_eq!(bearer, "alpha");
assert_eq!(scope.tenants, TenantScope::Set(ids([7])));
}
#[test]
fn parse_multi_tenant_entry() {
let (bearer, scope) = parse_token_entry("alpha:tenant=1,2,3").expect("parses");
assert_eq!(bearer, "alpha");
assert_eq!(scope.tenants, TenantScope::Set(ids([1, 2, 3])));
assert!(scope.allows(TenantId(2)));
assert!(!scope.allows(TenantId(4)));
}
#[test]
fn parse_multi_tenant_entry_with_whitespace() {
let (bearer, scope) = parse_token_entry("alpha:tenant= 1 , 2 , 3 ").expect("parses");
assert_eq!(bearer, "alpha");
assert_eq!(scope.tenants, TenantScope::Set(ids([1, 2, 3])));
}
#[test]
fn parse_empty_entry_errors() {
assert_eq!(
parse_token_entry("").unwrap_err(),
ScopeParseError::EmptyEntry
);
assert_eq!(
parse_token_entry(" ").unwrap_err(),
ScopeParseError::EmptyEntry
);
}
#[test]
fn parse_missing_bearer_errors() {
assert_eq!(
parse_token_entry(":tenant=1").unwrap_err(),
ScopeParseError::MissingBearer
);
}
#[test]
fn parse_empty_tenant_list_errors() {
assert_eq!(
parse_token_entry("alpha:tenant=").unwrap_err(),
ScopeParseError::EmptyTenantList
);
}
#[test]
fn parse_invalid_tenant_errors() {
let err = parse_token_entry("alpha:tenant=not-a-number").unwrap_err();
assert_eq!(err, ScopeParseError::InvalidTenant("not-a-number".into()));
}
#[test]
fn parse_unknown_key_errors() {
let err = parse_token_entry("alpha:qps=100").unwrap_err();
assert_eq!(err, ScopeParseError::UnknownKey("qps".into()));
}
#[test]
fn parse_missing_eq_errors() {
let err = parse_token_entry("alpha:tenant").unwrap_err();
assert_eq!(err, ScopeParseError::UnknownKey("tenant".into()));
}
#[test]
fn parse_wildcard_mixed_with_ids_errors() {
let err = parse_token_entry("alpha:tenant=1,*,2").unwrap_err();
assert_eq!(err, ScopeParseError::WildcardMixedWithIds);
}
#[test]
fn parse_tenant_list_with_trailing_comma_is_invalid() {
let err = parse_token_entry("alpha:tenant=1,").unwrap_err();
assert_eq!(err, ScopeParseError::InvalidTenant(String::new()));
}
#[test]
fn parse_env_empty_returns_default() {
let parsed = parse_tokens_env("");
assert!(parsed.token_scopes.is_empty());
assert_eq!(parsed.deprecated_count, 0);
let parsed = parse_tokens_env(" ");
assert!(parsed.token_scopes.is_empty());
assert_eq!(parsed.deprecated_count, 0);
}
#[test]
fn parse_env_single_bare_token() {
let parsed = parse_tokens_env("alpha");
assert_eq!(parsed.deprecated_count, 1);
assert_eq!(parsed.token_scopes.len(), 1);
assert!(parsed.token_scopes.get("alpha").unwrap().tenants.is_all());
}
#[test]
fn parse_env_three_bare_tokens() {
let parsed = parse_tokens_env("alpha,beta,gamma");
assert_eq!(parsed.deprecated_count, 3);
assert_eq!(parsed.token_scopes.len(), 3);
for name in ["alpha", "beta", "gamma"] {
assert!(
parsed.token_scopes.get(name).unwrap().tenants.is_all(),
"{name} not wildcard",
);
}
}
#[test]
fn parse_env_mixed_bare_and_scoped() {
let parsed = parse_tokens_env("alpha,beta:tenant=1,2,gamma:tenant=*");
assert_eq!(parsed.deprecated_count, 1, "only `alpha` is bare");
assert_eq!(parsed.token_scopes.len(), 3);
assert!(parsed.token_scopes["alpha"].tenants.is_all());
assert_eq!(
parsed.token_scopes["beta"].tenants,
TenantScope::Set(ids([1, 2]))
);
assert!(parsed.token_scopes["gamma"].tenants.is_all());
}
#[test]
fn parse_env_scoped_token_with_long_tenant_list() {
let parsed = parse_tokens_env("alpha:tenant=1,2,3,4,5,6,7");
assert_eq!(parsed.deprecated_count, 0);
assert_eq!(parsed.token_scopes.len(), 1);
let scope = &parsed.token_scopes["alpha"];
for t in 1..=7u64 {
assert!(scope.allows(TenantId(t)));
}
assert!(!scope.allows(TenantId(8)));
}
#[test]
fn parse_env_handles_whitespace_around_entries() {
let parsed = parse_tokens_env(" alpha , beta:tenant=1 , gamma:tenant=* ");
assert_eq!(parsed.token_scopes.len(), 3);
assert!(parsed.token_scopes.contains_key("alpha"));
assert!(parsed.token_scopes.contains_key("beta"));
assert!(parsed.token_scopes.contains_key("gamma"));
}
#[test]
fn parse_env_drops_malformed_entry_but_keeps_others() {
let parsed = parse_tokens_env("alpha,:tenant=1,gamma:tenant=*");
assert_eq!(parsed.token_scopes.len(), 2);
assert!(parsed.token_scopes.contains_key("alpha"));
assert!(parsed.token_scopes.contains_key("gamma"));
assert_eq!(parsed.deprecated_count, 1, "only `alpha` was bare");
}
#[test]
fn parse_env_two_scoped_tokens_with_continuation_lists() {
let parsed = parse_tokens_env("alpha:tenant=1,2,3,beta:tenant=4,5,6");
assert_eq!(parsed.deprecated_count, 0);
assert_eq!(parsed.token_scopes.len(), 2);
assert_eq!(
parsed.token_scopes["alpha"].tenants,
TenantScope::Set(ids([1, 2, 3]))
);
assert_eq!(
parsed.token_scopes["beta"].tenants,
TenantScope::Set(ids([4, 5, 6]))
);
}
#[test]
fn parse_env_bare_token_after_scoped_token() {
let parsed = parse_tokens_env("a:tenant=*,bare,b:tenant=4,5");
assert_eq!(parsed.token_scopes.len(), 3);
assert!(parsed.token_scopes["a"].tenants.is_all());
assert!(parsed.token_scopes["bare"].tenants.is_all());
assert_eq!(parsed.deprecated_count, 1, "`bare` is bare");
assert_eq!(
parsed.token_scopes["b"].tenants,
TenantScope::Set(ids([4, 5]))
);
}
#[test]
fn tenant_scope_allows_in_set() {
let s = TokenScope::from_tenants([TenantId(1), TenantId(2)]);
assert!(s.allows(TenantId(1)));
assert!(s.allows(TenantId(2)));
assert!(!s.allows(TenantId(3)));
}
#[test]
fn tenant_scope_wildcard_admits_arbitrary() {
let s = TokenScope::all();
assert!(s.allows(TenantId(0)));
assert!(s.allows(TenantId(u64::MAX)));
}
}