swarm-logs-auth: rate-limit verification and cache positive verdicts

Any request on logs.<swarm> carrying a well-formed swarm-agent bearer
made swarm-logs-auth do a bao cert login and KV read before it could
refuse, with no cap on the rate.

nginx now applies limit_req (zone swarm_logs_agent_auth, keyed by
client address, 5r/s burst 20 nodelay by default) on the internal
location that forwards to swarm-logs-auth, answering 403 when over,
since auth_request passes through only 401/403 and swarm-logs retries
a 401. The authelia path and the query location carry no limit.

swarm-logs-auth remembers bearers it has granted, keyed by SHA-256 of
the bearer, for --verdict-ttl seconds (deploy.victorialogs.authVerdictTtl,
60 by default), holding at most 4096 and evicting the oldest. Denials
are never remembered. A replaced or deleted queue secret therefore
keeps working at the log store for up to the TTL after its last check.

Closes #5190

swarm-logs-auth/Cargo.toml

@@ -11,6 +11,8 @@ path = "src/main.rs"
 anyhow.workspace = true
 axum.workspace = true
 clap.workspace = true
+# The verdict cache's key: a digest of the bearer, never the bearer.
+sha2.workspace = true
 swarm-agent-verify.workspace = true
 # The agent-token spelling, and the rustls provider the store client needs.
 swarm-queue-client.workspace = true

swarm-logs-auth/src/main.rs

@@ -9,11 +9,17 @@
 //! It reads the store under its own certificate, whose policy reaches every
 //! agent's `queue` secret and nothing else. It serves a unix socket whose
 //! directory only the gateway's nginx may enter.
+//!
+//! A grant is remembered for `--verdict-ttl` seconds ([`verdicts`]), so an
+//! agent's queries within that window cost no store lookup, and a secret that
+//! stops verifying is still accepted for up to that long.
+
+mod verdicts;
 
 use std::os::unix::fs::PermissionsExt as _;
 use std::path::PathBuf;
-use std::sync::Arc;
-use std::time::Duration;
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
 
 use anyhow::Context as _;
 use axum::extract::State;
@@ -21,12 +27,17 @@ use axum::http::{HeaderMap, StatusCode, header};
 use clap::Parser;
 use swarm_agent_verify::CredentialSource;
 use swarm_queue_client::agent_token::parse_agent_token;
+use verdicts::Verdicts;
 
 /// How long one store lookup may take before the request is denied. Well
 /// under nginx's 60s `proxy_read_timeout`, so a slow store reads as a 401
 /// rather than a gateway timeout.
 const LOOKUP_TIMEOUT: Duration = Duration::from_secs(5);
 
+/// Most grants held at once. One live secret per agent, so this is a bound on
+/// memory, not a size a swarm is expected to reach.
+const VERDICT_CAP: usize = 4096;
+
 #[derive(Parser)]
 #[command(about = "auth_request responder: checks an agent's queue secret against the store")]
 struct Args {
@@ -36,15 +47,28 @@ struct Args {
     /// Cert-auth role this responder logs in to the store as.
     #[arg(long, default_value = "swarm-logs-auth")]
     store_cert_role: String,
+    /// Seconds a verified bearer is granted again without a store lookup;
+    /// 0 looks up every request.
+    #[arg(long, default_value_t = 60)]
+    verdict_ttl: u64,
 }
 
-/// The answer for one subrequest's `Authorization` header.
-async fn decide(source: &impl CredentialSource, authorization: Option<&str>) -> StatusCode {
-    let Some(token) = authorization.and_then(|v| v.strip_prefix("Bearer ")) else {
+struct Responder<S> {
+    source: S,
+    verdicts: Mutex<Verdicts>,
+}
+
+/// The answer for one subrequest's `Authorization` header, at `now`.
+async fn decide(
+    responder: &Responder<impl CredentialSource>,
+    authorization: Option<&str>,
+    now: Instant,
+) -> StatusCode {
+    let Some(bearer) = authorization.and_then(|v| v.strip_prefix("Bearer ")) else {
         tracing::warn!("no bearer token; denying");
         return StatusCode::UNAUTHORIZED;
     };
-    let token = match parse_agent_token(token) {
+    let token = match parse_agent_token(bearer) {
         Some(Ok(token)) => token,
         Some(Err(e)) => {
             tracing::warn!(error = %e, "malformed agent token; denying");
@@ -55,22 +79,35 @@ async fn decide(source: &impl CredentialSource, authorization: Option<&str>) ->
             return StatusCode::UNAUTHORIZED;
         }
     };
-    if swarm_agent_verify::verify(source, &token, LOOKUP_TIMEOUT).await {
+    if lock(&responder.verdicts).granted(bearer, now) {
+        tracing::debug!(agent = token.agent, "verified recently");
+        return StatusCode::OK;
+    }
+    if swarm_agent_verify::verify(&responder.source, &token, LOOKUP_TIMEOUT).await {
         tracing::info!(agent = token.agent, "verified");
+        lock(&responder.verdicts).grant(bearer, now);
         StatusCode::OK
     } else {
         StatusCode::UNAUTHORIZED
     }
 }
 
+/// The cache holds no invariant a panicking holder could break mid-way, so a
+/// poisoned lock is used as is.
+fn lock(verdicts: &Mutex<Verdicts>) -> std::sync::MutexGuard<'_, Verdicts> {
+    verdicts
+        .lock()
+        .unwrap_or_else(std::sync::PoisonError::into_inner)
+}
+
 async fn handle(
-    State(store): State<Arc<swarm_agent_verify::Store>>,
+    State(responder): State<Arc<Responder<swarm_agent_verify::Store>>>,
     headers: HeaderMap,
 ) -> StatusCode {
     let authorization = headers
         .get(header::AUTHORIZATION)
         .and_then(|v| v.to_str().ok());
-    decide(store.as_ref(), authorization).await
+    decide(&responder, authorization, Instant::now()).await
 }
 
 #[tokio::main]
@@ -88,10 +125,13 @@ async fn main() -> anyhow::Result<()> {
 
     let settings = swarm_secret_client::client::Settings::from_env()
         .context("reading the secret store's address and this responder's identity")?;
-    let store = Arc::new(swarm_agent_verify::Store::new(
-        settings,
-        args.store_cert_role,
-    ));
+    let responder = Arc::new(Responder {
+        source: swarm_agent_verify::Store::new(settings, args.store_cert_role),
+        verdicts: Mutex::new(Verdicts::new(
+            Duration::from_secs(args.verdict_ttl),
+            VERDICT_CAP,
+        )),
+    });
 
     // `RuntimeDirectoryPreserve=yes` can leave the previous run's socket; the
     // directory is this unit's alone, so unlinking it is safe.
@@ -109,52 +149,116 @@ async fn main() -> anyhow::Result<()> {
         .with_context(|| format!("chmod {}", args.socket.display()))?;
     tracing::info!(socket = %args.socket.display(), "swarm-logs-auth listening");
 
-    let app = axum::Router::new().fallback(handle).with_state(store);
+    let app = axum::Router::new().fallback(handle).with_state(responder);
     axum::serve(listener, app).await.context("serving")
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
     use swarm_secret_client::queue::AgentCredential;
 
-    /// A store holding atlas's credential.
-    struct Atlas;
+    /// A store holding atlas's credential until `revoke`, counting lookups.
+    #[derive(Default)]
+    struct Atlas {
+        lookups: AtomicUsize,
+        revoked: AtomicBool,
+    }
 
-    impl CredentialSource for Atlas {
+    impl Atlas {
+        fn revoke(&self) {
+            self.revoked.store(true, Ordering::SeqCst);
+        }
+
+        fn lookups(&self) -> usize {
+            self.lookups.load(Ordering::SeqCst)
+        }
+    }
+
+    impl CredentialSource for &Atlas {
         async fn lookup(&self, agent: &str) -> anyhow::Result<Option<AgentCredential>> {
-            Ok((agent == "atlas").then(|| AgentCredential {
-                value: "Ab9_-zSECRET".to_owned(),
-                agent: "atlas".to_owned(),
-                minted_at: None,
-            }))
+            self.lookups.fetch_add(1, Ordering::SeqCst);
+            Ok(
+                (agent == "atlas" && !self.revoked.load(Ordering::SeqCst)).then(|| {
+                    AgentCredential {
+                        value: "Ab9_-zSECRET".to_owned(),
+                        agent: "atlas".to_owned(),
+                        minted_at: None,
+                    }
+                }),
+            )
+        }
+    }
+
+    const TTL: Duration = Duration::from_mins(1);
+    const GOOD: &str = "Bearer swarm-agent.atlas.Ab9_-zSECRET";
+    const WRONG: &str = "Bearer swarm-agent.atlas.Ab9_-zSECREX";
+
+    fn responder(store: &Atlas) -> Responder<&Atlas> {
+        Responder {
+            source: store,
+            verdicts: Mutex::new(Verdicts::new(TTL, 16)),
         }
     }
 
     #[tokio::test]
     async fn the_stored_secret_is_200() {
+        let store = Atlas::default();
         assert_eq!(
-            decide(&Atlas, Some("Bearer swarm-agent.atlas.Ab9_-zSECRET")).await,
+            decide(&responder(&store), Some(GOOD), Instant::now()).await,
             StatusCode::OK
         );
     }
 
     #[tokio::test]
     async fn everything_else_is_401() {
+        let store = Atlas::default();
+        let r = responder(&store);
         for authorization in [
             None,
             Some("swarm-agent.atlas.Ab9_-zSECRET"),
             Some("Basic swarm-agent.atlas.Ab9_-zSECRET"),
-            Some("Bearer swarm-agent.atlas.Ab9_-zSECREX"),
+            Some(WRONG),
             Some("Bearer swarm-agent.argus.Ab9_-zSECRET"),
             Some("Bearer swarm-agent.atlas"),
             Some("Bearer authelia_at_abc.def"),
         ] {
             assert_eq!(
-                decide(&Atlas, authorization).await,
+                decide(&r, authorization, Instant::now()).await,
                 StatusCode::UNAUTHORIZED,
                 "{authorization:?}"
             );
         }
     }
+
+    /// A grant is reused within the TTL — a revocation the store already
+    /// holds goes unseen — and checked against the store again after it.
+    #[tokio::test]
+    async fn a_grant_spares_the_store_until_the_ttl() {
+        let store = Atlas::default();
+        let r = responder(&store);
+        let t0 = Instant::now();
+        assert_eq!(decide(&r, Some(GOOD), t0).await, StatusCode::OK);
+        store.revoke();
+        let within = t0 + Duration::from_secs(59);
+        assert_eq!(decide(&r, Some(GOOD), within).await, StatusCode::OK);
+        assert_eq!(store.lookups(), 1);
+        assert_eq!(
+            decide(&r, Some(GOOD), t0 + TTL).await,
+            StatusCode::UNAUTHORIZED
+        );
+        assert_eq!(store.lookups(), 2);
+    }
+
+    #[tokio::test]
+    async fn a_denial_is_looked_up_every_time() {
+        let store = Atlas::default();
+        let r = responder(&store);
+        let t0 = Instant::now();
+        for _ in 0..3 {
+            assert_eq!(decide(&r, Some(WRONG), t0).await, StatusCode::UNAUTHORIZED);
+        }
+        assert_eq!(store.lookups(), 3);
+    }
 }

swarm-logs-auth/src/verdicts.rs

@@ -0,0 +1,131 @@
+//! Bearers verified recently enough to be granted again without a store
+//! lookup.
+//!
+//! Only grants are held. A denial is never remembered, so a credential that
+//! starts verifying is accepted on its next request; a credential that stops
+//! verifying (re-minted, or the agent destroyed) is still accepted until its
+//! entry expires, at most the TTL after its last successful lookup.
+//!
+//! Entries are keyed by the SHA-256 of the bearer, so the process never holds
+//! a presented secret past the request that carried it.
+
+use std::collections::HashMap;
+use std::time::{Duration, Instant};
+
+use sha2::{Digest, Sha256};
+
+type Key = [u8; 32];
+
+fn key(bearer: &str) -> Key {
+    Sha256::digest(bearer.as_bytes()).into()
+}
+
+/// Granted bearers, each until its expiry, at most `cap` of them.
+pub struct Verdicts {
+    ttl: Duration,
+    cap: usize,
+    expiries: HashMap<Key, Instant>,
+}
+
+impl Verdicts {
+    /// A TTL of zero holds nothing, so every request reaches the store.
+    pub fn new(ttl: Duration, cap: usize) -> Self {
+        Self {
+            ttl,
+            cap,
+            expiries: HashMap::new(),
+        }
+    }
+
+    /// Whether `bearer` was granted less than the TTL before `now`.
+    pub fn granted(&self, bearer: &str, now: Instant) -> bool {
+        self.expiries
+            .get(&key(bearer))
+            .is_some_and(|expiry| now < *expiry)
+    }
+
+    /// Records that `bearer` verified at `now`. When full, the entry closest
+    /// to expiry goes, which with one TTL for every entry is the oldest.
+    pub fn grant(&mut self, bearer: &str, now: Instant) {
+        if self.ttl.is_zero() {
+            return;
+        }
+        let key = key(bearer);
+        if !self.expiries.contains_key(&key)
+            && self.expiries.len() >= self.cap
+            && let Some(oldest) = self
+                .expiries
+                .iter()
+                .min_by_key(|(_, expiry)| **expiry)
+                .map(|(k, _)| *k)
+        {
+            self.expiries.remove(&oldest);
+        }
+        self.expiries.insert(key, now + self.ttl);
+    }
+
+    #[cfg(test)]
+    fn len(&self) -> usize {
+        self.expiries.len()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const TTL: Duration = Duration::from_mins(1);
+
+    #[test]
+    fn a_grant_holds_until_the_ttl() {
+        let t0 = Instant::now();
+        let mut v = Verdicts::new(TTL, 8);
+        assert!(!v.granted("swarm-agent.atlas.s", t0));
+        v.grant("swarm-agent.atlas.s", t0);
+        assert!(v.granted("swarm-agent.atlas.s", t0));
+        assert!(v.granted("swarm-agent.atlas.s", t0 + Duration::from_millis(59_999)));
+        assert!(!v.granted("swarm-agent.atlas.s", t0 + TTL));
+    }
+
+    #[test]
+    fn a_grant_covers_only_the_exact_bearer() {
+        let t0 = Instant::now();
+        let mut v = Verdicts::new(TTL, 8);
+        v.grant("swarm-agent.atlas.s", t0);
+        assert!(!v.granted("swarm-agent.atlas.t", t0));
+        assert!(!v.granted("swarm-agent.argus.s", t0));
+    }
+
+    #[test]
+    fn a_zero_ttl_holds_nothing() {
+        let t0 = Instant::now();
+        let mut v = Verdicts::new(Duration::ZERO, 8);
+        v.grant("swarm-agent.atlas.s", t0);
+        assert!(!v.granted("swarm-agent.atlas.s", t0));
+        assert_eq!(v.len(), 0);
+    }
+
+    #[test]
+    fn a_full_cache_evicts_the_oldest_grant() {
+        let t0 = Instant::now();
+        let mut v = Verdicts::new(TTL, 3);
+        for (i, bearer) in ["a", "b", "c", "d"].into_iter().enumerate() {
+            v.grant(bearer, t0 + Duration::from_secs(i as u64));
+        }
+        let now = t0 + Duration::from_secs(4);
+        assert_eq!(v.len(), 3);
+        assert!(!v.granted("a", now));
+        assert!(v.granted("b", now) && v.granted("c", now) && v.granted("d", now));
+    }
+
+    #[test]
+    fn regranting_a_held_bearer_in_a_full_cache_evicts_nothing() {
+        let t0 = Instant::now();
+        let mut v = Verdicts::new(TTL, 2);
+        v.grant("a", t0);
+        v.grant("b", t0);
+        v.grant("a", t0 + Duration::from_secs(1));
+        assert!(v.granted("a", t0 + Duration::from_secs(1)));
+        assert!(v.granted("b", t0 + Duration::from_secs(1)));
+    }
+}