swarm-nats-auth: answer callout requests concurrently, skip expired ones
Each request now runs in its own task, at most CONCURRENT_REQUESTS (32) at once, instead of being awaited inline by the subscription loop. A request whose exp (nats-server sets it to send time plus its authorization.timeout) has passed is skipped without a lookup or reply. Fixes #5434.
This commit is contained in:
atlas
mara
parent
bf848ba65f
commit
6971b4e4ed
6 changed files with 334 additions and 48 deletions
|
|
@ -30,6 +30,8 @@
|
|||
//! from the secret store, at the paths `xkey`, `nkey` and `oidc_secret` name.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use anyhow::Context;
|
||||
use clap::{Parser, Subcommand};
|
||||
|
|
@ -48,6 +50,13 @@ mod xkey;
|
|||
/// Subject the NATS server publishes authorization requests on.
|
||||
const AUTH_SUBJECT: &str = "$SYS.REQ.USER.AUTH";
|
||||
|
||||
/// Authorization requests answered at once. Each holds at most one store
|
||||
/// lookup or one introspection, for at most
|
||||
/// [`introspect::INTROSPECTION_TIMEOUT`]: with every lookup at that bound,
|
||||
/// 32 answer about 21 requests a second, and a flood of connects keeps at
|
||||
/// most 32 store logins and `IdP` calls open.
|
||||
const CONCURRENT_REQUESTS: usize = 32;
|
||||
|
||||
/// With no verb this is the responder, and its flags are required.
|
||||
#[derive(Debug, Parser)]
|
||||
#[command(
|
||||
|
|
@ -211,14 +220,14 @@ async fn main() -> anyhow::Result<()> {
|
|||
)
|
||||
.await
|
||||
}
|
||||
(None, Some(args)) => serve(&args).await,
|
||||
(None, Some(args)) => serve(args).await,
|
||||
(None, None) => unreachable!("clap requires the responder's flags when no verb is given"),
|
||||
}
|
||||
}
|
||||
|
||||
/// The responder: answers every authorization request until its
|
||||
/// subscription ends.
|
||||
async fn serve(args: &Args) -> anyhow::Result<()> {
|
||||
async fn serve(args: Args) -> anyhow::Result<()> {
|
||||
// The bucket is NOT a flag. `swarm_queue_client::status`'s own docs say
|
||||
// why: reader and writer must name the same bucket, and an option is a
|
||||
// way for two deployments to disagree about which one that is. This
|
||||
|
|
@ -253,14 +262,14 @@ async fn serve(args: &Args) -> anyhow::Result<()> {
|
|||
.await
|
||||
.context("logging in to the secret store for the callout keys and the OIDC client secret")?;
|
||||
let (callout_xkey, issuer, user_seed) = load_callout_keys(&keys).await?;
|
||||
let client_secret = load_client_secret(&keys, args).await?;
|
||||
let client_secret = load_client_secret(&keys, &args).await?;
|
||||
let store = swarm_agent_verify::Store::new(settings, args.store_cert_role.clone());
|
||||
let client = async_nats::ConnectOptions::with_nkey(user_seed)
|
||||
.name("swarm-nats-auth")
|
||||
.connect(&args.nats_url)
|
||||
.await
|
||||
.with_context(|| format!("connect to {}", args.nats_url))?;
|
||||
let mut requests = client
|
||||
let requests = client
|
||||
.subscribe(AUTH_SUBJECT)
|
||||
.await
|
||||
.with_context(|| format!("subscribe to {AUTH_SUBJECT}"))?;
|
||||
|
|
@ -270,10 +279,97 @@ async fn serve(args: &Args) -> anyhow::Result<()> {
|
|||
"swarm-nats-auth: connected, awaiting authorization requests"
|
||||
);
|
||||
|
||||
while let Some(msg) = requests.next().await {
|
||||
let Some((server_key, req)) = open_request(&callout_xkey, &msg) else {
|
||||
continue;
|
||||
let responder = Arc::new(Responder {
|
||||
args,
|
||||
policy,
|
||||
http,
|
||||
client_secret,
|
||||
store,
|
||||
callout_xkey,
|
||||
issuer,
|
||||
});
|
||||
answer_each(requests, CONCURRENT_REQUESTS, move |msg| {
|
||||
let responder = Arc::clone(&responder);
|
||||
let client = client.clone();
|
||||
async move {
|
||||
let Some((reply_to, reply)) = responder.answer(&msg).await else {
|
||||
return;
|
||||
};
|
||||
if let Err(e) = client.publish(reply_to, reply.into()).await {
|
||||
tracing::warn!(error = ?e, "failed to publish auth response");
|
||||
}
|
||||
}
|
||||
})
|
||||
.await;
|
||||
anyhow::bail!("subscription to {AUTH_SUBJECT} ended")
|
||||
}
|
||||
|
||||
/// Runs `answer` on each of `requests`, at most `permits` at once. A request
|
||||
/// waits in `requests` until a permit is free; the reply carries `server_id`
|
||||
/// and `user_nkey`, so the order answers go out in does not matter.
|
||||
async fn answer_each<R, F, Fut>(
|
||||
mut requests: impl futures_util::Stream<Item = R> + Unpin,
|
||||
permits: usize,
|
||||
answer: F,
|
||||
) where
|
||||
F: Fn(R) -> Fut,
|
||||
Fut: std::future::Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
let permits = Arc::new(tokio::sync::Semaphore::new(permits));
|
||||
while let Some(request) = requests.next().await {
|
||||
let Ok(permit) = Arc::clone(&permits).acquire_owned().await else {
|
||||
unreachable!("the semaphore is never closed");
|
||||
};
|
||||
let answering = answer(request);
|
||||
tokio::spawn(async move {
|
||||
answering.await;
|
||||
drop(permit);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether the server has stopped waiting for the answer to a request whose
|
||||
/// `exp` is `expires`. The server waits `authorization.timeout`
|
||||
/// (`swarm-nats.nix`) and rounds `exp` down to the second, so it has given
|
||||
/// up once the second after `exp` begins. Without `exp`, the request is
|
||||
/// answered.
|
||||
fn expired(expires: Option<i64>, now: SystemTime) -> bool {
|
||||
let Some(expires) = expires else {
|
||||
return false;
|
||||
};
|
||||
now.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.and_then(|since| i64::try_from(since.as_secs()).ok())
|
||||
.is_some_and(|now| now > expires)
|
||||
}
|
||||
|
||||
/// What answering a request needs, shared by every request in flight.
|
||||
struct Responder<S> {
|
||||
args: Args,
|
||||
policy: policy::Policy,
|
||||
http: reqwest::Client,
|
||||
client_secret: oidc_secret::OidcSecret,
|
||||
/// Where an agent token's secret is looked up.
|
||||
store: S,
|
||||
callout_xkey: nkeys::XKey,
|
||||
issuer: nkeys::KeyPair,
|
||||
}
|
||||
|
||||
impl<S: swarm_agent_verify::CredentialSource> Responder<S> {
|
||||
/// The sealed reply to `msg` and the subject it goes to, or `None` for a
|
||||
/// request that gets no reply. A request the server has stopped waiting
|
||||
/// for gets none, and no lookup.
|
||||
async fn answer(&self, msg: &async_nats::Message) -> Option<(async_nats::Subject, Vec<u8>)> {
|
||||
let (server_key, claims) = open_request(&self.callout_xkey, msg)?;
|
||||
let req = claims.nats;
|
||||
if expired(claims.exp, SystemTime::now()) {
|
||||
tracing::warn!(
|
||||
user_nkey = %req.user_nkey,
|
||||
server_id = %req.server_id.id,
|
||||
"auth request expired before it was answered; skipping"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
// No token is a denial, not an error: an anonymous connect is a
|
||||
// normal thing for a client to attempt and an abnormal thing to
|
||||
// grant. Neither the store nor introspection is reached until
|
||||
|
|
@ -290,7 +386,7 @@ async fn serve(args: &Args) -> anyhow::Result<()> {
|
|||
// the secret proved it; `granted` says which.
|
||||
Some(agent_token::Presented::Agent(token)) => (
|
||||
Some(format!("agent:{}", token.agent)),
|
||||
agent_token::authorize(&policy, &store, &token).await,
|
||||
agent_token::authorize(&self.policy, &self.store, &token).await,
|
||||
None,
|
||||
),
|
||||
Some(agent_token::Presented::Malformed(e)) => {
|
||||
|
|
@ -298,7 +394,14 @@ async fn serve(args: &Args) -> anyhow::Result<()> {
|
|||
(None, None, None)
|
||||
}
|
||||
Some(agent_token::Presented::Bearer(token)) => {
|
||||
introspected(&policy, &http, args, &client_secret, token).await
|
||||
introspected(
|
||||
&self.policy,
|
||||
&self.http,
|
||||
&self.args,
|
||||
&self.client_secret,
|
||||
token,
|
||||
)
|
||||
.await
|
||||
}
|
||||
None => (None, None, None),
|
||||
};
|
||||
|
|
@ -320,31 +423,27 @@ async fn serve(args: &Args) -> anyhow::Result<()> {
|
|||
// looks like ordinary denials.
|
||||
let Some(reply_to) = msg.reply.clone() else {
|
||||
tracing::warn!("auth request had no reply subject; dropping");
|
||||
continue;
|
||||
return None;
|
||||
};
|
||||
let token = match &permissions {
|
||||
Some(permissions) => respond::grant(
|
||||
&issuer,
|
||||
&args.account,
|
||||
&self.issuer,
|
||||
&self.args.account,
|
||||
&req.server_id.id,
|
||||
&req.user_nkey,
|
||||
permissions,
|
||||
expires,
|
||||
),
|
||||
None => respond::deny(&issuer, &req.server_id.id, &req.user_nkey),
|
||||
None => respond::deny(&self.issuer, &req.server_id.id, &req.user_nkey),
|
||||
};
|
||||
let reply = match xkey::seal(&callout_xkey, &server_key, token.as_bytes()) {
|
||||
Ok(reply) => reply,
|
||||
match xkey::seal(&self.callout_xkey, &server_key, token.as_bytes()) {
|
||||
Ok(reply) => Some((reply_to, reply)),
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "could not seal the auth response; dropping");
|
||||
continue;
|
||||
None
|
||||
}
|
||||
};
|
||||
if let Err(e) = client.publish(reply_to, reply.into()).await {
|
||||
tracing::warn!(error = ?e, "failed to publish auth response");
|
||||
}
|
||||
}
|
||||
anyhow::bail!("subscription to {AUTH_SUBJECT} ended")
|
||||
}
|
||||
|
||||
/// The request in `msg`, opened, and the server key its reply is sealed to.
|
||||
|
|
@ -357,7 +456,7 @@ async fn serve(args: &Args) -> anyhow::Result<()> {
|
|||
fn open_request(
|
||||
callout_xkey: &nkeys::XKey,
|
||||
msg: &async_nats::Message,
|
||||
) -> Option<(nkeys::XKey, request::AuthRequest)> {
|
||||
) -> Option<(nkeys::XKey, request::Claims)> {
|
||||
let header = msg
|
||||
.headers
|
||||
.as_ref()
|
||||
|
|
@ -568,6 +667,185 @@ mod tests {
|
|||
assert!(open_request(&own, &message(Some(headers), sealed)).is_some());
|
||||
}
|
||||
|
||||
/// How long [`SlowStore`] takes to answer, under the introspection bound.
|
||||
const LOOKUP: std::time::Duration = std::time::Duration::from_secs(1);
|
||||
|
||||
/// A store holding atlas's credential that takes [`LOOKUP`] to answer,
|
||||
/// counting its lookups.
|
||||
#[derive(Default)]
|
||||
struct SlowStore {
|
||||
lookups: std::sync::atomic::AtomicUsize,
|
||||
}
|
||||
|
||||
impl swarm_agent_verify::CredentialSource for SlowStore {
|
||||
async fn lookup(
|
||||
&self,
|
||||
agent: &str,
|
||||
) -> anyhow::Result<Option<swarm_secret_client::queue::AgentCredential>> {
|
||||
self.lookups
|
||||
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
tokio::time::sleep(LOOKUP).await;
|
||||
Ok(
|
||||
(agent == "atlas").then(|| swarm_secret_client::queue::AgentCredential {
|
||||
value: "test-fixture-token".to_owned(),
|
||||
agent: "atlas".to_owned(),
|
||||
minted_at: None,
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// A responder over a [`SlowStore`] that grants atlas's agent token.
|
||||
fn slow_responder() -> Responder<SlowStore> {
|
||||
let args = responder([
|
||||
"swarm-nats-auth",
|
||||
"--introspection-url",
|
||||
"https://auth.example/api/oidc/introspection",
|
||||
])
|
||||
.expect("the required flags alone must parse");
|
||||
let policy = policy::Policy::new(
|
||||
"hive-".to_owned(),
|
||||
"-agent".to_owned(),
|
||||
"hive-status".to_owned(),
|
||||
vec![],
|
||||
vec![],
|
||||
)
|
||||
.expect("valid")
|
||||
.with_agent_token_subjects(vec!["$SWARM.term.{agent}".to_owned()])
|
||||
.expect("valid");
|
||||
let client_secret =
|
||||
oidc_secret::OidcSecret::new(swarm_queue_client::stored_secret::StoredSecret::new(
|
||||
"swarm-nats",
|
||||
"unused".to_owned(),
|
||||
|| async { Err("never re-read".into()) },
|
||||
));
|
||||
Responder {
|
||||
args,
|
||||
policy,
|
||||
http: reqwest::Client::new(),
|
||||
client_secret,
|
||||
store: SlowStore::default(),
|
||||
callout_xkey: nkeys::XKey::new(),
|
||||
issuer: nkeys::KeyPair::new_account(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Now plus `offset` seconds, in whole Unix seconds.
|
||||
fn unix_now_plus(offset: i64) -> i64 {
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("after the epoch")
|
||||
.as_secs();
|
||||
i64::try_from(now).expect("fits") + offset
|
||||
}
|
||||
|
||||
/// Atlas's agent-token request with `exp`, sealed to `own` as the server
|
||||
/// seals it.
|
||||
fn agent_request(own: &nkeys::XKey, exp: i64) -> async_nats::Message {
|
||||
let b = |s: &str| data_encoding::BASE64URL_NOPAD.encode(s.as_bytes());
|
||||
let claims = format!(
|
||||
r#"{{"exp":{exp},"nats":{{"connect_opts":{{"auth_token":"swarm-agent.atlas.test-fixture-token"}},"user_nkey":"UCLIENT","server_id":{{"id":"NSERVER"}}}}}}"#
|
||||
);
|
||||
let jwt = format!("{}.{}.{}", b("{}"), b(&claims), b("sig"));
|
||||
let server = nkeys::XKey::new();
|
||||
let own_public = nkeys::XKey::from_public_key(&own.public_key()).expect("public");
|
||||
let mut headers = async_nats::HeaderMap::new();
|
||||
headers.insert(xkey::SERVER_XKEY_HEADER, server.public_key());
|
||||
async_nats::Message {
|
||||
subject: AUTH_SUBJECT.into(),
|
||||
reply: Some("_INBOX.reply".into()),
|
||||
payload: server
|
||||
.seal(jwt.as_bytes(), &own_public)
|
||||
.expect("sealed")
|
||||
.into(),
|
||||
headers: Some(headers),
|
||||
status: None,
|
||||
description: None,
|
||||
length: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// `n` live requests through [`answer_each`] at [`CONCURRENT_REQUESTS`]
|
||||
/// permits: how long until every one has its reply.
|
||||
async fn time_to_answer(n: usize) -> tokio::time::Duration {
|
||||
let responder = Arc::new(slow_responder());
|
||||
let requests: Vec<_> = (0..n)
|
||||
.map(|_| agent_request(&responder.callout_xkey, unix_now_plus(60)))
|
||||
.collect();
|
||||
let (answered, mut replies) = tokio::sync::mpsc::unbounded_channel();
|
||||
let started = tokio::time::Instant::now();
|
||||
answer_each(
|
||||
futures_util::stream::iter(requests),
|
||||
CONCURRENT_REQUESTS,
|
||||
move |msg| {
|
||||
let responder = Arc::clone(&responder);
|
||||
let answered = answered.clone();
|
||||
async move {
|
||||
let reply = responder.answer(&msg).await;
|
||||
answered
|
||||
.send(reply.is_some())
|
||||
.expect("the test is listening");
|
||||
}
|
||||
},
|
||||
)
|
||||
.await;
|
||||
for _ in 0..n {
|
||||
assert!(replies.recv().await.expect("a reply"), "answered");
|
||||
}
|
||||
started.elapsed()
|
||||
}
|
||||
|
||||
/// Up to the permit count, requests that each wait on the store finish in
|
||||
/// about one lookup, not one lookup each; one more waits for a permit.
|
||||
/// The 16 is fixed, so a permit count too small to answer that many
|
||||
/// logins together fails here.
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn requests_are_answered_concurrently_up_to_the_permit_count() {
|
||||
for (n, lookups) in [
|
||||
(16, 1),
|
||||
(CONCURRENT_REQUESTS, 1),
|
||||
(CONCURRENT_REQUESTS + 1, 2),
|
||||
] {
|
||||
let took = time_to_answer(n).await;
|
||||
let expected = LOOKUP * lookups;
|
||||
assert!(
|
||||
took >= expected && took < expected + LOOKUP / 2,
|
||||
"{n} requests took {took:?}, expected about {expected:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// A request whose `exp` has passed gets no lookup and no reply. The
|
||||
/// same request with `exp` ahead is looked up and answered.
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn a_request_the_server_gave_up_on_is_skipped_without_a_lookup() {
|
||||
let responder = slow_responder();
|
||||
let lookups = || {
|
||||
responder
|
||||
.store
|
||||
.lookups
|
||||
.load(std::sync::atomic::Ordering::SeqCst)
|
||||
};
|
||||
let stale = agent_request(&responder.callout_xkey, unix_now_plus(-1));
|
||||
assert!(responder.answer(&stale).await.is_none());
|
||||
assert_eq!(lookups(), 0);
|
||||
|
||||
let live = agent_request(&responder.callout_xkey, unix_now_plus(60));
|
||||
assert!(responder.answer(&live).await.is_some());
|
||||
assert_eq!(lookups(), 1);
|
||||
}
|
||||
|
||||
/// `exp` is rounded down, so the server is still waiting until the
|
||||
/// second after it begins.
|
||||
#[test]
|
||||
fn a_request_expires_at_the_second_after_its_exp() {
|
||||
let exp = 1_760_000_000;
|
||||
let at = |millis| UNIX_EPOCH + std::time::Duration::from_millis(millis);
|
||||
assert!(!expired(Some(exp), at(1_760_000_000_999)));
|
||||
assert!(expired(Some(exp), at(1_760_000_001_000)));
|
||||
assert!(!expired(None, at(1_760_000_001_000)));
|
||||
}
|
||||
|
||||
/// The control for the case above: an ordinary value parses through the
|
||||
/// same flag, so a failure there is about the leading `-` and not about
|
||||
/// the argument being unknown.
|
||||
|
|
|
|||
Loading…
Reference in a new issue