hyperhive/swarm-secret-client/src/stored_secret.rs

//! A client secret held in memory and read again from where it is kept when
//! the `IdP` refuses it.
//!
//! The secret can change under a running process: authelia's mint, in its
//! container, writes a new one when it re-mints. A refusal asks for a re-read
//! ([`StoredSecret::after_refusal`]); when the read lands a different secret,
//! the caller tries once more with it. One re-read runs at a time, and a
//! refusal while one is running waits for it rather than starting another.
//! The re-read is spawned, so it finishes even when every caller waiting on
//! it has given up, and the next refusal finds what it landed.
//!
//! A re-read that fails, or finds the secret the `IdP` refused, delays the
//! next one: [`FIRST_REREAD_DELAY`], doubled per such re-read up to
//! [`MAX_REREAD_DELAY`]. A refusal inside that delay reads nothing, so a
//! secret that stays wrong costs the store one read per delay, not one per
//! refusal.

use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex, RwLock};
use std::time::Duration;

use futures_util::FutureExt as _;
use futures_util::future::{BoxFuture, Shared};
use tokio::task::AbortHandle;
use tokio::time::Instant;

/// What a failed read of the secret reports. `anyhow::Error` converts into it.
pub type ReadError = Box<dyn std::error::Error + Send + Sync>;

/// The longest one re-read runs. A store client may set no timeout of its
/// own, and while a re-read runs every refusal waits on it.
pub const REREAD_TIMEOUT: Duration = Duration::from_secs(10);

/// The wait before the next re-read after one that brought no new secret.
pub const FIRST_REREAD_DELAY: Duration = Duration::from_secs(5);

/// The longest wait between two re-reads.
pub const MAX_REREAD_DELAY: Duration = Duration::from_mins(5);

type Read =
    dyn Fn() -> Pin<Box<dyn Future<Output = Result<String, ReadError>> + Send>> + Send + Sync;

/// A re-read in progress, shared by every caller waiting on it.
type Reread = Shared<BoxFuture<'static, ()>>;

/// The secret in use and how to read it again. Clones share one secret: a
/// re-read through any of them is seen by all. `Debug` names the client and
/// never the secret.
#[derive(Clone)]
pub struct StoredSecret {
    inner: Arc<Inner>,
}

struct Inner {
    client_id: String,
    read: Box<Read>,
    runtime: tokio::runtime::Handle,
    current: RwLock<String>,
    rereads: Mutex<Rereads>,
}

#[derive(Default)]
struct Rereads {
    running: Option<(AbortHandle, Reread)>,
    /// Re-reads in a row that brought no new secret.
    fruitless: u32,
    /// No re-read starts before this.
    not_before: Option<Instant>,
}

impl std::fmt::Debug for StoredSecret {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StoredSecret")
            .field("client_id", &self.inner.client_id)
            .finish_non_exhaustive()
    }
}

impl StoredSecret {
    /// `client_id`'s secret, `initial` until `read` replaces it. `client_id`
    /// names the client in logs.
    ///
    /// # Panics
    /// Outside a tokio runtime: re-reads are spawned on the one this is made
    /// in.
    pub fn new<F, Fut>(client_id: impl Into<String>, initial: String, read: F) -> Self
    where
        F: Fn() -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<String, ReadError>> + Send + 'static,
    {
        Self {
            inner: Arc::new(Inner {
                client_id: client_id.into(),
                read: Box::new(move || Box::pin(read())),
                runtime: tokio::runtime::Handle::current(),
                current: RwLock::new(initial),
                rereads: Mutex::default(),
            }),
        }
    }

    /// The secret in use.
    #[must_use]
    pub fn current(&self) -> String {
        self.inner.current.read().expect("unpoisoned").clone()
    }

    /// After the `IdP` refused `refused`: the secret to try once more, or
    /// `None` when there is no other one to try. Waits for a re-read, at most
    /// [`REREAD_TIMEOUT`]; a caller with a shorter budget bounds this itself.
    pub async fn after_refusal(&self, refused: &str) -> Option<String> {
        let current = self.current();
        if current != refused {
            return Some(current);
        }
        let Some(reread) = self.reread() else {
            tracing::debug!(
                client_id = %self.inner.client_id,
                "the IdP refused this client's secret inside the re-read backoff; not re-reading"
            );
            return None;
        };
        reread.await;
        let fresh = self.current();
        (fresh != refused).then_some(fresh)
    }

    /// [`after_refusal`](Self::after_refusal), for a thread with no tokio
    /// runtime of its own.
    ///
    /// # Panics
    /// Inside a tokio runtime, as [`tokio::runtime::Handle::block_on`] does.
    #[must_use]
    pub fn after_refusal_blocking(&self, refused: &str) -> Option<String> {
        self.inner.runtime.block_on(self.after_refusal(refused))
    }

    /// The running re-read, or a new one when none is running and the
    /// backoff allows it.
    fn reread(&self) -> Option<Reread> {
        let mut rereads = self.inner.rereads.lock().expect("unpoisoned");
        if let Some((handle, running)) = rereads.running.as_ref()
            && !handle.is_finished()
        {
            return Some(running.clone());
        }
        if rereads.not_before.is_some_and(|at| Instant::now() < at) {
            return None;
        }
        let inner = Arc::clone(&self.inner);
        let task = self
            .inner
            .runtime
            .spawn(async move { inner.reread().await });
        let handle = task.abort_handle();
        let client_id = self.inner.client_id.clone();
        let shared = task
            .map(move |done| {
                if let Err(e) = done {
                    tracing::warn!(
                        %client_id,
                        error = %e,
                        "the client secret re-read task died; keeping the one in use"
                    );
                }
            })
            .boxed()
            .shared();
        rereads.running = Some((handle, shared.clone()));
        Some(shared)
    }
}

impl Inner {
    /// One read of the secret, kept when it differs from the one in use. A
    /// read that brings no new secret pushes the next one back.
    async fn reread(&self) {
        let before = self.current.read().expect("unpoisoned").clone();
        let outcome = tokio::time::timeout(REREAD_TIMEOUT, (self.read)()).await;
        let mut rereads = self.rereads.lock().expect("unpoisoned");
        let failure = match outcome {
            Ok(Ok(secret)) if secret != before => {
                *self.current.write().expect("unpoisoned") = secret;
                rereads.fruitless = 0;
                rereads.not_before = None;
                tracing::info!(
                    client_id = %self.client_id,
                    "re-read a changed client secret after the IdP refused the one in use"
                );
                return;
            }
            Ok(Ok(_)) => "the stored client secret is the one the IdP refused".to_owned(),
            Ok(Err(e)) => format!("could not re-read the client secret: {}", chain(&*e)),
            Err(_) => format!("re-reading the client secret timed out after {REREAD_TIMEOUT:?}"),
        };
        rereads.fruitless = rereads.fruitless.saturating_add(1);
        let delay = delay_after(rereads.fruitless);
        rereads.not_before = Some(Instant::now() + delay);
        tracing::warn!(
            client_id = %self.client_id,
            next_reread_in = ?delay,
            "{failure}; keeping the one in use"
        );
    }
}

/// `error` and its sources, joined with `: `.
fn chain(error: &dyn std::error::Error) -> String {
    let mut rendered = error.to_string();
    let mut source = error.source();
    while let Some(cause) = source {
        rendered.push_str(": ");
        rendered.push_str(&cause.to_string());
        source = cause.source();
    }
    rendered
}

/// Whether an `IdP` answering `status` refused the client's credential (400 or
/// 401), the one answer a re-read can fix. A 429 or a 5xx says nothing about
/// the secret, and re-reading on one would spend the backoff on an outage.
#[must_use]
pub fn is_credential_refusal(status: reqwest::StatusCode) -> bool {
    matches!(
        status,
        reqwest::StatusCode::BAD_REQUEST | reqwest::StatusCode::UNAUTHORIZED
    )
}

/// [`FIRST_REREAD_DELAY`] doubled per fruitless re-read after the first,
/// capped at [`MAX_REREAD_DELAY`].
fn delay_after(fruitless: u32) -> Duration {
    let doublings = fruitless.saturating_sub(1).min(16);
    FIRST_REREAD_DELAY
        .saturating_mul(1 << doublings)
        .min(MAX_REREAD_DELAY)
}

#[cfg(test)]
pub(crate) mod tests {
    use std::sync::atomic::{AtomicUsize, Ordering};

    use super::*;

    /// A store holding `secret`, or failing, that takes `delay` to answer.
    pub(crate) struct Fake {
        pub(crate) secret: Arc<Mutex<Option<&'static str>>>,
        pub(crate) delay: Duration,
        pub(crate) reads: Arc<AtomicUsize>,
    }

    impl Fake {
        pub(crate) fn holding(secret: &'static str) -> Self {
            Self {
                secret: Arc::new(Mutex::new(Some(secret))),
                delay: Duration::ZERO,
                reads: Arc::default(),
            }
        }

        /// `client_id`'s secret, `initial` in use, re-read from this store.
        pub(crate) fn secret(&self, client_id: &str, initial: &str) -> StoredSecret {
            let (secret, delay, reads) = (
                Arc::clone(&self.secret),
                self.delay,
                Arc::clone(&self.reads),
            );
            StoredSecret::new(client_id, initial.to_owned(), move || {
                let (secret, reads) = (Arc::clone(&secret), Arc::clone(&reads));
                async move {
                    reads.fetch_add(1, Ordering::SeqCst);
                    tokio::time::sleep(delay).await;
                    let held = *secret.lock().expect("unpoisoned");
                    held.map(str::to_owned)
                        .ok_or_else(|| "store unreachable".into())
                }
            })
        }

        pub(crate) fn reads(&self) -> usize {
            self.reads.load(Ordering::SeqCst)
        }
    }

    thread_local! {
        static LOG: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
    }

    /// What `tracing` records on the calling thread. The subscriber is
    /// process-wide, because a scoped one loses to the callsite cache when
    /// another test reaches the same event first; the buffer is per thread,
    /// which is where a current-thread test runtime runs the spawned re-read.
    pub(crate) struct Logs;

    impl std::io::Write for Logs {
        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
            LOG.with(|log| log.borrow_mut().extend_from_slice(buf));
            Ok(buf.len())
        }

        fn flush(&mut self) -> std::io::Result<()> {
            Ok(())
        }
    }

    impl Logs {
        /// Starts this thread's capture empty.
        pub(crate) fn capture() {
            static INSTALLED: std::sync::Once = std::sync::Once::new();
            INSTALLED.call_once(|| {
                tracing_subscriber::fmt()
                    .with_writer(|| Logs)
                    .with_ansi(false)
                    .init();
            });
            LOG.with(|log| log.borrow_mut().clear());
        }

        pub(crate) fn text() -> String {
            LOG.with(|log| String::from_utf8(log.borrow().clone()).expect("utf-8"))
        }
    }

    #[tokio::test]
    async fn a_refused_secret_is_re_read_and_the_new_one_offered() {
        Logs::capture();
        let store = Fake::holding("new");
        let secret = store.secret("hive-alpha", "old");
        assert_eq!(secret.after_refusal("old").await.as_deref(), Some("new"));
        assert_eq!(secret.current(), "new");
        assert_eq!(store.reads(), 1);

        let text = Logs::text();
        assert!(text.contains("re-read a changed client secret"), "{text}");
        assert!(text.contains("hive-alpha"), "{text}");
        assert!(!text.contains("old") && !text.contains("new"), "{text}");
    }

    /// A refusal of a secret another caller already replaced offers the
    /// replacement without reading again.
    #[tokio::test]
    async fn a_refusal_of_a_replaced_secret_offers_the_replacement() {
        let store = Fake::holding("new");
        let secret = store.secret("hive-alpha", "old");
        assert_eq!(secret.after_refusal("old").await.as_deref(), Some("new"));
        assert_eq!(secret.after_refusal("old").await.as_deref(), Some("new"));
        assert_eq!(store.reads(), 1);
    }

    /// A secret that stays wrong: the first refusal reads, refusals inside
    /// the backoff read nothing, and the backoff doubles per fruitless read.
    #[tokio::test(start_paused = true)]
    async fn a_permanently_refused_secret_backs_off() {
        Logs::capture();
        let store = Fake::holding("old");
        let secret = store.secret("hive-alpha", "old");

        assert_eq!(secret.after_refusal("old").await, None);
        assert_eq!(store.reads(), 1);
        let text = Logs::text();
        assert!(text.contains("is the one the IdP refused"), "{text}");
        assert!(text.contains("hive-alpha"), "{text}");
        assert!(!text.contains("\"old\""), "{text}");

        for _ in 0..10 {
            assert_eq!(secret.after_refusal("old").await, None);
        }
        assert_eq!(store.reads(), 1);

        tokio::time::sleep(FIRST_REREAD_DELAY).await;
        assert_eq!(secret.after_refusal("old").await, None);
        assert_eq!(store.reads(), 2);

        tokio::time::sleep(FIRST_REREAD_DELAY).await;
        assert_eq!(secret.after_refusal("old").await, None);
        assert_eq!(
            store.reads(),
            2,
            "the second fruitless read doubles the wait"
        );

        tokio::time::sleep(FIRST_REREAD_DELAY).await;
        assert_eq!(secret.after_refusal("old").await, None);
        assert_eq!(store.reads(), 3);
    }

    /// Once the store holds a new secret, the first refusal past the backoff
    /// picks it up.
    #[tokio::test(start_paused = true)]
    async fn a_secret_published_during_the_backoff_is_read_after_it() {
        let store = Fake::holding("old");
        let secret = store.secret("hive-alpha", "old");
        assert_eq!(secret.after_refusal("old").await, None);

        *store.secret.lock().expect("unpoisoned") = Some("new");
        assert_eq!(secret.after_refusal("old").await, None);
        tokio::time::sleep(FIRST_REREAD_DELAY).await;
        assert_eq!(secret.after_refusal("old").await.as_deref(), Some("new"));
        assert_eq!(store.reads(), 2);
    }

    #[tokio::test]
    async fn a_failed_re_read_keeps_the_secret_in_use() {
        Logs::capture();
        let store = Fake::holding("new");
        *store.secret.lock().expect("unpoisoned") = None;
        let secret = store.secret("hive-alpha", "old");
        assert_eq!(secret.after_refusal("old").await, None);
        assert_eq!(secret.current(), "old");

        let text = Logs::text();
        assert!(
            text.contains("could not re-read the client secret"),
            "{text}"
        );
        assert!(text.contains("store unreachable"), "{text}");
        assert!(!text.contains("is the one the IdP refused"), "{text}");
        assert!(!text.contains("old"), "no secret in the log: {text}");
    }

    /// A refusal while a re-read is still running waits on that one rather
    /// than starting another.
    #[tokio::test(start_paused = true)]
    async fn a_refusal_during_a_running_re_read_joins_it() {
        let mut store = Fake::holding("new");
        store.delay = Duration::from_secs(2);
        let secret = store.secret("hive-alpha", "old");
        let (first, second) =
            tokio::join!(secret.after_refusal("old"), secret.after_refusal("old"));
        assert_eq!(first.as_deref(), Some("new"));
        assert_eq!(second.as_deref(), Some("new"));
        assert_eq!(store.reads(), 1);
    }

    /// A caller that gives up leaves the re-read running; the next refusal
    /// finds what it landed without reading again.
    #[tokio::test(start_paused = true)]
    async fn a_re_read_outlives_a_caller_that_gave_up() {
        let mut store = Fake::holding("new");
        store.delay = Duration::from_secs(3);
        let secret = store.secret("hive-alpha", "old");
        let gave_up =
            tokio::time::timeout(Duration::from_secs(1), secret.after_refusal("old")).await;
        assert!(gave_up.is_err());

        tokio::time::sleep(Duration::from_secs(3)).await;
        assert_eq!(secret.current(), "new");
        assert_eq!(secret.after_refusal("old").await.as_deref(), Some("new"));
        assert_eq!(store.reads(), 1);
    }

    /// A re-read that stalls is given up at [`REREAD_TIMEOUT`], keeping the
    /// secret in use, and counts as fruitless.
    #[tokio::test(start_paused = true)]
    async fn a_stalled_re_read_is_dropped_at_its_timeout() {
        let mut store = Fake::holding("new");
        store.delay = REREAD_TIMEOUT * 6;
        let secret = store.secret("hive-alpha", "old");
        assert_eq!(secret.after_refusal("old").await, None);
        assert_eq!(secret.current(), "old");

        tokio::time::sleep(FIRST_REREAD_DELAY).await;
        assert_eq!(secret.after_refusal("old").await, None);
        assert_eq!(store.reads(), 2);
    }

    #[test]
    fn the_backoff_doubles_up_to_its_cap() {
        assert_eq!(delay_after(1), FIRST_REREAD_DELAY);
        assert_eq!(delay_after(2), FIRST_REREAD_DELAY * 2);
        assert_eq!(delay_after(40), MAX_REREAD_DELAY);
    }

    #[tokio::test]
    async fn a_debug_print_does_not_carry_the_secret() {
        let secret = Fake::holding("new").secret("hive-alpha", "s3cret");
        let printed = format!("{secret:?}");
        assert!(!printed.contains("s3cret"), "{printed}");
        assert!(printed.contains("hive-alpha"), "{printed}");
    }
}