hyperhive/swarm-controller/src/forge/api.rs

//! The forge API under the controller's token, re-read from the swarm secret
//! store when the forge refuses it.
//!
//! [`Api::call`] runs one request. When the forge answers 401, it asks the
//! [`StoredSecret`] for a token to try instead and runs the request once more
//! with it; a second 401 is returned as it came. The `StoredSecret` runs one
//! store read at a time and backs off after a read that brings no new token,
//! so a burst of refusals costs the store one read.

use std::sync::{Arc, Mutex};

use forgejo_api::{ApiErrorKind, Auth, Forgejo, ForgejoError};
use futures_util::future::BoxFuture;
use reqwest::StatusCode;
use swarm_queue_client::stored_secret::StoredSecret;

/// One forge request, as [`Api::call`] runs it.
pub(crate) type Call<'a, R> = BoxFuture<'a, Result<R, ForgejoError>>;

pub(crate) struct Api {
    url: url::Url,
    token: StoredSecret,
    /// A client sending the token it is paired with. Rebuilt when `token`
    /// holds a different one.
    built: Mutex<(String, Arc<Forgejo>)>,
}

impl Api {
    /// A client for the forge at `url`, sending `token`'s current value.
    ///
    /// # Errors
    /// `url` is not http(s), or the token is not ASCII.
    pub(crate) fn new(url: url::Url, token: StoredSecret) -> Result<Self, ForgejoError> {
        let current = token.current();
        let api = Arc::new(Forgejo::new(Auth::Token(&current), url.clone())?);
        Ok(Self {
            url,
            token,
            built: Mutex::new((current, api)),
        })
    }

    /// The token in use and a client sending it.
    fn current(&self) -> Result<(String, Arc<Forgejo>), ForgejoError> {
        let token = self.token.current();
        let mut built = self.built.lock().expect("unpoisoned");
        if built.0 != token {
            let api = Arc::new(Forgejo::new(Auth::Token(&token), self.url.clone())?);
            *built = (token.clone(), api);
        }
        Ok((token, Arc::clone(&built.1)))
    }

    /// `request` against the forge, run a second time with a re-read token
    /// when the forge refuses the first. The forge acts on nothing it answers
    /// 401 to, so running a non-idempotent request again is safe.
    ///
    /// Boxed, so a caller awaiting several of these does not carry each one's
    /// `R` inline.
    pub(crate) fn call<'s, R, F>(&'s self, request: F) -> Call<'s, R>
    where
        F: Fn(&Forgejo) -> Call<'_, R> + Send + 's,
        R: Send + 's,
    {
        Box::pin(async move {
            let (token, api) = self.current()?;
            match request(&api).await {
                Err(e) if refuses_the_token(&e) => {
                    if self.token.after_refusal(&token).await.is_none() {
                        return Err(e);
                    }
                    let (_, api) = self.current()?;
                    request(&api).await
                }
                done => done,
            }
        })
    }
}

/// The HTTP status `e` reports, for the shapes a 401 can arrive in: endpoints
/// that declare a 401 parse its body into [`ApiErrorKind::Unauthorized`], the
/// rest report the bare status.
fn status(e: &ForgejoError) -> Option<StatusCode> {
    match e {
        ForgejoError::UnexpectedStatusCode(s) => Some(*s),
        ForgejoError::ApiError(e) => match e.error_kind() {
            ApiErrorKind::Unauthorized => Some(StatusCode::UNAUTHORIZED),
            ApiErrorKind::Other(s) => Some(*s),
            _ => None,
        },
        _ => None,
    }
}

/// Whether the forge refused the token. Only a 401: to the forge a 400 is a
/// malformed request, and a re-read it started would find the same token and
/// put a real 401 right after it inside the re-read backoff.
fn refuses_the_token(e: &ForgejoError) -> bool {
    status(e) == Some(StatusCode::UNAUTHORIZED)
}

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

    use super::*;

    /// A forge answering 200 to `accepted` and 401 to any other token,
    /// counting the requests it serves. The `held`-th request to arrive, from
    /// 0, is answered only after one with `accepted` is.
    async fn forge_accepting(
        accepted: &'static str,
        held: Option<usize>,
    ) -> (url::Url, Arc<AtomicUsize>) {
        use axum::http::{HeaderMap, StatusCode, header::AUTHORIZATION, header::CONTENT_TYPE};

        let served = Arc::new(AtomicUsize::new(0));
        let counted = Arc::clone(&served);
        let (accepted_once, seen_accepted) = tokio::sync::watch::channel(false);
        let app = axum::Router::new().fallback(move |headers: HeaderMap| {
            let counted = Arc::clone(&counted);
            let (accepted_once, mut seen_accepted) = (accepted_once.clone(), seen_accepted.clone());
            async move {
                let n = counted.fetch_add(1, Ordering::SeqCst);
                if held == Some(n) {
                    seen_accepted.wait_for(|seen| *seen).await.unwrap();
                }
                let sent = headers.get(AUTHORIZATION).and_then(|v| v.to_str().ok());
                if sent == Some(&format!("token {accepted}")) {
                    accepted_once.send_replace(true);
                    (StatusCode::OK, [(CONTENT_TYPE, "application/json")], "{}")
                } else {
                    let body = r#"{"message":"token is required"}"#;
                    (
                        StatusCode::UNAUTHORIZED,
                        [(CONTENT_TYPE, "application/json")],
                        body,
                    )
                }
            }
        });
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let url = url::Url::parse(&format!("http://{}/", listener.local_addr().unwrap())).unwrap();
        tokio::spawn(async move { axum::serve(listener, app).await });
        (url, served)
    }

    /// A token starting at `initial` whose store holds `stored`, counting the
    /// reads, each of which takes a second.
    fn token(initial: &str, stored: &'static str) -> (StoredSecret, Arc<AtomicUsize>) {
        let reads = Arc::new(AtomicUsize::new(0));
        let counted = Arc::clone(&reads);
        let secret = StoredSecret::new("test", initial.to_owned(), move || {
            let counted = Arc::clone(&counted);
            async move {
                counted.fetch_add(1, Ordering::SeqCst);
                tokio::time::sleep(std::time::Duration::from_secs(1)).await;
                Ok(stored.to_owned())
            }
        });
        (secret, reads)
    }

    /// A client for `url` holding `token`, whose store read always fails.
    pub(crate) fn fixed(url: url::Url, token: &str) -> Api {
        let token = StoredSecret::new("test", token.to_owned(), || async {
            Err("no store in this test".into())
        });
        Api::new(url, token).unwrap()
    }

    async fn whoami(api: &Api) -> Result<forgejo_api::structs::User, ForgejoError> {
        api.call(|api| api.user_get_current().into_future()).await
    }

    #[tokio::test(start_paused = true)]
    async fn a_refused_token_is_re_read_and_the_request_sent_again_with_the_new_one() {
        let (url, served) = forge_accepting("new", None).await;
        let (token, reads) = token("old", "new");
        let api = Api::new(url, token).unwrap();
        whoami(&api).await.expect("accepted with the re-read token");
        assert_eq!(reads.load(Ordering::SeqCst), 1);
        assert_eq!(served.load(Ordering::SeqCst), 2);
        whoami(&api).await.expect("the new token stays in use");
        assert_eq!(reads.load(Ordering::SeqCst), 1);
        assert_eq!(served.load(Ordering::SeqCst), 3);
    }

    #[tokio::test(start_paused = true)]
    async fn a_re_read_token_the_forge_also_refuses_is_returned_as_the_error() {
        let (url, served) = forge_accepting("neither", None).await;
        let (token, reads) = token("old", "new");
        let api = Api::new(url, token).unwrap();
        let err = whoami(&api).await.expect_err("both tokens refused");
        assert!(refuses_the_token(&err), "{err}");
        assert_eq!(reads.load(Ordering::SeqCst), 1);
        assert_eq!(served.load(Ordering::SeqCst), 2);
    }

    #[tokio::test(start_paused = true)]
    async fn a_refused_token_still_in_the_store_is_not_sent_again() {
        let (url, served) = forge_accepting("new", None).await;
        let (token, reads) = token("old", "old");
        let api = Api::new(url, token).unwrap();
        let err = whoami(&api).await.expect_err("no other token to try");
        assert!(refuses_the_token(&err), "{err}");
        assert_eq!(reads.load(Ordering::SeqCst), 1);
        assert_eq!(served.load(Ordering::SeqCst), 1);
    }

    #[tokio::test(start_paused = true)]
    async fn concurrent_refusals_share_one_store_read() {
        let (url, served) = forge_accepting("new", None).await;
        let (token, reads) = token("old", "new");
        let api = Api::new(url, token).unwrap();
        let all = futures_util::future::join_all((0..5).map(|_| whoami(&api))).await;
        assert!(all.iter().all(Result::is_ok), "{all:?}");
        assert_eq!(reads.load(Ordering::SeqCst), 1);
        assert_eq!(served.load(Ordering::SeqCst), 10);
    }

    // Real time: the held request waits on a socket, and paused time would
    // jump the bound below while it does.
    #[tokio::test]
    async fn a_refusal_landing_after_the_re_read_is_sent_again_without_another_read() {
        let (url, served) = forge_accepting("new", Some(1)).await;
        let (token, reads) = token("old", "new");
        let api = Api::new(url, token).unwrap();
        let both = async { tokio::join!(whoami(&api), whoami(&api)) };
        let (a, b) = tokio::time::timeout(std::time::Duration::from_secs(10), both)
            .await
            .expect("the held request is answered once the other is accepted");
        a.expect("accepted with the re-read token");
        b.expect("accepted with the re-read token");
        assert_eq!(reads.load(Ordering::SeqCst), 1);
        assert_eq!(served.load(Ordering::SeqCst), 4);
    }

    #[test]
    fn only_a_401_counts_as_refusing_the_token() {
        let unauthorized = forgejo_api::ApiError {
            message: None,
            kind: ApiErrorKind::Unauthorized,
        };
        assert!(refuses_the_token(&ForgejoError::ApiError(unauthorized)));
        assert!(refuses_the_token(&ForgejoError::UnexpectedStatusCode(
            StatusCode::UNAUTHORIZED
        )));
        let other = forgejo_api::ApiError {
            message: None,
            kind: ApiErrorKind::Other(StatusCode::UNAUTHORIZED),
        };
        assert!(refuses_the_token(&ForgejoError::ApiError(other)));
        for refused in [StatusCode::BAD_REQUEST, StatusCode::FORBIDDEN] {
            assert!(!refuses_the_token(&ForgejoError::UnexpectedStatusCode(
                refused
            )));
        }
    }
}