babelforce-manager-sdk 0.44.0

//! Transparent bearer-token lifecycle for grant-based [`Auth`]: the connect-time token is kept
//! fresh — re-granted (single-flight) when a call would use it within [`EXPIRY_MARGIN`] of
//! expiry — and the per-spec configurations are rebuilt to carry the current token.

use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};

use crate::auth::Auth;
use crate::error::ManagerError;

/// Re-grant this long before nominal expiry (mirrors the TS/Go SDKs' 30-second early refresh).
const EXPIRY_MARGIN: Duration = Duration::from_secs(30);

/// The `/oauth/token` response fields the token lifecycle needs.
#[derive(serde::Deserialize)]
struct TokenResponse {
    access_token: String,
    #[serde(default)]
    expires_in: Option<i64>,
    #[serde(default)]
    refresh_token: Option<String>,
}

struct TokenState {
    /// Bumped on every re-grant; lets a [`SharedCfg`] detect that its cached config is stale.
    epoch: u64,
    token: String,
    /// `None` — the token is treated as non-expiring (no positive `expires_in` in the response).
    expires_at: Option<Instant>,
    /// The latest rotated refresh token (refresh-token grants rotate it on every use).
    refresh_token: Option<String>,
}

impl TokenState {
    fn apply(&mut self, resp: TokenResponse) {
        self.epoch += 1;
        self.token = resp.access_token;
        // A missing or non-positive `expires_in` means the token never triggers a re-grant.
        self.expires_at = match resp.expires_in {
            Some(secs) if secs > 0 => Some(Instant::now() + Duration::from_secs(secs as u64)),
            _ => None,
        };
        if resp.refresh_token.is_some() {
            self.refresh_token = resp.refresh_token;
        }
    }

    fn near_expiry(&self) -> bool {
        self.expires_at
            .is_some_and(|at| Instant::now() + EXPIRY_MARGIN >= at)
    }
}

/// Resolves and renews the bearer token for the grant-based [`Auth`] flavours (password,
/// client-credentials, refresh-token). [`Auth::Bearer`] never constructs one — a static token is
/// used as-is for the client's lifetime.
pub(crate) struct TokenManager {
    base: String,
    auth: Auth,
    http: reqwest_middleware::ClientWithMiddleware,
    state: tokio::sync::Mutex<TokenState>,
}

impl TokenManager {
    /// Resolve the connect-time grant eagerly (fail fast) and return a manager that keeps the
    /// token fresh. Returns `None` for [`Auth::Bearer`].
    pub(crate) async fn connect(
        base: &str,
        auth: &Auth,
        http: reqwest_middleware::ClientWithMiddleware,
    ) -> Result<Option<Arc<TokenManager>>, ManagerError> {
        if matches!(auth, Auth::Bearer { .. }) {
            return Ok(None);
        }
        let mut mgr = TokenManager {
            base: base.to_string(),
            auth: auth.clone(),
            http,
            state: tokio::sync::Mutex::new(TokenState {
                epoch: 0,
                token: String::new(),
                expires_at: None,
                refresh_token: None,
            }),
        };
        let resp = grant(&mgr.base, &mgr.auth, &mgr.http, None).await?;
        mgr.state.get_mut().apply(resp);
        Ok(Some(Arc::new(mgr)))
    }

    /// The current `(epoch, token)`, re-granting first when the token is within
    /// [`EXPIRY_MARGIN`] of expiry. The async lock is held across the re-grant, so concurrent
    /// callers needing a refresh await the same in-flight token request (single-flight) — and
    /// their futures stay cancellable while they wait.
    pub(crate) async fn current(&self) -> Result<(u64, String), ManagerError> {
        let mut st = self.state.lock().await;
        if st.near_expiry() {
            let rotated = st.refresh_token.clone();
            let resp = grant(&self.base, &self.auth, &self.http, rotated.as_deref()).await?;
            st.apply(resp);
        }
        Ok((st.epoch, st.token.clone()))
    }
}

/// POST an OAuth2 grant to `{base}/oauth/token` and parse the token response. For
/// [`Auth::RefreshToken`], `rotated_refresh_token` (the one issued by the previous grant)
/// replaces the connect-time seed once available.
async fn grant(
    base: &str,
    auth: &Auth,
    http: &reqwest_middleware::ClientWithMiddleware,
    rotated_refresh_token: Option<&str>,
) -> Result<TokenResponse, ManagerError> {
    let params: Vec<(&str, &str)> = match auth {
        Auth::Password {
            user,
            pass,
            client_id,
        } => vec![
            ("grant_type", "password"),
            ("username", user),
            ("password", pass),
            ("client_id", client_id.as_deref().unwrap_or("manager")),
        ],
        Auth::ClientCredentials {
            client_id,
            client_secret,
        } => vec![
            ("grant_type", "client_credentials"),
            ("client_id", client_id),
            ("client_secret", client_secret),
        ],
        Auth::RefreshToken {
            refresh_token,
            client_id,
            client_secret,
        } => {
            let mut params = vec![
                ("grant_type", "refresh_token"),
                (
                    "refresh_token",
                    rotated_refresh_token.unwrap_or(refresh_token),
                ),
                ("client_id", client_id),
            ];
            // Confidential clients also present their secret on every exchange (mirrors TS).
            if let Some(secret) = client_secret {
                params.push(("client_secret", secret));
            }
            params
        }
        Auth::Bearer { .. } => {
            return Err(ManagerError::InvalidArgument(
                "static bearer auth has no grant to resolve".into(),
            ))
        }
    };
    // The token POST rides the shared retry middleware like every other request (non-idempotent,
    // so it is retried only on 429) — matching the TS/Go SDKs, whose token fetches go through
    // their retry-wrapped fetch/transport.
    let resp = http
        .post(format!("{base}/oauth/token"))
        .form(&params)
        .send()
        .await
        .map_err(|e| ManagerError::Network(e.to_string()))?;
    let status = resp.status().as_u16();
    let body = resp
        .text()
        .await
        .map_err(|e| ManagerError::Network(e.to_string()))?;
    if !(200..300).contains(&status) {
        return Err(ManagerError::from_response(status, body));
    }
    serde_json::from_str(&body).map_err(|e| ManagerError::Decode(e.to_string()))
}

/// Rebuild a generated per-spec `Configuration` with a new bearer token. Implemented for each
/// spec's (structurally identical) config type.
pub(crate) trait TokenCfg: Clone {
    fn with_token(&self, token: String) -> Self;
}

macro_rules! impl_token_cfg {
    ($($t:ty),+ $(,)?) => {$(
        impl TokenCfg for $t {
            fn with_token(&self, token: String) -> Self {
                let mut cfg = self.clone();
                cfg.oauth_access_token = Some(token);
                cfg
            }
        }
    )+};
}

impl_token_cfg!(
    crate::gen::manager::apis::configuration::Configuration,
    crate::gen::user::apis::configuration::Configuration,
    crate::gen::auth::apis::configuration::Configuration,
    crate::gen::task_automation::apis::configuration::Configuration,
    crate::gen::task_schedule::apis::configuration::Configuration,
);

/// A per-spec configuration cell that resources consult before every call: hands out the current
/// config, rebuilding it (token swap) once the [`TokenManager`] has re-granted.
pub(crate) struct SharedCfg<C> {
    inner: Arc<Inner<C>>,
}

struct Inner<C> {
    manager: Option<Arc<TokenManager>>,
    /// The token epoch the cached config was built for, and the config itself.
    cell: RwLock<(u64, Arc<C>)>,
}

impl<C> Clone for SharedCfg<C> {
    fn clone(&self) -> Self {
        SharedCfg {
            inner: self.inner.clone(),
        }
    }
}

impl<C: TokenCfg> SharedCfg<C> {
    /// `epoch` must be the [`TokenManager::current`] epoch whose token `cfg` carries (0 when
    /// there is no manager).
    pub(crate) fn new(cfg: C, manager: Option<Arc<TokenManager>>, epoch: u64) -> Self {
        SharedCfg {
            inner: Arc::new(Inner {
                manager,
                cell: RwLock::new((epoch, Arc::new(cfg))),
            }),
        }
    }

    /// The configuration for the next call — carrying a fresh token when the current one is near
    /// expiry. Static ([`Auth::Bearer`]) configs are returned as-is.
    pub(crate) async fn get(&self) -> Result<Arc<C>, ManagerError> {
        let Some(manager) = &self.inner.manager else {
            return Ok(self.read().1);
        };
        let (epoch, token) = manager.current().await?;
        let (cached_epoch, cached) = self.read();
        if cached_epoch == epoch {
            return Ok(cached);
        }
        let mut cell = self.inner.cell.write().expect("cfg cell poisoned");
        if cell.0 != epoch {
            *cell = (epoch, Arc::new(cell.1.with_token(token)));
        }
        Ok(cell.1.clone())
    }

    fn read(&self) -> (u64, Arc<C>) {
        let cell = self.inner.cell.read().expect("cfg cell poisoned");
        (cell.0, cell.1.clone())
    }
}