sausage/host/src/cred/oauth2.rs

//! OAuth2 refresh exchange, endpoint policy, and a revision-aware process-local cache.
use super::{
    items::{Item, rule_matches},
    types::FieldType,
};
use std::{
    collections::{BTreeMap, BTreeSet},
    sync::Mutex,
    time::{Duration, Instant},
};
use zeroize::Zeroizing;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Error;
impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("OAuth refresh failed or returned an invalid response")
    }
}
impl std::error::Error for Error {}

/// A single HTTP exchange. Adapters must honor the deadline and never follow redirects.
pub trait Http: Sync {
    fn send_once(
        &self,
        request: &Request,
        deadline: Instant,
    ) -> impl std::future::Future<Output = Result<Response, Error>> + Send;
}
/// The status and body of one response, without any redirect exchange.
pub struct Response {
    pub status: http::StatusCode,
    pub body: Zeroizing<Vec<u8>>,
}
/// Contains credentials. Deliberately does not implement Debug or Serialize.
pub struct Request {
    pub url: String,
    pub method: &'static str,
    pub headers: Vec<(&'static str, String)>,
    pub body: Zeroizing<String>,
}
pub fn request(
    token_url: &str,
    refresh_token: &str,
    client_id: &str,
    client_secret: &str,
) -> Result<Request, Error> {
    let target = url::Url::parse(token_url).map_err(|_| Error)?;
    if target.scheme() != "https"
        || !target.username().is_empty()
        || target.password().is_some()
        || target.fragment().is_some()
        || [refresh_token, client_id, client_secret].contains(&"")
    {
        return Err(Error);
    }
    let body = Zeroizing::new(
        url::form_urlencoded::Serializer::new(String::new())
            .append_pair("grant_type", "refresh_token")
            .append_pair("refresh_token", refresh_token)
            .append_pair("client_id", client_id)
            .append_pair("client_secret", client_secret)
            .finish(),
    );
    Ok(Request {
        url: target.into(),
        method: "POST",
        headers: vec![
            ("Content-Type", "application/x-www-form-urlencoded".into()),
            ("Content-Length", body.len().to_string()),
        ],
        body,
    })
}
fn allowed(item: &Item, request: &Request) -> bool {
    let Ok(target) = url::Url::parse(&request.url) else {
        return false;
    };
    target.scheme() == "https"
        && target.username().is_empty()
        && target.password().is_none()
        && target.fragment().is_none()
        && request.method == "POST"
        && target.host_str().is_some_and(|host| {
            item.rules
                .iter()
                .any(|rule| rule_matches(rule, target.scheme(), host, target.path(), "POST"))
        })
}
#[derive(Clone)]
pub struct Token {
    pub value: Zeroizing<String>,
    pub scopes: Option<Vec<String>>,
    pub expires: Instant,
}
pub fn decode(bytes: &[u8]) -> Result<Token, Error> {
    #[derive(serde::Deserialize)]
    struct Response {
        access_token: Zeroizing<String>,
        token_type: String,
        expires_in: u64,
        scope: Option<String>,
    }
    let r: Response = serde_json::from_slice(bytes).map_err(|_| Error)?;
    // Validate header bytes without a second, non-zeroizing token allocation.
    if r.access_token.is_empty()
        || !r.token_type.eq_ignore_ascii_case("bearer")
        || r.expires_in == 0
        || r.access_token
            .bytes()
            .any(|b| b != b'\t' && (b < 32 || b == 127))
    {
        return Err(Error);
    }
    Ok(Token {
        value: r.access_token,
        scopes: r
            .scope
            .map(|s| s.split_whitespace().map(str::to_owned).collect()),
        expires: Instant::now() + Duration::from_secs(r.expires_in.min(86400).saturating_sub(30)),
    })
}
#[derive(Default)]
struct State {
    tokens: BTreeMap<(String, i64), Token>,
    busy: BTreeSet<String>,
    epoch: u64,
}
/// No mutex guard crosses HTTP await. Clear this cache on vault lock or item changes.
#[derive(Default)]
pub struct Cache {
    state: Mutex<State>,
}
struct Flight<'a> {
    cache: &'a Cache,
    item: String,
}
impl Drop for Flight<'_> {
    fn drop(&mut self) {
        if let Ok(mut state) = self.cache.state.lock() {
            state.busy.remove(&self.item);
        }
    }
}
impl Cache {
    pub fn clear(&self) -> Result<(), Error> {
        let mut state = self.state.lock().map_err(|_| Error)?;
        state.tokens.clear();
        state.epoch += 1;
        Ok(())
    }
    pub async fn refresh(
        &self,
        item: &Item,
        request: &Request,
        http: &impl Http,
        deadline: Instant,
    ) -> Result<Token, Error> {
        if !allowed(item, request) || Instant::now() >= deadline {
            return Err(Error);
        }
        let key = (item.id.clone(), item.version);
        let epoch;
        {
            let mut state = self.state.lock().map_err(|_| Error)?;
            if let Some(t) = state
                .tokens
                .get(&key)
                .filter(|t| t.expires > Instant::now())
            {
                return Ok(t.clone());
            }
            state.tokens.retain(|(id, _), _| id != &item.id);
            if !state.busy.insert(item.id.clone()) {
                return Err(Error);
            }
            epoch = state.epoch;
        }
        let _flight = Flight {
            cache: self,
            item: item.id.clone(),
        };
        let response = http.send_once(request, deadline).await?;
        if !response.status.is_success() {
            return Err(Error);
        }
        let token = decode(&response.body)?;
        let mut state = self.state.lock().map_err(|_| Error)?;
        if state.epoch != epoch || Instant::now() >= deadline {
            return Err(Error);
        }
        state.tokens.insert(key, token.clone());
        Ok(token)
    }

    /// Injection entry point. Replace the refresh field before any transform sees these values.
    /// The caller must reserve any additional grant use before this exchange.
    pub async fn fields_for_injection(
        &self,
        item: &Item,
        mut fields: BTreeMap<String, Zeroizing<String>>,
        http: &impl Http,
        deadline: Instant,
    ) -> Result<BTreeMap<String, Zeroizing<String>>, Error> {
        let mut refresh = item
            .fields
            .iter()
            .filter(|f| f.field_type == FieldType::Oauth2RefreshToken);
        let field = refresh.next().ok_or(Error)?;
        if field.public || refresh.next().is_some() {
            return Err(Error);
        }
        let value = |name: &str| {
            fields
                .get(name)
                .map(|s| s.as_str())
                .filter(|s| !s.is_empty())
                .ok_or(Error)
        };
        let req = request(
            value("token_url")?,
            value(&field.name)?,
            value("client_id")?,
            value("client_secret")?,
        )?;
        let token = self.refresh(item, &req, http, deadline).await?;
        fields.insert(field.name.clone(), token.value);
        Ok(fields)
    }
}