use std::any::TypeId;
use std::collections::HashMap;
use std::fmt;
use std::io;
use std::path::Path;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64;
const OPERATOR_TOKEN: &str = "operator";
const REPLICATION_TOKEN: &str = "replication";
pub trait Role: Copy + Eq + fmt::Debug + Send + Sync + 'static {
const ROLES: &'static [(&'static str, Self)];
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NoRoles {}
impl Role for NoRoles {
const ROLES: &'static [(&'static str, Self)] = &[];
}
pub fn validate_roles<R: Role>() -> Result<(), String> {
let roles = R::ROLES;
if roles.len() > RoleId::CAPACITY {
return Err(format!(
"role table: {} roles, more than the {} a role index holds",
roles.len(),
RoleId::CAPACITY
));
}
for (i, &(token, role)) in roles.iter().enumerate() {
if !is_valid_token(token) {
return Err(format!(
"role table: '{token}' is not a valid role token (lowercase ASCII \
letters, digits, '-' and '_', starting with a letter)"
));
}
if token == OPERATOR_TOKEN || token == REPLICATION_TOKEN {
return Err(format!(
"role table: '{token}' is the runtime's own role token"
));
}
for &(earlier_token, earlier_role) in &roles[..i] {
if earlier_token == token {
return Err(format!("role table: token '{token}' is listed twice"));
}
if earlier_role == role {
return Err(format!(
"role table: {role:?} is listed twice, as '{earlier_token}' and '{token}'"
));
}
}
}
Ok(())
}
fn is_valid_token(token: &str) -> bool {
let mut bytes = token.bytes();
matches!(bytes.next(), Some(b'a'..=b'z'))
&& bytes.all(|b| matches!(b, b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_'))
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RoleId(
u8,
);
impl RoleId {
const CAPACITY: usize = u8::MAX as usize + 1;
pub(crate) fn index(self) -> usize {
usize::from(self.0)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ClientRole<R> {
Operator,
App(R),
}
impl<R> ClientRole<R> {
pub fn is_operator(&self) -> bool {
matches!(self, ClientRole::Operator)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KeyRole {
Replication,
Client(ClientRole<RoleId>),
}
impl KeyRole {
pub fn is_replication(self) -> bool {
matches!(self, KeyRole::Replication)
}
pub fn client(self) -> Option<ClientRole<RoleId>> {
match self {
KeyRole::Replication => None,
KeyRole::Client(role) => Some(role),
}
}
}
#[derive(Debug)]
pub struct AuthorizedKeys {
keys: HashMap<[u8; 32], KeyRole>,
app_tokens: Box<[&'static str]>,
role_type: TypeId,
}
impl AuthorizedKeys {
pub fn load<R: Role>(path: &Path) -> io::Result<Self> {
let content = std::fs::read_to_string(path)?;
Self::parse::<R>(&content).map_err(|e| io::Error::other(format!("{path:?}: {e}")))
}
pub fn parse<R: Role>(content: &str) -> Result<Self, String> {
validate_roles::<R>()?;
let mut keys = HashMap::new();
for (line_num, line) in content.lines().enumerate() {
let line = line.trim();
if line.is_empty() || line.starts_with('#') {
continue;
}
let mut parts = line.split_whitespace();
let role_token = parts
.next()
.ok_or_else(|| format!("line {}: missing role", line_num + 1))?;
let key_b64 = parts
.next()
.ok_or_else(|| format!("line {}: missing public key", line_num + 1))?;
let role = role_for_token::<R>(role_token).ok_or_else(|| {
let expected: Vec<&str> = [OPERATOR_TOKEN, REPLICATION_TOKEN]
.into_iter()
.chain(R::ROLES.iter().map(|&(token, _)| token))
.collect();
format!(
"line {}: unknown role '{role_token}' (expected {})",
line_num + 1,
expected.join(", ")
)
})?;
let key_bytes = BASE64
.decode(key_b64)
.map_err(|e| format!("line {}: invalid base64: {e}", line_num + 1))?;
if key_bytes.len() != 32 {
return Err(format!(
"line {}: public key must be 32 bytes, got {}",
line_num + 1,
key_bytes.len()
));
}
let mut key = [0u8; 32];
key.copy_from_slice(&key_bytes);
if keys.insert(key, role).is_some() {
return Err(format!(
"line {}: public key already listed on an earlier line",
line_num + 1
));
}
}
Ok(Self {
keys,
app_tokens: R::ROLES.iter().map(|&(token, _)| token).collect(),
role_type: TypeId::of::<R>(),
})
}
pub fn lookup(&self, public_key: &[u8; 32]) -> Option<KeyRole> {
self.keys.get(public_key).copied()
}
pub fn token(&self, role: KeyRole) -> &'static str {
match role {
KeyRole::Replication => REPLICATION_TOKEN,
KeyRole::Client(ClientRole::Operator) => OPERATOR_TOKEN,
KeyRole::Client(ClientRole::App(id)) => {
self.app_tokens.get(id.index()).copied().unwrap_or("?")
}
}
}
pub(crate) fn is_for<R: Role>(&self) -> bool {
self.role_type == TypeId::of::<R>()
}
pub fn len(&self) -> usize {
self.keys.len()
}
pub fn is_empty(&self) -> bool {
self.keys.is_empty()
}
}
fn role_for_token<R: Role>(token: &str) -> Option<KeyRole> {
match token {
OPERATOR_TOKEN => Some(KeyRole::Client(ClientRole::Operator)),
REPLICATION_TOKEN => Some(KeyRole::Replication),
_ => R::ROLES
.iter()
.position(|&(app_token, _)| app_token == token)
.map(|index| {
let id = u8::try_from(index).expect("role table validated to fit a u8 index");
KeyRole::Client(ClientRole::App(RoleId(id)))
}),
}
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
const ZERO_KEY: &str = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum DeskRole {
Trader,
Custodian,
ReadOnly,
}
impl Role for DeskRole {
const ROLES: &'static [(&'static str, Self)] = &[
("trader", DeskRole::Trader),
("custodian", DeskRole::Custodian),
("readonly", DeskRole::ReadOnly),
];
}
fn listing<R: Role>(token: &str) -> Result<AuthorizedKeys, String> {
AuthorizedKeys::parse::<R>(&format!("{token} {ZERO_KEY} test\n"))
}
macro_rules! role_table {
($name:ident, $table:expr) => {
#[allow(dead_code)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum $name {
A,
B,
}
impl Role for $name {
const ROLES: &'static [(&'static str, Self)] = $table;
}
};
}
#[test]
fn a_valid_role_table_is_accepted() {
validate_roles::<DeskRole>().unwrap();
validate_roles::<NoRoles>().unwrap();
role_table!(Mixed, &[("desk-2", Mixed::A), ("read_only9", Mixed::B)]);
validate_roles::<Mixed>().unwrap();
}
#[test]
fn a_runtime_token_in_the_role_table_is_refused() {
role_table!(ClaimsOperator, &[("operator", ClaimsOperator::A)]);
role_table!(ClaimsReplication, &[("replication", ClaimsReplication::A)]);
for err in [
validate_roles::<ClaimsOperator>().unwrap_err(),
validate_roles::<ClaimsReplication>().unwrap_err(),
] {
assert!(err.contains("is the runtime's own role token"), "{err}");
}
}
#[test]
fn a_token_listed_twice_is_refused() {
role_table!(Twice, &[("desk", Twice::A), ("desk", Twice::B)]);
let err = validate_roles::<Twice>().unwrap_err();
assert_eq!(err, "role table: token 'desk' is listed twice");
}
#[test]
fn a_role_listed_twice_is_refused() {
role_table!(Alias, &[("desk", Alias::A), ("trading-desk", Alias::A)]);
let err = validate_roles::<Alias>().unwrap_err();
assert_eq!(
err,
"role table: A is listed twice, as 'desk' and 'trading-desk'"
);
}
#[test]
fn a_token_outside_the_charset_is_refused() {
role_table!(Capital, &[("Trader", Capital::A)]);
role_table!(Empty, &[("", Empty::A)]);
role_table!(Space, &[("read only", Space::A)]);
role_table!(Comment, &[("#desk", Comment::A)]);
role_table!(LeadingDigit, &[("9desk", LeadingDigit::A)]);
role_table!(Colon, &[("desk:1", Colon::A)]);
for err in [
validate_roles::<Capital>().unwrap_err(),
validate_roles::<Empty>().unwrap_err(),
validate_roles::<Space>().unwrap_err(),
validate_roles::<Comment>().unwrap_err(),
validate_roles::<LeadingDigit>().unwrap_err(),
validate_roles::<Colon>().unwrap_err(),
] {
assert!(err.contains("is not a valid role token"), "{err}");
}
}
#[test]
fn a_table_longer_than_a_role_id_indexes_is_refused() {
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct Numbered(u16);
const TABLE: [(&str, Numbered); RoleId::CAPACITY + 1] = {
let mut table = [("desk", Numbered(0)); RoleId::CAPACITY + 1];
let mut i = 0;
while i < table.len() {
table[i].1 = Numbered(i as u16);
i += 1;
}
table
};
impl Role for Numbered {
const ROLES: &'static [(&'static str, Self)] = &TABLE;
}
let err = validate_roles::<Numbered>().unwrap_err();
assert_eq!(
err,
"role table: 257 roles, more than the 256 a role index holds"
);
}
#[test]
fn parsing_validates_the_role_table_even_for_an_empty_file() {
role_table!(ClaimsOperator, &[("operator", ClaimsOperator::A)]);
let err = AuthorizedKeys::parse::<ClaimsOperator>("").unwrap_err();
assert!(err.contains("is the runtime's own role token"), "{err}");
}
#[test]
fn client_role_is_operator() {
assert!(ClientRole::<DeskRole>::Operator.is_operator());
assert!(!ClientRole::App(DeskRole::Trader).is_operator());
}
#[test]
fn only_replication_has_no_client_role() {
assert_eq!(KeyRole::Replication.client(), None);
assert!(KeyRole::Replication.is_replication());
let operator = KeyRole::Client(ClientRole::Operator);
assert_eq!(operator.client(), Some(ClientRole::Operator));
assert!(!operator.is_replication());
}
#[test]
fn runtime_tokens_parse_to_the_runtime_roles() {
for keys in [
[
listing::<DeskRole>("operator"),
listing::<DeskRole>("replication"),
],
[
listing::<NoRoles>("operator"),
listing::<NoRoles>("replication"),
],
] {
let [operator, replication] = keys.map(Result::unwrap);
assert_eq!(
operator.lookup(&[0u8; 32]),
Some(KeyRole::Client(ClientRole::Operator))
);
assert_eq!(replication.lookup(&[0u8; 32]), Some(KeyRole::Replication));
}
}
#[test]
fn every_app_role_round_trips_through_its_token() {
for (index, &(token, _)) in DeskRole::ROLES.iter().enumerate() {
let keys = listing::<DeskRole>(token).unwrap();
let role = keys.lookup(&[0u8; 32]).unwrap();
let Some(ClientRole::App(id)) = role.client() else {
panic!("{token} must parse to an application role, got {role:?}");
};
assert_eq!(id.index(), index);
assert_eq!(keys.token(role), token);
}
}
#[test]
fn the_runtime_roles_are_named_by_their_tokens() {
let keys = AuthorizedKeys::parse::<NoRoles>("").unwrap();
assert_eq!(keys.token(KeyRole::Replication), "replication");
assert_eq!(
keys.token(KeyRole::Client(ClientRole::Operator)),
"operator"
);
}
#[test]
fn no_roles_admits_only_the_runtime_tokens() {
let err = listing::<NoRoles>("trader").unwrap_err();
assert_eq!(
err,
"line 1: unknown role 'trader' (expected operator, replication)"
);
}
#[test]
fn a_token_in_another_case_is_unknown() {
let err = listing::<DeskRole>("Trader").unwrap_err();
assert!(err.contains("unknown role 'Trader'"), "{err}");
}
#[test]
fn a_table_knows_the_role_type_it_was_parsed_for() {
let keys = AuthorizedKeys::parse::<DeskRole>("").unwrap();
assert!(keys.is_for::<DeskRole>());
assert!(!keys.is_for::<NoRoles>());
}
#[test]
fn parse_valid_keys_file() {
let content = "\
# Auth keys file
operator AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA= ops-team
trader AQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQE= market-maker-1
readonly AgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgI= monitoring
";
let keys = AuthorizedKeys::parse::<DeskRole>(content).unwrap();
assert_eq!(keys.len(), 3);
let admin_key = BASE64
.decode("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=")
.unwrap();
let mut k = [0u8; 32];
k.copy_from_slice(&admin_key);
assert_eq!(keys.lookup(&k), Some(KeyRole::Client(ClientRole::Operator)));
}
#[test]
fn parse_skips_comments_and_blanks() {
let content = "\
# comment
# indented comment
operator AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA= test
";
let keys = AuthorizedKeys::parse::<NoRoles>(content).unwrap();
assert_eq!(keys.len(), 1);
}
#[test]
fn parse_rejects_unknown_role_naming_every_valid_one() {
let err = listing::<DeskRole>("superuser").unwrap_err();
assert_eq!(
err,
"line 1: unknown role 'superuser' \
(expected operator, replication, trader, custodian, readonly)"
);
}
#[test]
fn parse_rejects_wrong_key_length() {
let content = "operator AQID test\n"; let result = AuthorizedKeys::parse::<NoRoles>(content);
assert!(result.is_err());
assert!(result.unwrap_err().contains("32 bytes"));
}
#[test]
fn lookup_missing_key_returns_none() {
let keys = AuthorizedKeys::parse::<NoRoles>("").unwrap();
assert!(keys.lookup(&[0u8; 32]).is_none());
}
#[test]
fn duplicate_key_is_refused() {
for content in [
"\
operator AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA= first
readonly AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA= second
",
"\
trader AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA= first
trader AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA= again
",
] {
let err = AuthorizedKeys::parse::<DeskRole>(content).unwrap_err();
assert!(
err.starts_with("line 2:") && err.contains("already listed"),
"unexpected error: {err}"
);
}
}
#[test]
fn empty_file_produces_empty_keys() {
let keys = AuthorizedKeys::parse::<NoRoles>("").unwrap();
assert!(keys.is_empty());
assert_eq!(keys.len(), 0);
assert!(keys.lookup(&[0u8; 32]).is_none());
}
#[test]
fn parse_rejects_invalid_base64() {
let content = "operator not-valid-base64!!! test\n";
let result = AuthorizedKeys::parse::<NoRoles>(content);
assert!(result.is_err());
assert!(result.unwrap_err().contains("invalid base64"));
}
#[test]
fn parse_rejects_missing_key_field() {
let content = "admin\n";
let result = AuthorizedKeys::parse::<NoRoles>(content);
assert!(result.is_err());
assert!(result.unwrap_err().contains("missing public key"));
}
#[test]
fn comments_only_file_produces_empty_keys() {
let content = "\
# only comments
# nothing else
# indented
";
let keys = AuthorizedKeys::parse::<NoRoles>(content).unwrap();
assert!(keys.is_empty());
}
#[test]
fn load_nonexistent_file_is_error() {
let result =
AuthorizedKeys::load::<NoRoles>(std::path::Path::new("/nonexistent/path/keys.txt"));
assert!(result.is_err());
}
#[test]
fn load_from_file() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("keys.txt");
std::fs::write(
&path,
"trader AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA= test\n",
)
.unwrap();
let keys = AuthorizedKeys::load::<DeskRole>(&path).unwrap();
assert_eq!(keys.len(), 1);
}
}