Watch
0
0
Fork
You've already forked hyperhive
0
hyperhive/hive-agent/src/swarm_queue.rs
atlas d5de1e4943 swarm UI: log an agent into claude from its card; the credentials go to bao and the agent picks them up
A new swarm service, `swarm-claude-login`, in its own container on the host
that sets `deploy.claude-login.enable` (defaulted on by `allSwarmServices`
where the host has the minter's store leaf, like the review classifier).

- swarm-ui: an agent card showing "login needed" gets a "log in" button.
  Its dialog starts the login, shows the URL claude printed, and takes the
  pasted code.
- swarm-controller: `POST/DELETE /api/agents/{name}/claude-login` and
  `POST /api/agents/{name}/claude-login/code`, behind the swarm UI's
  authelia-gated `/api/`, relay the steps to the minter as requests on
  `$SWARM.claude-login.request`.
- swarm-claude-login: runs `claude auth login` over piped stdio with HOME
  and CLAUDE_CONFIG_DIR in a per-login scratch dir, scrapes the first
  https URL, writes the code to stdin, and stores the `.credentials.json`
  at `swarm/agents/<agent>/claude-credentials` with a KV v2 write, a new
  version over any earlier one. The scratch dir goes on every exit path; a
  login waits 15 minutes for its code.
- hive-agent: while parked in `wait_for_login`, reads that leaf every 30 s
  under the agent's own cert and installs it only when its KV version is
  newer than the one recorded in `<harness>/claude-credentials-version`.
  An equal or older version is never put back, because claude rotates the
  refresh token and the stored copy is stale after the first refresh.
- swarm-secret-client: `claude_login` paths and `SecretStore::read_versioned`.
- Credentials: a `claude-login` store leaf minted by glue-bao-tls, a
  `swarm-claude-login` policy (create/update on
  `swarm/agents/+/claude-credentials`, read on its own queue client
  secret), an authelia machine client, and a `--claude-login-client`
  grant in the queue's auth-callout responder. The unit runs as a static
  system user.

The hive's own web-UI login is unchanged.

Refs #5268
2026-10-15 13:26:29 +02:00

763 lines
28 KiB
Rust

//! This agent's swarm-queue connection, made with the agent's own credential.
//!
//! The credential is minted per agent at swarm level. This process reads it
//! from the swarm secret store under the agent's own certificate and holds it
//! in memory only ([`AgentSecret`]), and publishes on the agent's own
//! hive-free subjects.
//!
//! **The connection is required.** [`init`] fails when the queue's address or
//! the store's coordinates are missing, and [`client`] retries the first
//! connect on [`CONNECT_BACKOFF`] and then exits the process, so the harness
//! unit fails and systemd restarts it.
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;
use anyhow::{Context as _, anyhow};
use tokio::sync::OnceCell;
use swarm_secret_client::{
SecretStore,
client::{DEFAULT_CERT_MOUNT, ENV_ADDR, ENV_CACERT, Settings},
policy, queue,
};
/// Where the swarm queue listens, as this container reaches it.
const ENV_NATS_URL: &str = "HIVE_AGENT_NATS_URL";
/// Names the agent whose queue secret this process reads. The cert-auth role
/// and the store path are both built from it by `swarm_secret_client`.
const ENV_AGENT_NAME: &str = "HIVE_AGENT_NAME";
/// How long one read of the store may take. The read runs inside a queue
/// connection attempt, and a store that never answers would otherwise stall
/// every reconnect behind it.
const STORE_READ_TIMEOUT: Duration = Duration::from_secs(10);
/// The waits between first-connect attempts: 5s doubling, about five minutes
/// in all. After the last attempt fails the process exits.
const CONNECT_BACKOFF: [Duration; 6] = [
Duration::from_secs(5),
Duration::from_secs(10),
Duration::from_secs(20),
Duration::from_secs(40),
Duration::from_secs(80),
Duration::from_secs(160),
];
/// Where to connect, resolved once at boot by [`init`].
static AGENT: OnceLock<AgentPath> = OnceLock::new();
/// The one connection every publisher in this process shares. Separate from
/// [`AGENT`] because resolving the coordinates is synchronous boot work and
/// connecting is not — see [`client`].
static CLIENT: OnceCell<async_nats::Client> = OnceCell::const_new();
/// What connecting with this agent's own credential needs.
#[derive(Debug, PartialEq, Eq)]
struct AgentPath {
url: String,
store: Arc<StoreSource>,
}
/// This agent's queue secret as the store holds it: where the store is, the
/// role this agent logs in as, and the path it reads.
#[derive(Debug, PartialEq, Eq)]
struct StoreSource {
settings: Settings,
role: String,
path: String,
}
/// Where [`AgentSecret`] reads the secret from. A trait so the retry rules can
/// be exercised without a store.
trait SecretSource: Send + Sync + 'static {
/// The store path, for log lines and errors. Never the value.
fn origin(&self) -> &str;
/// The value stored now, or `None` when nothing is stored.
fn fetch(&self) -> impl Future<Output = anyhow::Result<Option<String>>> + Send;
}
impl SecretSource for Arc<StoreSource> {
fn origin(&self) -> &str {
&self.path
}
/// A fresh login per read: reads happen once per connection attempt, and a
/// token held between them would be one more secret in memory to expire.
async fn fetch(&self) -> anyhow::Result<Option<String>> {
let store = SecretStore::connect(&self.settings, &self.role, DEFAULT_CERT_MOUNT)
.await
.context("logging in to the swarm secret store as this agent")?;
// `read_optional`: only a 404 is "nothing stored". A 403 stays an
// error, so a policy that stopped covering the path is not reported
// as a secret that was deleted.
let credential: Option<queue::AgentCredential> = store
.read_optional(&self.path)
.await
.with_context(|| format!("reading {} from the store", self.path))?;
Ok(credential.map(|c| c.value))
}
}
/// This agent's queue secret, held in memory between connection attempts.
///
/// No `Debug`: [`Held`] carries the secret.
struct AgentSecret<S> {
source: S,
timeout: Duration,
held: Mutex<Held>,
}
#[derive(Default)]
struct Held {
value: Option<String>,
/// `value` was read by [`AgentSecret::prime`] and no attempt has
/// presented it yet.
unpresented: bool,
}
impl<S: SecretSource> AgentSecret<S> {
fn new(source: S) -> Self {
Self {
source,
timeout: STORE_READ_TIMEOUT,
held: Mutex::default(),
}
}
fn held(&self) -> std::sync::MutexGuard<'_, Held> {
// A poisoned lock only means another attempt panicked mid-update; the
// worst it left is a stale value, which the next read replaces.
self.held
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
/// Read the secret before the first connection attempt. An `Err` here
/// fails that attempt.
async fn prime(&self) -> anyhow::Result<()> {
let value = self
.read()
.await?
.with_context(|| format!("nothing is stored at {}", self.source.origin()))?;
*self.held() = Held {
value: Some(value),
unpresented: true,
};
Ok(())
}
/// The secret to present on one connection attempt.
///
/// The value [`prime`](Self::prime) read serves the first attempt. Every
/// later attempt reads the store again, so after the queue refuses a
/// re-minted secret the next attempt presents the current one. A store
/// that cannot be read leaves the held value in use; a store that holds
/// nothing drops it. Either way the attempt is one read, and the next one
/// waits out the connection's reconnect backoff.
async fn for_attempt(&self) -> anyhow::Result<String> {
{
let mut held = self.held();
if std::mem::take(&mut held.unpresented)
&& let Some(value) = &held.value
{
return Ok(value.clone());
}
}
match self.read().await {
Ok(Some(value)) => {
self.held().value = Some(value.clone());
Ok(value)
}
Ok(None) => {
self.held().value = None;
tracing::warn!(
path = %self.source.origin(),
"nothing is stored at this agent's queue credential path; its queue \
connection retries under backoff"
);
Err(anyhow!("nothing is stored at {}", self.source.origin()))
}
Err(e) => {
let held = self.held().value.clone();
tracing::warn!(
error = format!("{e:#}"),
presenting_held = held.is_some(),
"re-reading this agent's queue credential from the store failed"
);
held.ok_or(e)
}
}
}
/// One bounded read. An empty stored value is no value.
async fn read(&self) -> anyhow::Result<Option<String>> {
let value = tokio::time::timeout(self.timeout, self.source.fetch())
.await
.map_err(|_| anyhow!("the store did not answer within {:?}", self.timeout))??;
Ok(value.map(|v| v.trim().to_owned()).filter(|v| !v.is_empty()))
}
}
/// Where this agent reads its own queue secret from.
///
/// Taken through a lookup so it can be asserted without mutating the process
/// environment. The shape follows `hive_matrix_mcp::credential`'s
/// `store_settings`, which reads the same variables for the same identity.
///
/// # Errors
/// No store address, or one without the agent name or the identity beside it.
fn store_source(get: impl Fn(&str) -> Option<String>) -> anyhow::Result<StoreSource> {
let (settings, role, agent) = agent_store(get)?;
Ok(StoreSource {
settings,
role,
path: queue::agent_queue_path(&agent)?,
})
}
/// Where the store is, the cert-auth role this agent logs in as, and the
/// agent's name, from `get`.
///
/// # Errors
/// No store address, or one without the agent name or the identity beside it.
pub(crate) fn agent_store(
get: impl Fn(&str) -> Option<String>,
) -> anyhow::Result<(Settings, String, String)> {
if get(ENV_ADDR).is_none_or(|v| v.is_empty()) {
anyhow::bail!(
"{ENV_ADDR} is unset or empty, so there is no secret store to read this agent's \
secrets from"
);
}
let agent = get(ENV_AGENT_NAME)
.filter(|v| !v.is_empty())
.with_context(|| format!("{ENV_AGENT_NAME} is unset or empty"))?;
let settings = Settings::from_lookup(|k| get(k).filter(|v| k != ENV_CACERT || ca_is_usable(v)))
.context("reading the swarm secret store's coordinates from the environment")?;
let role = policy::agent_object_name(&agent)?;
Ok((settings, role, agent))
}
/// Whether the CA bundle at `path` is a file with bytes in it. `BAO_CACERT`
/// names a systemd credential that may not have been delivered; absent means
/// the container's own trust store.
fn ca_is_usable(path: &str) -> bool {
std::fs::metadata(path).is_ok_and(|m| m.len() > 0)
}
/// What connecting needs: the queue's address and a store to read this agent's
/// secret from.
///
/// # Errors
/// Either one missing.
fn agent_path(get: impl Fn(&str) -> Option<String>) -> anyhow::Result<AgentPath> {
let url = get(ENV_NATS_URL)
.filter(|v| !v.is_empty())
.with_context(|| format!("{ENV_NATS_URL} is unset or empty"))?;
Ok(AgentPath {
url,
store: Arc::new(store_source(&get)?),
})
}
/// Resolve this agent's queue coordinates and record them for [`client`].
///
/// # Errors
/// A coordinate is missing. The harness exits on it: an agent without a queue
/// connection is invisible to the operator at swarm level.
pub fn init() -> anyhow::Result<()> {
let agent = agent_path(|k| std::env::var(k).ok())
.context("resolving this agent's swarm queue coordinates")?;
tracing::info!(
url = %agent.url,
path = %agent.store.path,
"this agent connects to the swarm queue with its own credential from the secret store"
);
let _ = AGENT.set(agent);
Ok(())
}
/// The shared queue connection, made on first call and memoized for the rest
/// of the process.
///
/// One connection per process, not per publisher: the agent authenticates as
/// one client, so a second `connect` would be a second live connection for the
/// same identity rather than a second credential.
///
/// The first call retries the connect on [`CONNECT_BACKOFF`], and every other
/// caller waits on it. When no attempt succeeds the process exits with status
/// 1, so this never returns without a client.
pub async fn client() -> async_nats::Client {
CLIENT.get_or_init(connect_or_exit).await.clone()
}
async fn connect_or_exit() -> async_nats::Client {
let connected = match AGENT.get() {
Some(agent) => retry_first_connect(|| connect_once(agent), tokio::time::sleep).await,
None => Err(anyhow!("swarm_queue::init has not run")),
};
connected.unwrap_or_else(|e| {
tracing::error!(
error = format!("{e:#}"),
"no swarm queue connection; exiting, since this agent cannot run without one"
);
std::process::exit(1)
})
}
/// One connection attempt with this agent's own credential.
async fn connect_once(agent: &AgentPath) -> anyhow::Result<async_nats::Client> {
let secret = AgentSecret::new(Arc::clone(&agent.store));
let name = crate::identity::label();
let client = connect_as_agent(&agent.url, name, secret).await?;
tracing::info!(
url = %agent.url,
"connected to the swarm queue with this agent's own credential"
);
Ok(client)
}
/// Run `connect` until it succeeds, sleeping each wait in [`CONNECT_BACKOFF`]
/// between attempts, and return the last error once every attempt has failed.
async fn retry_first_connect<T, C, CFut, S, SFut>(mut connect: C, mut sleep: S) -> anyhow::Result<T>
where
C: FnMut() -> CFut,
CFut: Future<Output = anyhow::Result<T>>,
S: FnMut(Duration) -> SFut,
SFut: Future<Output = ()>,
{
let mut waits = CONNECT_BACKOFF.iter();
let mut attempt = 1;
loop {
let e = match connect().await {
Ok(client) => return Ok(client),
Err(e) => e,
};
let Some(&wait) = waits.next() else {
return Err(e.context(format!("{attempt} connection attempts failed")));
};
tracing::warn!(
error = format!("{e:#}"),
attempt,
retry_in = ?wait,
"connecting to the swarm queue failed"
);
sleep(wait).await;
attempt += 1;
}
}
/// Connect presenting this agent's own token as `name`.
///
/// The secret is read before the first attempt, and any failure of that read
/// or of the first attempt, a refusal included, is returned for
/// [`retry_first_connect`] to retry. Once connected, each reconnect attempt
/// takes its secret from [`AgentSecret::for_attempt`].
async fn connect_as_agent<S: SecretSource>(
url: &str,
name: String,
secret: AgentSecret<S>,
) -> anyhow::Result<async_nats::Client> {
secret.prime().await?;
let secret = Arc::new(secret);
let agent = name.clone();
// No extra trust anchor: the container trusts the store's root, which
// `hive_c0re::meta` embeds at build time.
swarm_queue_client::connect_with_token(url, None, &agent, move || {
let secret = Arc::clone(&secret);
let name = name.clone();
// Spawned so the future handed back is only a `JoinHandle`, which is
// `Sync` as async-nats's auth callback requires; the store read is not.
let attempt = tokio::spawn(async move {
let value = secret.for_attempt().await.map_err(|e| format!("{e:#}"))?;
swarm_queue_client::agent_token::format_agent_token(&name, &value)
.map_err(|e| e.to_string())
});
async move { attempt.await.map_err(|e| e.to_string())? }
})
.await
.map_err(|e| anyhow!(swarm_queue_client::chain(&e)))
}
#[cfg(test)]
mod tests {
use std::collections::VecDeque;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use anyhow::anyhow;
use super::{
AgentPath, AgentSecret, CONNECT_BACKOFF, SecretSource, StoreSource, agent_path,
connect_as_agent, retry_first_connect, store_source,
};
/// A lookup standing in for a container that was handed a store and the
/// queue's address.
fn with_store(k: &str) -> Option<String> {
match k {
"BAO_ADDR" => Some("https://bao.t.local:8200".to_owned()),
"BAO_CLIENT_CERT" => Some("/run/credentials/x/cert".to_owned()),
"BAO_CLIENT_KEY" => Some("/run/credentials/x/key".to_owned()),
"HIVE_AGENT_NAME" => Some("a1".to_owned()),
"HIVE_AGENT_NATS_URL" => Some("nats://10.42.0.1:4222".to_owned()),
_ => None,
}
}
fn source() -> StoreSource {
store_source(with_store).expect("a complete store environment")
}
#[test]
fn a_store_address_reads_the_agents_own_path_as_its_own_role() {
let s = source();
assert_eq!(s.path, "swarm/agents/a1/queue");
assert_eq!(s.role, "hive-agent-a1");
}
#[test]
fn a_missing_store_address_is_an_error() {
for addr in [None, Some(String::new())] {
let e = store_source(|k| {
if k == "BAO_ADDR" {
addr.clone()
} else {
with_store(k)
}
})
.expect_err("no store to read the credential from");
assert!(format!("{e:#}").contains("BAO_ADDR"), "{e:#}");
}
}
#[test]
fn a_store_address_without_the_agent_or_its_identity_is_an_error() {
for var in ["HIVE_AGENT_NAME", "BAO_CLIENT_CERT", "BAO_CLIENT_KEY"] {
let e = store_source(|k| if k == var { None } else { with_store(k) })
.expect_err("a half-delivered store");
assert!(format!("{e:#}").contains(var), "dropping {var}: {e:#}");
}
}
/// `BAO_CACERT` is always named and may point at nothing; that is the
/// container's own trust store, not a CA path that fails the handshake.
#[test]
fn an_undelivered_ca_is_dropped() {
let with_missing_ca = store_source(|k| {
if k == "BAO_CACERT" {
Some("/nonexistent/ca".to_owned())
} else {
with_store(k)
}
})
.expect("complete");
assert_eq!(with_missing_ca, source());
}
#[test]
fn a_store_and_the_queue_address_are_all_an_agent_connects_with() {
assert_eq!(
agent_path(with_store).expect("complete"),
AgentPath {
url: "nats://10.42.0.1:4222".to_owned(),
store: Arc::new(source()),
}
);
}
#[test]
fn without_a_store_or_the_queue_address_an_agent_has_nothing_to_connect_with() {
for var in ["HIVE_AGENT_NATS_URL", "BAO_ADDR"] {
let e = agent_path(|k| if k == var { None } else { with_store(k) })
.expect_err("a coordinate is missing");
assert!(format!("{e:#}").contains(var), "dropping {var}: {e:#}");
}
}
/// Runs `retry_first_connect` against a connect that fails `failures`
/// times and then succeeds, recording each wait instead of sleeping it.
async fn retry_with(failures: usize) -> (anyhow::Result<usize>, usize, Vec<Duration>) {
let mut calls = 0;
let mut waits = Vec::new();
let result = retry_first_connect(
|| {
calls += 1;
let n = calls;
async move {
if n > failures {
Ok(n)
} else {
Err(anyhow!("attempt {n} refused"))
}
}
},
|wait| {
waits.push(wait);
async {}
},
)
.await;
(result, calls, waits)
}
#[tokio::test]
async fn a_first_connect_that_succeeds_waits_for_nothing() {
let (result, calls, waits) = retry_with(0).await;
assert_eq!(result.expect("connected"), 1);
assert_eq!(calls, 1);
assert!(waits.is_empty(), "{waits:?}");
}
/// 5s doubling, one wait per failed attempt, and the attempt after the
/// last wait is the one that connects.
#[tokio::test]
async fn a_failed_connect_is_retried_on_a_doubling_backoff() {
let (result, calls, waits) = retry_with(3).await;
assert_eq!(result.expect("connected on the fourth attempt"), 4);
assert_eq!(calls, 4);
assert_eq!(
waits,
[5, 10, 20].map(Duration::from_secs),
"each wait doubles the one before"
);
}
/// About five minutes of waits, then the last attempt's error rather than
/// another wait.
#[tokio::test]
async fn the_retry_gives_up_with_the_last_error_after_about_five_minutes() {
let (result, calls, waits) = retry_with(usize::MAX).await;
let e = result.expect_err("every attempt failed");
assert_eq!(calls, CONNECT_BACKOFF.len() + 1);
assert_eq!(waits, CONNECT_BACKOFF);
let total: Duration = waits.iter().sum();
assert!(
(Duration::from_mins(4)..=Duration::from_mins(6)).contains(&total),
"{total:?}"
);
assert!(
format!("{e:#}").contains(&format!("attempt {calls} refused")),
"{e:#}"
);
}
enum Answer {
Value(&'static str),
Absent,
Fails,
Hangs,
}
/// A store that answers each read with the next scripted answer and
/// counts the reads.
struct Scripted {
answers: Mutex<VecDeque<Answer>>,
reads: AtomicUsize,
}
fn scripted(answers: impl IntoIterator<Item = Answer>) -> Arc<Scripted> {
Arc::new(Scripted {
answers: Mutex::new(answers.into_iter().collect()),
reads: AtomicUsize::new(0),
})
}
impl SecretSource for Arc<Scripted> {
fn origin(&self) -> &'static str {
"swarm/agents/a1/queue"
}
async fn fetch(&self) -> anyhow::Result<Option<String>> {
self.reads.fetch_add(1, Ordering::SeqCst);
let next = self.answers.lock().expect("lock").pop_front();
match next.expect("a read nobody scripted") {
Answer::Value(v) => Ok(Some(v.to_owned())),
Answer::Absent => Ok(None),
Answer::Fails => Err(anyhow!("the store is unreachable")),
Answer::Hangs => std::future::pending().await,
}
}
}
fn reads(s: &Arc<Scripted>) -> usize {
s.reads.load(Ordering::SeqCst)
}
/// (a) The first attempt presents what was just read, and does not read
/// the store a second time for it.
#[tokio::test]
async fn the_first_attempt_presents_the_value_read_before_it() {
let store = scripted([Answer::Value("v1\n")]);
let secret = AgentSecret::new(Arc::clone(&store));
secret.prime().await.expect("the store holds a secret");
assert_eq!(secret.for_attempt().await.expect("primed"), "v1");
assert_eq!(reads(&store), 1);
}
/// (b) After the queue refuses a re-minted secret, the next attempt reads
/// the store and presents what it holds now: one read per attempt.
#[tokio::test]
async fn each_later_attempt_presents_what_the_store_holds_now() {
let store = scripted([
Answer::Value("v1"),
Answer::Value("v1"),
Answer::Value("v2"),
]);
let secret = AgentSecret::new(Arc::clone(&store));
secret.prime().await.expect("the store holds a secret");
assert_eq!(secret.for_attempt().await.expect("primed"), "v1");
assert_eq!(secret.for_attempt().await.expect("re-read"), "v1");
assert_eq!(secret.for_attempt().await.expect("re-read"), "v2");
assert_eq!(reads(&store), 3);
}
/// (b) Bounded: a store that never answers costs one attempt the read
/// timeout, and the attempt goes ahead with the value it already holds.
#[tokio::test]
async fn a_store_that_does_not_answer_is_given_up_on() {
let store = scripted([Answer::Value("v1"), Answer::Hangs, Answer::Hangs]);
let mut secret = AgentSecret::new(Arc::clone(&store));
secret.timeout = Duration::from_millis(20);
secret.prime().await.expect("the store holds a secret");
assert_eq!(secret.for_attempt().await.expect("primed"), "v1");
assert_eq!(
secret
.for_attempt()
.await
.expect("the held value stands in"),
"v1"
);
assert_eq!(reads(&store), 2);
let cold = AgentSecret::new(scripted([Answer::Hangs]));
let mut cold = cold;
cold.timeout = Duration::from_millis(20);
let e = cold.prime().await.expect_err("nothing to stand in");
assert!(format!("{e:#}").contains("did not answer"), "{e:#}");
}
/// (b) A secret the store no longer holds (the agent was destroyed) is not
/// presented again, even when a later read fails.
#[tokio::test]
async fn a_secret_removed_from_the_store_is_dropped() {
let store = scripted([Answer::Value("v1"), Answer::Absent, Answer::Fails]);
let secret = AgentSecret::new(Arc::clone(&store));
secret.prime().await.expect("the store holds a secret");
assert_eq!(secret.for_attempt().await.expect("primed"), "v1");
let gone = secret.for_attempt().await.expect_err("nothing stored");
assert!(
format!("{gone:#}").contains("swarm/agents/a1/queue"),
"{gone:#}"
);
secret
.for_attempt()
.await
.expect_err("the dropped value does not come back");
}
/// (c) Nothing to connect with before the first attempt fails that attempt.
/// The queue is never dialled: nothing listens at the address given.
#[tokio::test]
async fn no_secret_before_the_first_connect_fails_the_attempt() {
for (answer, says) in [
(Answer::Fails, "unreachable"),
(Answer::Absent, "nothing is stored"),
(Answer::Value(" \n"), "nothing is stored"),
] {
let secret = AgentSecret::new(scripted([answer]));
let e = connect_as_agent("nats://127.0.0.1:9", "a1".to_owned(), secret)
.await
.expect_err("no secret to present");
assert!(format!("{e:#}").contains(says), "{e:#}");
}
}
/// Serve one NATS handshake on `sock`: send INFO, read CONNECT up to the
/// PING, and return the token presented.
async fn handshake(sock: &mut tokio::net::TcpStream, port: u16) -> String {
use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _};
let info = format!(
"INFO {{\"server_id\":\"fake\",\"version\":\"2.10.0\",\"go\":\"go\",\
\"host\":\"127.0.0.1\",\"port\":{port},\"headers\":true,\
\"max_payload\":1048576,\"proto\":1,\"auth_required\":true,\
\"nonce\":\"n0nce\"}}\r\n"
);
sock.write_all(info.as_bytes()).await.expect("send INFO");
let mut lines = tokio::io::BufReader::new(sock).lines();
let mut token = String::new();
while let Some(line) = lines.next_line().await.expect("read") {
if let Some(json) = line.strip_prefix("CONNECT ") {
let v: serde_json::Value = serde_json::from_str(json).expect("CONNECT is JSON");
token = v["auth_token"].as_str().unwrap_or_default().to_owned();
} else if line.starts_with("PING") {
break;
}
}
token
}
/// (b) end to end through async-nats: connected with `v1`, the connection
/// drops, the reconnect presenting `v1` is refused, and the next attempt
/// reads `v2` and is accepted.
#[tokio::test]
async fn a_refused_reconnect_reconnects_with_the_secret_read_after_it() {
use tokio::io::AsyncWriteExt as _;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind");
let port = listener.local_addr().expect("addr").port();
let (seen_tx, mut seen_rx) = tokio::sync::mpsc::unbounded_channel();
tokio::spawn(async move {
let mut kept = Vec::new();
for n in 0.. {
let (mut sock, _) = listener.accept().await.expect("accept");
let token = handshake(&mut sock, port).await;
let accept = n == 0 || token.rsplit_once('.').is_some_and(|(_, s)| s == "v2");
let reply: &[u8] = if accept {
b"PONG\r\n"
} else {
b"-ERR 'Authorization Violation'\r\n"
};
let _ = sock.write_all(reply).await;
let _ = seen_tx.send(token);
// The first connection is dropped to force a reconnect; the
// one accepted after the refusal stays open.
if accept && n > 0 {
kept.push(sock);
}
}
});
let store = scripted([
Answer::Value("v1"),
Answer::Value("v1"),
Answer::Value("v2"),
]);
let secret = AgentSecret::new(Arc::clone(&store));
let url = format!("nats://127.0.0.1:{port}");
let _client = connect_as_agent(&url, "a1".to_owned(), secret)
.await
.expect("the first connect is accepted");
let mut seen = Vec::new();
while seen.len() < 3 {
let token = tokio::time::timeout(Duration::from_secs(20), seen_rx.recv())
.await
.expect("the client keeps reconnecting")
.expect("the server is running");
seen.push(token);
}
let secrets: Vec<&str> = seen
.iter()
.map(|t| t.rsplit_once('.').map_or("", |(_, s)| s))
.collect();
assert_eq!(secrets, ["v1", "v1", "v2"]);
assert_eq!(reads(&store), 3);
}
}