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.

swarm-nats-auth/README.md

@@ -73,6 +73,9 @@ timeout`**, with a test asserting the relation — a responder that answers
   It is one request's whole budget: when the `IdP` refuses the responder's own
   secret, the re-read from the store and the one retry fit inside it, or the
   request is denied and the re-read lands for the next one.
+- **The responder answers requests concurrently, up to a fixed cap**, so one
+  slow lookup doesn't hold the requests behind it. A request whose `exp` has
+  passed gets no lookup and no reply: the server has stopped waiting for it.
 - **Every issued user JWT expires**, at the presented OIDC token's own `exp`
   or an hour after issue, whichever is sooner. nats-server closes the
   connection at that point, and the client's reconnect comes back through

swarm-nats-auth/src/agent_token.rs

@@ -37,9 +37,9 @@ pub fn classify(token: &str) -> Presented<'_> {
 /// The grant for an agent token, or `None` for a denial.
 ///
 /// The store lookup is bounded by introspection's budget, under the server's
-/// `authorization.timeout`. The callout loop answers one request at a time, so
-/// a store that does not answer holds every request behind it, OIDC ones
-/// included, for at most that long, and this then denies.
+/// `authorization.timeout`. A store that does not answer holds one of the
+/// responder's `CONCURRENT_REQUESTS` permits for at most that long, and this
+/// then denies.
 pub async fn authorize(
     policy: &Policy,
     source: &impl CredentialSource,
@@ -133,8 +133,8 @@ mod tests {
         );
     }
 
-    /// The callout loop answers one request at a time, so a store that never
-    /// answers must cost at most the bound, and then deny.
+    /// A store that never answers holds a request's permit, so it must cost
+    /// at most the bound, and then deny.
     #[tokio::test]
     async fn a_store_that_never_answers_is_denied_within_the_bound() {
         let bound = crate::introspect::INTROSPECTION_TIMEOUT;

swarm-nats-auth/src/main.rs

@@ -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.

swarm-nats-auth/src/oidc_secret.rs

@@ -87,7 +87,7 @@ impl OidcSecret {
         }))
     }
 
-    fn new(secret: StoredSecret) -> Self {
+    pub(crate) fn new(secret: StoredSecret) -> Self {
         Self { secret }
     }
 

swarm-nats-auth/src/request.rs

@@ -48,19 +48,24 @@ pub struct ServerId {
     pub id: String,
 }
 
-/// Split a NATS JWT and deserialize its claim set's `nats` object.
+/// The claims of a request that this responder acts on. `iss`/`sub` are the
+/// server's business, and claiming to understand fields we do not act on
+/// would invite someone to trust them.
+#[derive(Debug, Deserialize)]
+pub struct Claims {
+    /// When the server stops waiting for the answer, in whole Unix seconds:
+    /// nats-server sets it to the time it sent the request plus its
+    /// `authorization.timeout`, rounded down.
+    #[serde(default)]
+    pub exp: Option<i64>,
+    pub nats: AuthRequest,
+}
+
+/// Split a NATS JWT and deserialize its claim set.
 ///
 /// The signature is **not** verified here — see the module docs on
 /// `main.rs` for why that is a decision and not an oversight.
-pub fn decode(payload: &[u8]) -> Result<AuthRequest> {
-    /// Only the `nats` object is modelled; `iss`/`sub`/`exp` are the server's
-    /// business, and claiming to understand fields we do not act on would
-    /// invite someone to trust them.
-    #[derive(Deserialize)]
-    struct Claims {
-        nats: AuthRequest,
-    }
-
+pub fn decode(payload: &[u8]) -> Result<Claims> {
     let text = std::str::from_utf8(payload).context("auth request is not utf-8")?;
     let mut parts = text.split('.');
     // Exactly three segments: a fourth means this is not the shape we think
@@ -73,9 +78,7 @@ pub fn decode(payload: &[u8]) -> Result<AuthRequest> {
     };
     let raw = base64_url_decode(claims).context("claim segment is not base64url")?;
 
-    let claims: Claims =
-        serde_json::from_slice(&raw).context("claim segment is not the expected JSON")?;
-    Ok(claims.nats)
+    serde_json::from_slice(&raw).context("claim segment is not the expected JSON")
 }
 
 /// base64url without padding, as JWT uses. Written out rather than pulled in:
@@ -101,15 +104,16 @@ mod tests {
     fn decodes_a_token_request() {
         let req = decode(
             jwt_with(
-                r#"{"nats":{"connect_opts":{"auth_token":"t0k"},
+                r#"{"exp":1760000002,"nats":{"connect_opts":{"auth_token":"t0k"},
                    "user_nkey":"UABC","server_id":{"id":"NDEADBEEF"}}}"#,
             )
             .as_bytes(),
         )
         .expect("decode");
-        assert_eq!(req.connect_opts.auth_token.as_deref(), Some("t0k"));
-        assert_eq!(req.user_nkey, "UABC");
-        assert_eq!(req.server_id.id, "NDEADBEEF");
+        assert_eq!(req.nats.connect_opts.auth_token.as_deref(), Some("t0k"));
+        assert_eq!(req.nats.user_nkey, "UABC");
+        assert_eq!(req.nats.server_id.id, "NDEADBEEF");
+        assert_eq!(req.exp, Some(1_760_000_002));
     }
 
     #[test]
@@ -122,7 +126,8 @@ mod tests {
                 .as_bytes(),
         )
         .expect("decode");
-        assert!(req.connect_opts.auth_token.is_none());
+        assert!(req.nats.connect_opts.auth_token.is_none());
+        assert_eq!(req.exp, None);
     }
 
     #[test]

swarm-nats-auth/src/xkey.rs

@@ -240,7 +240,7 @@ pub(crate) mod tests {
         let (server_key, opened) = open(&own, Some(&server.public_key()), &sealed).expect("opens");
         assert_eq!(server_key.public_key(), server.public_key());
         let req = crate::request::decode(&opened).expect("a request");
-        assert_eq!(req.user_nkey, "UCLIENT");
+        assert_eq!(req.nats.user_nkey, "UCLIENT");
 
         let issuer = nkeys::KeyPair::new_account();
         let permissions = crate::policy::Permissions {
@@ -251,8 +251,8 @@ pub(crate) mod tests {
         let token = crate::respond::grant(
             &issuer,
             "APP",
-            &req.server_id.id,
-            &req.user_nkey,
+            &req.nats.server_id.id,
+            &req.nats.user_nkey,
             &permissions,
             None,
         );