use std::collections::HashMap;
use std::sync::atomic::Ordering;
use std::sync::{Arc, Mutex, OnceLock, Weak};
use std::time::Duration;
use anyhow::{Context, Result, bail};
use tracing::{debug, info, warn};
use crate::secrets;
const ENV_TOKEN: &str = "KANADE_NATS_TOKEN";
const REG_VALUE: &str = "NatsToken";
const ENV_USER: &str = "KANADE_NATS_USER";
const ENV_PASSWORD: &str = "KANADE_NATS_PASSWORD";
const REG_USER: &str = "NatsUser";
const REG_PASSWORD: &str = "NatsPassword";
const PROBE_TIMEOUT: Duration = Duration::from_secs(3);
const PROBE_NAME: &str = "auth-probe";
const PING_INTERVAL: Duration = Duration::from_secs(60);
const CHECK_INTERVAL: Duration = Duration::from_secs(5);
const STALL_BOUND: Duration = Duration::from_secs(3 * PING_INTERVAL.as_secs());
const HEALTH_TIMEOUT: Duration = Duration::from_secs(5);
const RECONNECT_DELAY_MAX_SECS: u64 = 4;
const CONNECT_TIMEOUT_SECS: u64 = 5;
const ATTEMPT_MAX_SECS: u64 =
RECONNECT_DELAY_MAX_SECS + PROBE_TIMEOUT.as_secs() + 1 + CONNECT_TIMEOUT_SECS;
pub const RESUME_BOUND: Duration = Duration::from_secs(2 * ATTEMPT_MAX_SECS + 4);
pub const EXIT_BOUND: Duration = Liveness::PRODUCTION.exit_bound();
const _: () = {
assert!(STALL_BOUND.as_secs() >= PING_INTERVAL.as_secs() + 60);
assert!(RESUME_BOUND.as_secs() <= STALL_BOUND.as_secs());
assert!(EXIT_BOUND.as_secs() < 300);
};
#[derive(Debug, Clone, Copy)]
pub struct Liveness {
pub check_interval: Duration,
pub stall_bound: Duration,
}
impl Liveness {
pub const PRODUCTION: Liveness = Liveness {
check_interval: CHECK_INTERVAL,
stall_bound: STALL_BOUND,
};
pub const fn exit_bound(&self) -> Duration {
Duration::from_secs(
self.stall_bound.as_secs()
+ self.check_interval.as_secs()
+ 2 * HEALTH_TIMEOUT.as_secs(),
)
}
fn resume_grace(&self) -> Duration {
RESUME_BOUND.min(self.stall_bound)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Verdict {
Wait,
Stalled,
}
struct StallJudge {
liveness: Liveness,
deadline: tokio::time::Instant,
outage: bool,
next_witness: Option<tokio::time::Instant>,
}
impl StallJudge {
fn new(liveness: Liveness, now: tokio::time::Instant) -> Self {
Self {
liveness,
deadline: now + liveness.stall_bound,
outage: false,
next_witness: None,
}
}
fn progressed(&mut self, now: tokio::time::Instant) {
self.deadline = now + self.liveness.stall_bound;
self.outage = false;
self.next_witness = None;
}
fn witness_due(&self, now: tokio::time::Instant) -> bool {
now >= self.deadline && self.next_witness.is_none_or(|at| now >= at)
}
fn witnessed(&mut self, now: tokio::time::Instant, outcome: ProbeOutcome) -> Verdict {
if outcome == ProbeOutcome::Unreachable {
self.outage = true;
self.next_witness = Some(now + self.liveness.resume_grace());
return Verdict::Wait;
}
if std::mem::take(&mut self.outage) {
self.deadline = self.deadline.max(now + self.liveness.resume_grace());
self.next_witness = None;
return Verdict::Wait;
}
Verdict::Stalled
}
}
const AUTH_REJECTION_LIMIT: usize = 5;
const AUTH_REJECTION_WINDOW: Duration = Duration::from_secs(10);
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(),
)
}
#[derive(Clone, PartialEq, Eq)]
struct UserCredential {
name: String,
password: String,
}
impl std::fmt::Debug for UserCredential {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("UserCredential")
.field("name", &self.name)
.field("password", &"<redacted>")
.finish()
}
}
fn pair(
name: Option<String>,
password: Option<String>,
source: &str,
) -> Result<Option<UserCredential>> {
match (name, password) {
(Some(name), Some(password)) => Ok(Some(UserCredential { name, password })),
(None, None) => Ok(None),
(Some(_), None) => bail!("NATS user is set but its password is missing ({source})"),
(None, Some(_)) => bail!("NATS password is set but its user name is missing ({source})"),
}
}
fn resolve_user_with(
role: NatsRole,
read_reg: impl Fn(&str, &str) -> Option<String>,
env_user: Option<String>,
env_password: Option<String>,
) -> Result<Option<UserCredential>> {
let subkey = role.reg_subkey();
let reg_user = read_reg(&subkey, REG_USER);
let reg_password = read_reg(&subkey, REG_PASSWORD);
if reg_user.is_some() || reg_password.is_some() {
let source = format!(r"HKLM\{subkey}: {REG_USER} / {REG_PASSWORD}");
return pair(reg_user, reg_password, &source);
}
let source = format!("${ENV_USER} / ${ENV_PASSWORD}");
pair(
env_user.filter(|v| !v.is_empty()),
env_password.filter(|v| !v.is_empty()),
&source,
)
}
fn resolve_user(role: NatsRole) -> Result<Option<UserCredential>> {
resolve_user_with(
role,
secrets::read_hklm_value,
std::env::var(ENV_USER).ok(),
std::env::var(ENV_PASSWORD).ok(),
)
}
#[derive(Clone, Default)]
pub struct NatsCredentials {
token: Option<String>,
user: Option<UserCredential>,
}
impl NatsCredentials {
pub fn new(token: Option<String>, user: Option<(String, String)>) -> Self {
Self {
token,
user: user.map(|(name, password)| UserCredential { name, password }),
}
}
fn resolve(role: NatsRole) -> Result<Self> {
Ok(Self {
token: resolve_token(role),
user: resolve_user(role)?,
})
}
}
impl std::fmt::Debug for NatsCredentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("NatsCredentials")
.field("token", &self.token.as_ref().map(|_| "<redacted>"))
.field("user", &self.user)
.finish()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Choice {
User,
Token,
}
impl Choice {
fn label(self) -> &'static str {
match self {
Choice::User => "user",
Choice::Token => "token",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ProbeOutcome {
Accepted,
Rejected,
Unreachable,
}
impl ProbeOutcome {
fn label(self) -> &'static str {
match self {
ProbeOutcome::Accepted => "accepted",
ProbeOutcome::Rejected => "rejected",
ProbeOutcome::Unreachable => "unreachable",
}
}
}
#[derive(Default)]
struct Live {
decision: Mutex<Option<Choice>>,
refusals: Mutex<Option<(std::time::Instant, usize)>>,
}
impl Live {
fn observe(&self, ev: &async_nats::Event) {
let mut refusals = self.refusals.lock().unwrap_or_else(|e| e.into_inner());
match ev {
async_nats::Event::Connected | async_nats::Event::Disconnected => *refusals = None,
async_nats::Event::ClientError(async_nats::ClientError::Other(kind))
if *kind == async_nats::ConnectErrorKind::AuthorizationViolation.to_string() =>
{
let (first, n) = refusals.unwrap_or((std::time::Instant::now(), 0));
*refusals = Some((first, n + 1));
}
async_nats::Event::ClientError(async_nats::ClientError::Other(kind))
if *kind == async_nats::ConnectErrorKind::Authentication.to_string() => {}
async_nats::Event::ClientError(_) => *refusals = None,
_ => {}
}
}
fn auth_failed(&self) -> bool {
match *self.refusals.lock().unwrap_or_else(|e| e.into_inner()) {
Some((first, n)) => {
n >= AUTH_REJECTION_LIMIT && first.elapsed() >= AUTH_REJECTION_WINDOW
}
None => false,
}
}
fn get(&self) -> Option<Choice> {
*self.decision.lock().unwrap_or_else(|e| e.into_inner())
}
fn set(&self, c: Choice) {
*self.decision.lock().unwrap_or_else(|e| e.into_inner()) = Some(c);
}
}
fn role_lives() -> &'static Mutex<HashMap<&'static str, Arc<Live>>> {
static LIVE: OnceLock<Mutex<HashMap<&'static str, Arc<Live>>>> = OnceLock::new();
LIVE.get_or_init(Default::default)
}
fn publish_role_live(role: NatsRole, live: &Live) {
let reported = live_for(role);
*reported.decision.lock().unwrap_or_else(|e| e.into_inner()) = live.get();
}
fn live_for(role: NatsRole) -> Arc<Live> {
role_lives()
.lock()
.unwrap_or_else(|e| e.into_inner())
.entry(role.as_str())
.or_default()
.clone()
}
struct ConnectionHealth {
live: Arc<Live>,
url: String,
creds: NatsCredentials,
}
struct HealthEntry {
stats: Weak<async_nats::client::Statistics>,
health: Arc<ConnectionHealth>,
}
fn health_entries() -> &'static Mutex<Vec<HealthEntry>> {
static HEALTH: OnceLock<Mutex<Vec<HealthEntry>>> = OnceLock::new();
HEALTH.get_or_init(Default::default)
}
fn register_health(client: &async_nats::Client, health: Arc<ConnectionHealth>) {
let mut entries = health_entries().lock().unwrap_or_else(|e| e.into_inner());
entries.retain(|entry| entry.stats.strong_count() > 0);
entries.push(HealthEntry {
stats: Arc::downgrade(&client.statistics()),
health,
});
}
fn connection_health(client: &async_nats::Client) -> Option<Arc<ConnectionHealth>> {
let stats = client.statistics();
health_entries()
.lock()
.unwrap_or_else(|e| e.into_inner())
.iter()
.find(|entry| entry.stats.ptr_eq(&Arc::downgrade(&stats)))
.map(|entry| entry.health.clone())
}
async fn witness(health: &ConnectionHealth) -> ProbeOutcome {
let attempt = async {
let mut opts = async_nats::ConnectOptions::new()
.name("auth-witness")
.connection_timeout(HEALTH_TIMEOUT);
match &health.creds.user {
Some(user) => match probe_user(&health.url, user).await {
ProbeOutcome::Accepted => {
opts = opts.user_and_password(user.name.clone(), user.password.clone());
}
ProbeOutcome::Rejected => match &health.creds.token {
Some(token) => opts = opts.token(token.clone()),
None => return ProbeOutcome::Rejected,
},
ProbeOutcome::Unreachable => return ProbeOutcome::Unreachable,
},
None => {
if let Some(token) = &health.creds.token {
opts = opts.token(token.clone());
}
}
}
match opts.connect(&health.url).await {
Ok(_) => ProbeOutcome::Accepted,
Err(error) if error.kind() == async_nats::ConnectErrorKind::AuthorizationViolation => {
ProbeOutcome::Rejected
}
Err(_) => ProbeOutcome::Unreachable,
}
};
tokio::time::timeout(HEALTH_TIMEOUT, attempt)
.await
.unwrap_or(ProbeOutcome::Unreachable)
}
fn select(
outcome: ProbeOutcome,
_last_worked: Option<Choice>,
have_token: bool,
) -> std::result::Result<Choice, &'static str> {
match outcome {
ProbeOutcome::Accepted => Ok(Choice::User),
ProbeOutcome::Rejected if have_token => Ok(Choice::Token),
ProbeOutcome::Rejected => {
Err("the broker rejected the NATS user and no NatsToken is provisioned")
}
ProbeOutcome::Unreachable => {
Err("NATS credential probe inconclusive; deferring this attempt")
}
}
}
async fn probe_user(url: &str, user: &UserCredential) -> ProbeOutcome {
let (url, name, password) = (url.to_string(), user.name.clone(), user.password.clone());
let probe = tokio::spawn(async move {
let attempt = async_nats::ConnectOptions::new()
.name(PROBE_NAME)
.connection_timeout(PROBE_TIMEOUT)
.user_and_password(name, password)
.connect(url);
match tokio::time::timeout(PROBE_TIMEOUT + Duration::from_secs(1), attempt).await {
Ok(Ok(_client)) => ProbeOutcome::Accepted,
Ok(Err(e)) if e.kind() == async_nats::ConnectErrorKind::AuthorizationViolation => {
ProbeOutcome::Rejected
}
Ok(Err(_)) | Err(_) => ProbeOutcome::Unreachable,
}
});
probe.await.unwrap_or(ProbeOutcome::Unreachable)
}
async fn choose<P, Fut>(
role: NatsRole,
have_token: bool,
live: &Live,
probe: P,
) -> std::result::Result<Choice, &'static str>
where
P: FnOnce() -> Fut,
Fut: std::future::Future<Output = ProbeOutcome>,
{
let outcome = probe().await;
let previous = live.get();
let choice = select(outcome, previous, have_token).inspect_err(|_| {
if outcome == ProbeOutcome::Rejected {
live.observe(&async_nats::Event::ClientError(
async_nats::ClientError::Other(
async_nats::ConnectErrorKind::AuthorizationViolation.to_string(),
),
));
} else {
*live.refusals.lock().unwrap_or_else(|e| e.into_inner()) = None;
debug!(
role = role.as_str(),
"NATS credential probe inconclusive; deferring this attempt"
);
}
})?;
if outcome != ProbeOutcome::Unreachable {
live.set(choice);
live_for(role).set(choice);
}
if previous == Some(choice) {
debug!(
role = role.as_str(),
credential = choice.label(),
probe = outcome.label(),
"NATS credential selected"
);
} else {
info!(
role = role.as_str(),
credential = choice.label(),
probe = outcome.label(),
"NATS credential selected"
);
}
Ok(choice)
}
enum AuthPlan {
Static(Option<String>),
Select {
token: Option<String>,
user: UserCredential,
},
}
fn auth_plan(creds: NatsCredentials) -> AuthPlan {
match creds.user {
None => AuthPlan::Static(creds.token),
Some(user) => AuthPlan::Select {
token: creds.token,
user,
},
}
}
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 {
token: Option<String>,
user: Option<String>,
live: Option<Arc<Live>>,
}
#[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 mut held = Vec::new();
if self.token.is_some() {
held.push("<redacted token>".to_string());
}
if let Some(name) = &self.user {
held.push(format!("user {name}"));
}
let rendered = if held.is_empty() {
"<none>".to_string()
} else {
held.join(" + ")
};
f.debug_struct("CredentialProbe")
.field("presented", &rendered)
.field("kind", &self.kind())
.finish()
}
}
impl CredentialProbe {
pub fn for_role(role: NatsRole) -> Self {
Self {
token: resolve_token(role),
user: resolve_user(role).ok().flatten().map(|u| u.name),
live: Some(live_for(role)),
}
}
pub fn from_token(token: Option<String>) -> Self {
Self {
token,
user: None,
live: None,
}
}
pub fn from_user(name: impl Into<String>) -> Self {
Self {
token: None,
user: Some(name.into()),
live: None,
}
}
pub fn kind(&self) -> CredentialKind {
let decided = self.live.as_ref().and_then(|l| l.get());
match decided {
Some(Choice::User) if self.user.is_some() => return CredentialKind::User,
Some(Choice::Token) if self.token.is_some() => return CredentialKind::Token,
_ => {}
}
if self.live.is_none() && self.user.is_some() {
CredentialKind::User
} else if self.token.is_some() {
CredentialKind::Token
} else {
CredentialKind::None
}
}
pub fn is_ours(&self, candidate: &str) -> bool {
self.token.as_deref() == Some(candidate)
}
}
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<()>>,
NatsCredentials::resolve(role)?,
Arc::new(Live::default()),
)
.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),
NatsCredentials::resolve(role)?,
Arc::new(Live::default()),
)
.await
}
pub async fn connect_with_credentials(
role: NatsRole,
url: &str,
creds: NatsCredentials,
) -> Result<async_nats::Client> {
connect_inner(
role,
url,
None,
None::<fn(async_nats::Event) -> std::future::Ready<()>>,
creds,
Arc::new(Live::default()),
)
.await
}
pub async fn connect_with_credentials_and_event_callback<F, Fut>(
role: NatsRole,
url: &str,
creds: NatsCredentials,
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, None, Some(cb), creds, Arc::new(Live::default())).await
}
async fn connect_inner<F, Fut>(
role: NatsRole,
url: &str,
identity: Option<&str>,
cb: Option<F>,
creds: NatsCredentials,
live: Arc<Live>,
) -> Result<async_nats::Client>
where
F: Fn(async_nats::Event) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = ()> + Send + Sync + 'static,
{
let _ = rustls::crypto::ring::default_provider().install_default();
let health = Arc::new(ConnectionHealth {
live: live.clone(),
url: url.to_string(),
creds: creds.clone(),
});
publish_role_live(role, &live);
let tracker = live.clone();
let opts = match auth_plan(creds) {
AuthPlan::Static(token) => {
let opts = async_nats::ConnectOptions::new();
match token {
Some(token) => opts.token(token),
None => opts,
}
}
AuthPlan::Select { token, user } => {
let url = url.to_string();
async_nats::ConnectOptions::with_auth_callback(move |_nonce| {
let (url, token, user, live) =
(url.clone(), token.clone(), user.clone(), live.clone());
async move {
let choice = choose(role, token.is_some(), &live, || probe_user(&url, &user))
.await
.map_err(async_nats::AuthError::new)?;
let mut auth = async_nats::Auth::new();
match choice {
Choice::User => {
auth.username = Some(user.name);
auth.password = Some(user.password);
}
Choice::Token => auth.token = token,
}
Ok(auth)
}
})
}
};
let opts = opts
.retry_on_initial_connect()
.connection_timeout(HEALTH_TIMEOUT)
.ping_interval(PING_INTERVAL)
.name(client_name(role, identity));
let opts = opts.event_callback(move |ev| {
tracker.observe(&ev);
let forwarded = cb.as_ref().map(|cb| cb(ev));
async move {
if let Some(forwarded) = forwarded {
forwarded.await;
}
}
});
let client = opts
.connect(url)
.await
.with_context(|| format!("connect to NATS at {url}"))?;
register_health(&client, health);
Ok(client)
}
pub async fn is_dead(client: &async_nats::Client) -> bool {
match tokio::time::timeout(Duration::from_secs(5), client.flush()).await {
Ok(Err(e)) => match e.kind() {
async_nats::client::FlushErrorKind::SendError => true,
async_nats::client::FlushErrorKind::FlushError => false,
},
Ok(Ok(())) | Err(_) => false,
}
}
pub async fn wait_until_dead_with(role: NatsRole, client: &async_nats::Client, liveness: Liveness) {
let health = connection_health(client);
let mut tick = tokio::time::interval(liveness.check_interval);
tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
let stats = client.statistics();
let mut received = stats.in_bytes.load(Ordering::Relaxed);
let mut judge = StallJudge::new(liveness, tokio::time::Instant::now());
loop {
tick.tick().await;
let now_received = stats.in_bytes.load(Ordering::Relaxed);
if client.connection_state() == async_nats::connection::State::Connected
&& now_received != received
{
judge.progressed(tokio::time::Instant::now());
}
received = now_received;
if health.as_ref().is_some_and(|h| h.live.auth_failed()) {
warn!(
role = role.as_str(),
"the NATS broker keeps refusing this role's credential; the client cannot connect"
);
return;
}
if is_dead(client).await {
warn!(
role = role.as_str(),
"NATS connection task has terminated; the client can no longer talk to the broker"
);
return;
}
let Some(health) = &health else { continue };
if !judge.witness_due(tokio::time::Instant::now()) {
continue;
}
let outcome = witness(health).await;
if judge.witnessed(tokio::time::Instant::now(), outcome) != Verdict::Stalled {
continue;
}
if client.connection_state() == async_nats::connection::State::Connected
&& stats.in_bytes.load(Ordering::Relaxed) != received
{
judge.progressed(tokio::time::Instant::now());
continue;
}
warn!(
role = role.as_str(),
witness = outcome.label(),
"NATS client made no receive progress while a fresh credential witness reached the broker; exiting for a supervised restart"
);
return;
}
}
pub async fn wait_until_dead_every(
role: NatsRole,
client: &async_nats::Client,
interval: Duration,
) {
wait_until_dead_with(
role,
client,
Liveness {
check_interval: interval,
..Liveness::PRODUCTION
},
)
.await
}
pub async fn wait_until_dead(role: NatsRole, client: &async_nats::Client) {
wait_until_dead_with(role, client, Liveness::PRODUCTION).await
}
pub fn exit_on_dead_with(role: NatsRole, client: &async_nats::Client, liveness: Liveness) {
let client = client.clone();
tokio::spawn(async move {
wait_until_dead_with(role, &client, liveness).await;
std::process::exit(1);
});
}
pub fn exit_on_dead(role: NatsRole, client: &async_nats::Client) {
exit_on_dead_with(role, client, Liveness::PRODUCTION);
}
#[cfg(test)]
mod tests {
use super::*;
fn secs(n: u64) -> Duration {
Duration::from_secs(n)
}
fn simulate(
judge: &mut StallJudge,
start: tokio::time::Instant,
until: Duration,
progress_every: Option<Duration>,
outcome: ProbeOutcome,
witnesses: &mut Vec<Duration>,
) -> Option<Duration> {
let step = judge.liveness.check_interval;
let mut elapsed = Duration::ZERO;
let mut last_progress = Duration::ZERO;
while elapsed < until {
elapsed += step;
let now = start + elapsed;
if progress_every.is_some_and(|p| elapsed - last_progress >= p) {
last_progress = elapsed;
judge.progressed(now);
}
if judge.witness_due(now) {
witnesses.push(elapsed);
if judge.witnessed(now, outcome) == Verdict::Stalled {
return Some(elapsed);
}
}
}
None
}
#[test]
fn the_production_bounds_leave_a_healthy_idle_connection_a_wide_margin() {
assert!(STALL_BOUND >= PING_INTERVAL + secs(60));
assert!(EXIT_BOUND < secs(300));
assert_eq!(PING_INTERVAL, secs(60));
assert_eq!(STALL_BOUND, secs(180));
assert_eq!(EXIT_BOUND, secs(195));
}
#[test]
fn an_idle_healthy_client_is_never_judged_stalled() {
let start = tokio::time::Instant::now();
let mut judge = StallJudge::new(Liveness::PRODUCTION, start);
let mut witnesses = Vec::new();
let verdict = simulate(
&mut judge,
start,
secs(24 * 3600),
Some(PING_INTERVAL),
ProbeOutcome::Accepted,
&mut witnesses,
);
assert_eq!(verdict, None);
assert!(
witnesses.is_empty(),
"a witness connected to a healthy broker"
);
}
#[test]
fn an_offline_broker_never_exits_and_is_polled_at_the_resume_interval() {
let start = tokio::time::Instant::now();
let mut judge = StallJudge::new(Liveness::PRODUCTION, start);
let mut witnesses = Vec::new();
let verdict = simulate(
&mut judge,
start,
secs(3 * 3600),
None,
ProbeOutcome::Unreachable,
&mut witnesses,
);
assert_eq!(verdict, None);
assert_eq!(witnesses[0], STALL_BOUND);
assert!(witnesses.windows(2).all(|w| w[1] - w[0] >= RESUME_BOUND));
}
#[test]
fn a_stall_after_an_outage_is_judged_against_the_normal_bound() {
let start = tokio::time::Instant::now();
let mut judge = StallJudge::new(Liveness::PRODUCTION, start);
let mut witnesses = Vec::new();
assert_eq!(
simulate(
&mut judge,
start,
secs(3600),
None,
ProbeOutcome::Unreachable,
&mut witnesses
),
None
);
assert!(judge.outage);
let back = start + secs(3600);
witnesses.clear();
let exit = simulate(
&mut judge,
back,
EXIT_BOUND,
None,
ProbeOutcome::Accepted,
&mut witnesses,
);
assert!(!judge.outage);
let exit = exit.expect("a silent client with a reachable broker must exit");
assert!(exit <= EXIT_BOUND, "{exit:?}");
let mut judge = StallJudge::new(Liveness::PRODUCTION, start);
let mut witnesses = Vec::new();
simulate(
&mut judge,
start,
secs(3600),
None,
ProbeOutcome::Unreachable,
&mut witnesses,
);
assert!(judge.outage);
let resumed = start + secs(3600);
judge.progressed(resumed);
assert!(!judge.outage);
assert!(!judge.witness_due(resumed + STALL_BOUND - secs(1)));
witnesses.clear();
let exit = simulate(
&mut judge,
resumed,
STALL_BOUND + secs(60),
None,
ProbeOutcome::Rejected,
&mut witnesses,
);
assert_eq!(exit, Some(STALL_BOUND));
}
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"),
);
}
fn user(r: Result<Option<UserCredential>>) -> Option<(String, String)> {
r.unwrap().map(|u| (u.name, u.password))
}
#[test]
fn a_users_registry_pair_comes_from_the_roles_own_key() {
let registry = reg(&[
(r"SOFTWARE\kanade\backend\NatsUser", "kanade-backend"),
(r"SOFTWARE\kanade\backend\NatsPassword", "pw"),
]);
assert_eq!(
user(resolve_user_with(NatsRole::Backend, ®istry, None, None)),
Some(("kanade-backend".into(), "pw".into()))
);
for role in [NatsRole::Agent, NatsRole::Cli] {
assert_eq!(user(resolve_user_with(role, ®istry, None, None)), None);
}
let agent = reg(&[
(r"SOFTWARE\kanade\agent\NatsUser", "a"),
(r"SOFTWARE\kanade\agent\NatsPassword", "b"),
]);
assert_eq!(
user(resolve_user_with(NatsRole::Cli, &agent, None, None)),
None
);
}
#[test]
fn the_registry_pair_outranks_the_environment_pair() {
let registry = reg(&[
(r"SOFTWARE\kanade\cli\NatsUser", "reg-user"),
(r"SOFTWARE\kanade\cli\NatsPassword", "reg-pw"),
]);
assert_eq!(
user(resolve_user_with(
NatsRole::Cli,
®istry,
Some("env-user".into()),
Some("env-pw".into())
)),
Some(("reg-user".into(), "reg-pw".into()))
);
}
#[test]
fn the_environment_serves_a_user_pair_when_the_registry_has_none() {
assert_eq!(
user(resolve_user_with(
NatsRole::Cli,
reg(&[]),
Some("u".into()),
Some("p".into())
)),
Some(("u".into(), "p".into()))
);
assert_eq!(
user(resolve_user_with(
NatsRole::Cli,
reg(&[]),
Some(String::new()),
Some(String::new())
)),
None
);
assert_eq!(
user(resolve_user_with(NatsRole::Cli, reg(&[]), None, None)),
None
);
}
#[test]
fn half_a_pair_is_an_error_not_a_fallback() {
let err = |r: Result<Option<UserCredential>>| format!("{:#}", r.unwrap_err());
let only_user = reg(&[(r"SOFTWARE\kanade\cli\NatsUser", "u")]);
assert!(
err(resolve_user_with(NatsRole::Cli, &only_user, None, None))
.contains("password is missing")
);
let only_pw = reg(&[(r"SOFTWARE\kanade\cli\NatsPassword", "secret-pw")]);
let e = err(resolve_user_with(NatsRole::Cli, &only_pw, None, None));
assert!(e.contains("user name is missing"), "{e}");
assert!(!e.contains("secret-pw"), "{e}");
assert!(
resolve_user_with(
NatsRole::Cli,
&only_user,
Some("u".into()),
Some("p".into())
)
.is_err()
);
assert!(resolve_user_with(NatsRole::Cli, reg(&[]), Some("u".into()), None).is_err());
assert!(resolve_user_with(NatsRole::Cli, reg(&[]), None, Some("p".into())).is_err());
}
#[test]
fn credentials_never_print_a_secret() {
let creds = NatsCredentials::new(
Some("tok-secret".into()),
Some(("kanade-agent".into(), "pw-secret".into())),
);
let rendered = format!("{creds:?}");
assert!(!rendered.contains("tok-secret"), "{rendered}");
assert!(!rendered.contains("pw-secret"), "{rendered}");
assert!(rendered.contains("kanade-agent"), "{rendered}");
}
#[test]
fn no_user_means_the_static_token_path_with_no_selection() {
match auth_plan(NatsCredentials::new(Some("t".into()), None)) {
AuthPlan::Static(Some(t)) => assert_eq!(t, "t"),
_ => panic!("a token-only host must take the static path"),
}
assert!(matches!(
auth_plan(NatsCredentials::new(None, None)),
AuthPlan::Static(None)
));
assert!(matches!(
auth_plan(NatsCredentials::new(
Some("t".into()),
Some(("u".into(), "p".into()))
)),
AuthPlan::Select { .. }
));
}
#[test]
fn selection_follows_the_probe() {
use Choice::{Token, User};
use ProbeOutcome::*;
assert_eq!(select(Accepted, None, true), Ok(User));
assert_eq!(select(Accepted, Some(Token), false), Ok(User));
assert_eq!(select(Rejected, Some(User), true), Ok(Token));
let missing = select(Rejected, None, false).unwrap_err();
assert!(missing.contains("NatsToken"), "{missing}");
assert!(select(Unreachable, None, true).is_err());
assert!(select(Unreachable, Some(Token), true).is_err());
assert!(select(Unreachable, Some(User), true).is_err());
}
#[tokio::test]
async fn only_a_probe_answer_is_remembered() {
let live = Live::default();
let role = NatsRole::Cli;
let run = |o| choose(role, true, &live, move || async move { o });
assert!(run(ProbeOutcome::Unreachable).await.is_err());
assert_eq!(live.get(), None);
assert_eq!(run(ProbeOutcome::Rejected).await, Ok(Choice::Token));
assert_eq!(live.get(), Some(Choice::Token));
assert!(run(ProbeOutcome::Unreachable).await.is_err());
assert_eq!(run(ProbeOutcome::Accepted).await, Ok(Choice::User));
assert_eq!(live.get(), Some(Choice::User));
let err = choose(role, false, &live, || async { ProbeOutcome::Rejected })
.await
.unwrap_err();
assert!(err.contains("NatsToken"));
assert_eq!(live.get(), Some(Choice::User));
}
#[test]
fn the_probe_reports_the_shape_in_use_right_now() {
let live = Arc::new(Live::default());
let probe = CredentialProbe {
token: Some("tok".into()),
user: Some("kanade-agent".into()),
live: Some(live.clone()),
};
assert_eq!(probe.kind(), CredentialKind::Token);
live.set(Choice::User);
assert_eq!(probe.kind(), CredentialKind::User);
live.set(Choice::Token);
assert_eq!(probe.kind(), CredentialKind::Token);
assert!(probe.is_ours("tok"));
let user_only = CredentialProbe {
token: None,
user: Some("u".into()),
live: Some(Arc::new(Live::default())),
};
assert_eq!(user_only.kind(), CredentialKind::None);
user_only.live.as_ref().unwrap().set(Choice::User);
assert_eq!(user_only.kind(), CredentialKind::User);
}
#[test]
fn a_probe_holding_both_credentials_prints_neither_secret() {
let probe = CredentialProbe {
token: Some("tok-secret".into()),
user: Some("kanade-agent".into()),
live: None,
};
let rendered = format!("{probe:?}");
assert!(!rendered.contains("tok-secret"), "{rendered}");
assert!(rendered.contains("kanade-agent"), "{rendered}");
}
#[test]
fn sustained_refusals_fail_the_client_and_a_connect_clears_them() {
let refused = || {
async_nats::Event::ClientError(async_nats::ClientError::Other(
async_nats::ConnectErrorKind::AuthorizationViolation.to_string(),
))
};
let live = Live::default();
for _ in 0..AUTH_REJECTION_LIMIT {
live.observe(&refused());
}
assert!(!live.auth_failed());
let aged = std::time::Instant::now() - AUTH_REJECTION_WINDOW - Duration::from_secs(1);
*live.refusals.lock().unwrap() = Some((aged, AUTH_REJECTION_LIMIT));
assert!(live.auth_failed());
live.observe(&async_nats::Event::Connected);
assert!(!live.auth_failed());
*live.refusals.lock().unwrap() = Some((aged, AUTH_REJECTION_LIMIT));
live.observe(&async_nats::Event::ClientError(
async_nats::ClientError::Other("io".into()),
));
assert!(!live.auth_failed());
live.observe(&refused());
assert!(!live.auth_failed());
}
#[tokio::test]
async fn deferred_probes_do_not_count_as_credential_refusals() {
let live = Live::default();
for _ in 0..AUTH_REJECTION_LIMIT {
assert!(
choose(NatsRole::Cli, true, &live, || async {
ProbeOutcome::Unreachable
})
.await
.is_err()
);
live.observe(&async_nats::Event::ClientError(
async_nats::ClientError::Other(
async_nats::ConnectErrorKind::Authentication.to_string(),
),
));
}
assert!(live.refusals.lock().unwrap().is_none());
for _ in 0..AUTH_REJECTION_LIMIT {
assert!(
choose(NatsRole::Cli, false, &live, || async {
ProbeOutcome::Rejected
})
.await
.is_err()
);
live.observe(&async_nats::Event::ClientError(
async_nats::ClientError::Other(
async_nats::ConnectErrorKind::Authentication.to_string(),
),
));
}
assert_eq!(
live.refusals.lock().unwrap().as_ref().unwrap().1,
AUTH_REJECTION_LIMIT
);
assert!(!live.auth_failed(), "quick refusals must still allow retry");
assert!(
choose(NatsRole::Cli, true, &live, || async {
ProbeOutcome::Unreachable
})
.await
.is_err()
);
assert!(
live.refusals.lock().unwrap().is_none(),
"an outage interrupts the refusal run"
);
}
#[tokio::test]
async fn same_role_connections_have_independent_health_state() {
let first = connect_with_credentials(
NatsRole::Agent,
"nats://127.0.0.1:1",
NatsCredentials::default(),
)
.await
.unwrap();
let second = connect_with_credentials(
NatsRole::Agent,
"nats://127.0.0.1:1",
NatsCredentials::default(),
)
.await
.unwrap();
let first_health = connection_health(&first).unwrap();
let second_health = connection_health(&second).unwrap();
assert!(!Arc::ptr_eq(&first_health.live, &second_health.live));
assert!(Arc::ptr_eq(
&first_health,
&connection_health(&first.clone()).unwrap()
));
*first_health.live.refusals.lock().unwrap() = Some((
std::time::Instant::now() - AUTH_REJECTION_WINDOW,
AUTH_REJECTION_LIMIT,
));
assert!(first_health.live.auth_failed());
assert!(!second_health.live.auth_failed());
}
#[test]
fn connect_leaves_a_process_crypto_provider_installed() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
connect(NatsRole::Agent, "nats://127.0.0.1:1")
.await
.unwrap();
});
assert!(rustls::crypto::CryptoProvider::get_default().is_some());
}
}