use anyhow::{Context, Result};
use crate::secrets;
const ENV_TOKEN: &str = "KANADE_NATS_TOKEN";
const REG_VALUE: &str = "NatsToken";
const NAME_PREFIX: &str = "kanade-";
const NAME_SEP: char = '/';
const REG_SHARED_SUBKEY: &str = r"SOFTWARE\kanade\agent";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NatsRole {
Agent,
Backend,
Cli,
}
impl NatsRole {
pub fn as_str(self) -> &'static str {
match self {
NatsRole::Agent => "agent",
NatsRole::Backend => "backend",
NatsRole::Cli => "cli",
}
}
fn reg_subkey(self) -> String {
format!(r"SOFTWARE\kanade\{}", self.as_str())
}
}
fn resolve_token_with(
role: NatsRole,
read_reg: impl Fn(&str, &str) -> Option<String>,
env: Option<String>,
) -> Option<String> {
if let Some(t) = read_reg(&role.reg_subkey(), REG_VALUE) {
return Some(t);
}
if let Some(t) = read_reg(REG_SHARED_SUBKEY, REG_VALUE) {
return Some(t);
}
env.filter(|t| !t.is_empty())
}
fn resolve_token(role: NatsRole) -> Option<String> {
resolve_token_with(
role,
secrets::read_hklm_value,
std::env::var(ENV_TOKEN).ok(),
)
}
pub fn client_name(role: NatsRole, identity: Option<&str>) -> String {
match identity.map(str::trim).filter(|s| !s.is_empty()) {
Some(id) => format!("{NAME_PREFIX}{}{NAME_SEP}{id}", role.as_str()),
None => format!("{NAME_PREFIX}{}", role.as_str()),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ClientName<'a> {
pub role: &'a str,
pub identity: Option<&'a str>,
}
pub fn parse_client_name(name: &str) -> Option<ClientName<'_>> {
let rest = name.strip_prefix(NAME_PREFIX)?;
Some(match rest.split_once(NAME_SEP) {
Some((role, id)) if !role.is_empty() && !id.is_empty() => ClientName {
role,
identity: Some(id),
},
Some((role, _)) if !role.is_empty() => ClientName {
role,
identity: None,
},
Some(_) => return None,
None if !rest.is_empty() => ClientName {
role: rest,
identity: None,
},
None => return None,
})
}
pub struct CredentialProbe {
presented: Credential,
}
enum Credential {
None,
Token(String),
User { name: String },
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CredentialKind {
None,
Token,
User,
}
impl std::fmt::Debug for CredentialProbe {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let rendered = match &self.presented {
Credential::None => "<none>".to_string(),
Credential::Token(_) => "<redacted token>".to_string(),
Credential::User { name } => format!("user {name}"),
};
f.debug_struct("CredentialProbe")
.field("presented", &rendered)
.finish()
}
}
impl CredentialProbe {
pub fn for_role(role: NatsRole) -> Self {
Self {
presented: match resolve_token(role) {
Some(t) => Credential::Token(t),
None => Credential::None,
},
}
}
pub fn from_token(token: Option<String>) -> Self {
Self {
presented: match token {
Some(t) => Credential::Token(t),
None => Credential::None,
},
}
}
pub fn from_user(name: impl Into<String>) -> Self {
Self {
presented: Credential::User { name: name.into() },
}
}
pub fn kind(&self) -> CredentialKind {
match &self.presented {
Credential::None => CredentialKind::None,
Credential::Token(_) => CredentialKind::Token,
Credential::User { .. } => CredentialKind::User,
}
}
pub fn is_ours(&self, candidate: &str) -> bool {
match &self.presented {
Credential::Token(t) => t == candidate,
Credential::None | Credential::User { .. } => false,
}
}
}
pub async fn connect(role: NatsRole, url: &str) -> Result<async_nats::Client> {
connect_inner(
role,
url,
None,
None::<fn(async_nats::Event) -> std::future::Ready<()>>,
)
.await
}
pub async fn connect_with_event_callback<F, Fut>(
role: NatsRole,
url: &str,
identity: Option<&str>,
cb: F,
) -> Result<async_nats::Client>
where
F: Fn(async_nats::Event) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = ()> + Send + Sync + 'static,
{
connect_inner(role, url, identity, Some(cb)).await
}
async fn connect_inner<F, Fut>(
role: NatsRole,
url: &str,
identity: Option<&str>,
cb: Option<F>,
) -> Result<async_nats::Client>
where
F: Fn(async_nats::Event) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = ()> + Send + Sync + 'static,
{
let opts = async_nats::ConnectOptions::new()
.retry_on_initial_connect()
.name(client_name(role, identity));
let opts = match resolve_token(role) {
Some(token) => opts.token(token),
None => opts,
};
let opts = match cb {
Some(cb) => opts.event_callback(cb),
None => opts,
};
opts.connect(url)
.await
.with_context(|| format!("connect to NATS at {url}"))
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
fn reg(entries: &[(&str, &str)]) -> impl Fn(&str, &str) -> Option<String> {
let map: HashMap<String, String> = entries
.iter()
.map(|(k, v)| ((*k).to_string(), (*v).to_string()))
.collect();
move |subkey: &str, value: &str| map.get(&format!(r"{subkey}\{value}")).cloned()
}
#[test]
fn role_subkeys_are_distinct_and_agent_matches_the_shared_path() {
assert_eq!(NatsRole::Backend.reg_subkey(), r"SOFTWARE\kanade\backend");
assert_eq!(NatsRole::Cli.reg_subkey(), r"SOFTWARE\kanade\cli");
assert_eq!(NatsRole::Agent.reg_subkey(), REG_SHARED_SUBKEY);
}
#[test]
fn an_unmigrated_fleet_keeps_presenting_the_shared_token() {
let registry = reg(&[(r"SOFTWARE\kanade\agent\NatsToken", "shared")]);
for role in [NatsRole::Agent, NatsRole::Backend, NatsRole::Cli] {
assert_eq!(
resolve_token_with(role, ®istry, None).as_deref(),
Some("shared"),
"{role:?} must keep working before its own key is provisioned"
);
}
}
#[test]
fn a_role_key_wins_over_the_shared_one() {
let registry = reg(&[
(r"SOFTWARE\kanade\agent\NatsToken", "shared"),
(r"SOFTWARE\kanade\backend\NatsToken", "backend-only"),
]);
assert_eq!(
resolve_token_with(NatsRole::Backend, ®istry, None).as_deref(),
Some("backend-only")
);
assert_eq!(
resolve_token_with(NatsRole::Cli, ®istry, None).as_deref(),
Some("shared")
);
}
#[test]
fn removing_a_role_key_falls_back_rather_than_failing() {
let registry = reg(&[(r"SOFTWARE\kanade\agent\NatsToken", "shared")]);
assert_eq!(
resolve_token_with(NatsRole::Backend, ®istry, None).as_deref(),
Some("shared")
);
}
#[test]
fn the_registry_outranks_the_environment() {
let registry = reg(&[(r"SOFTWARE\kanade\agent\NatsToken", "shared")]);
assert_eq!(
resolve_token_with(NatsRole::Agent, ®istry, Some("from-env".into())).as_deref(),
Some("shared")
);
}
#[test]
fn the_environment_serves_hosts_with_no_registry_at_all() {
let empty = reg(&[]);
assert_eq!(
resolve_token_with(NatsRole::Cli, &empty, Some("from-env".into())).as_deref(),
Some("from-env")
);
assert_eq!(
resolve_token_with(NatsRole::Cli, &empty, Some(String::new())),
None
);
assert_eq!(resolve_token_with(NatsRole::Cli, &empty, None), None);
}
#[test]
fn an_identity_round_trips_through_the_connection_name() {
for pc in ["PC001", "minipc", "Web%01", "ws-9"] {
let name = client_name(NatsRole::Agent, Some(pc));
let parsed = parse_client_name(&name).expect("our own name must parse");
assert_eq!(parsed.role, "agent");
assert_eq!(parsed.identity, Some(pc));
}
}
#[test]
fn a_role_without_an_identity_keeps_the_pre_1270_name() {
assert_eq!(client_name(NatsRole::Backend, None), "kanade-backend");
let parsed = parse_client_name("kanade-agent").unwrap();
assert_eq!(parsed.role, "agent");
assert_eq!(parsed.identity, None);
assert_eq!(client_name(NatsRole::Agent, Some(" ")), "kanade-agent");
}
#[test]
fn foreign_connections_do_not_parse_as_kanade_ones() {
assert!(parse_client_name("NATS CLI Version 0.1.5").is_none());
assert!(parse_client_name("").is_none());
assert!(parse_client_name("kanade-").is_none());
assert!(parse_client_name("kanade-/PC001").is_none());
assert_eq!(parse_client_name("kanade-agent/").unwrap().identity, None);
}
#[test]
fn an_unknown_role_is_preserved_rather_than_dropped() {
let parsed = parse_client_name("kanade-relay/PC001").unwrap();
assert_eq!(parsed.role, "relay");
assert_eq!(parsed.identity, Some("PC001"));
}
#[test]
fn the_probe_recognises_only_the_credential_we_present() {
let probe = CredentialProbe::from_token(Some("shared".into()));
assert_eq!(probe.kind(), CredentialKind::Token);
assert!(probe.is_ours("shared"));
assert!(!probe.is_ours("something-else"));
assert!(!probe.is_ours(""));
}
#[test]
fn a_probe_with_no_credential_matches_nothing() {
let probe = CredentialProbe::from_token(None);
assert_eq!(probe.kind(), CredentialKind::None);
assert!(!probe.is_ours(""));
assert!(!probe.is_ours("anything"));
}
#[test]
fn a_username_is_not_a_secret_we_can_recognise() {
let probe = CredentialProbe::from_user("kanade-backend");
assert_eq!(probe.kind(), CredentialKind::User);
assert!(!probe.is_ours("kanade-backend"));
}
#[test]
fn the_probe_never_prints_the_credential() {
let probe = CredentialProbe::from_token(Some("super-secret-token".into()));
let rendered = format!("{probe:?}");
assert!(!rendered.contains("super-secret-token"), "{rendered}");
assert!(rendered.contains("redacted"), "{rendered}");
assert!(
format!("{:?}", CredentialProbe::from_token(None)).contains("none"),
"the no-credential case should be visible, just not the value"
);
assert!(
format!("{:?}", CredentialProbe::from_user("kanade-backend"))
.contains("kanade-backend"),
);
}
}