spiral-auth 0.1.12

use std::fmt::{Debug, Formatter};
use std::time::Duration;

use async_trait::async_trait;
use cloud_detect::ProviderId;
use reqwest_middleware::ClientWithMiddleware as Client;
use tokio::sync::Mutex;
use tokio::time::sleep;
use tracing::{debug, warn};
use url::Url;

use super::{ClientGrant, cloud_provider_id, jwt_bearer, request_token, token_endpoint};
use crate::http::Accepted as _;
use crate::{
    CredentialError, CredentialResult, OAuth2ClientAuthentication, OAuth2Grant, OAuth2TokenRequest,
    OAuth2TokenResponse, Token,
};

/// Authenticates a workload with its GCP host's metadata-server ID token, as a JWT-bearer
/// grant at the authorization server.
/// <https://cloud.google.com/docs/authentication/get-id-token>
///
/// You can get an ID token from the metadata server when your code is running on the following
/// Google Cloud services:
///     - Compute Engine
///     - App Engine standard environment
///     - App Engine flexible environment
///     - Cloud Run functions
///     - Cloud Run
///     - Google Kubernetes Engine
///     - Cloud Build
pub struct GcpAuthProvider {
    client: Client,
    token_endpoint: Url,
    workload_id: String,
    resource: Option<String>,
}

impl Debug for GcpAuthProvider {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GcpAuthProvider").finish()
    }
}

impl GcpAuthProvider {
    /// Obtain the workload's tokens at `token_endpoint`, the authorization server's.
    pub fn new(client: impl Into<Client>, token_endpoint: Url, workload_id: String) -> Self {
        GcpAuthProvider {
            client: client.into(),
            token_endpoint,
            workload_id,
            resource: None,
        }
    }

    /// Request a token for the API resource named by `resource` (RFC 8707).
    pub fn with_resource(mut self, resource: impl Into<String>) -> Self {
        self.resource = Some(resource.into());
        self
    }
}

impl GcpAuthProvider {
    /// A usable token, or `None` when this provider has no credentials to offer.
    pub async fn token(&self) -> CredentialResult<Option<Token>> {
        if !matches!(cloud_provider_id().await, Some(ProviderId::GCP)) {
            return Ok(None);
        }

        // RFC 7523 §3: the assertion's audience is the token endpoint it is presented to.
        let gcp_token = metadata_identity_token(&self.client, self.token_endpoint.as_str()).await?;
        let token = jwt_bearer(
            &self.client,
            &self.token_endpoint,
            self.workload_id.clone(),
            gcp_token.expose_secret().to_string(),
            self.resource.clone(),
        )
        .await?;

        debug!("GcpAuthProvider obtained a workload token");
        Ok(Some(token))
    }
}

impl_token_source!(GcpAuthProvider);

/// The metadata server's Google ID tokens.
const METADATA_IDENTITY_URL: &str =
    "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity";

/// How long before its expiry a host ID token is replaced. A token the authorization server
/// issues for an assertion ends no later than the assertion, and lasts at most 15 minutes, so an
/// assertion with more than this left yields full-length tokens.
const IDENTITY_RENEWAL_MARGIN: Duration = Duration::from_secs(20 * 60);

/// The pauses between attempts at a metadata request that failed transiently. GKE's metadata
/// server answers 499 when its own ID-token request is cancelled, and Google's clients retry it.
const METADATA_RETRY_DELAYS: [Duration; 2] =
    [Duration::from_millis(100), Duration::from_millis(400)];

/// Authenticates one registered service client with the GCP host identity.
pub struct GcpServiceAuthProvider {
    client: Client,
    token_endpoint: Url,
    client_id: String,
    resource: Option<String>,
    identity: HostIdentity,
}

/// This host's Google ID token for one audience, fetched once and shared until it nears expiry.
/// RFC 7523 lets a JWT-bearer assertion be presented until it expires, and the metadata server
/// cancels some requests when every grant asks it for a new token.
struct HostIdentity {
    client: Client,
    url: Url,
    current: Mutex<Option<Token>>,
}

impl HostIdentity {
    fn new(client: Client, metadata_url: &str, audience: &str) -> CredentialResult<Self> {
        Ok(Self {
            client,
            url: identity_url(metadata_url, audience)?,
            current: Mutex::new(None),
        })
    }

    /// The current token, renewed once it nears expiry. A failed renewal presents the current
    /// token while it remains usable, and the next call tries again.
    async fn token(&self) -> CredentialResult<Token> {
        let mut current = self.current.lock().await;
        if let Some(token) = current.as_ref()
            && token.remaining_lifetime()? > IDENTITY_RENEWAL_MARGIN
        {
            return Ok(token.clone());
        }
        match fetch_identity_token(&self.client, &self.url).await {
            Ok(token) => {
                *current = Some(token.clone());
                Ok(token)
            }
            Err(error) => match current.as_ref() {
                Some(token) if !token.is_expired()? => {
                    warn!(%error, "renewing the GCP host identity failed; presenting the current one");
                    Ok(token.clone())
                }
                _ => Err(error),
            },
        }
    }
}

impl GcpServiceAuthProvider {
    /// Obtain tokens for the client `client_id` from the authorization server at `auth_url`,
    /// presenting this GCP host's Google ID token as an RFC 7523 JWT-bearer grant. `resource`
    /// names the RFC 8707 resource the token is for.
    pub fn new(
        client: impl Into<Client>,
        auth_url: &Url,
        client_id: String,
        resource: Option<String>,
    ) -> CredentialResult<Self> {
        let client = client.into();
        let token_endpoint = token_endpoint(auth_url)?;
        // RFC 7523 §3: the assertion's audience is the token endpoint it is presented to.
        let identity = HostIdentity::new(
            client.clone(),
            METADATA_IDENTITY_URL,
            token_endpoint.as_str(),
        )?;
        Ok(Self {
            client,
            token_endpoint,
            client_id,
            resource,
            identity,
        })
    }

    /// A token for the service client.
    pub async fn token(&self) -> CredentialResult<Option<Token>> {
        // Servers build this provider only on GCP, so it does not probe for the cloud: the probe
        // caches its first answer, and one that times out at startup would leave the service
        // without tokens for the life of its process.
        let assertion = self.identity.token().await?;
        let token = jwt_bearer(
            &self.client,
            &self.token_endpoint,
            self.client_id.clone(),
            assertion.expose_secret().to_string(),
            self.resource.clone(),
        )
        .await?;

        debug!(client_id = %self.client_id, "exchanged GCP host identity for service token");
        Ok(Some(token))
    }
}

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

impl_token_source!(GcpServiceAuthProvider);

#[async_trait]
impl ClientGrant for GcpServiceAuthProvider {
    async fn grant(
        &self,
        resource: &str,
        authorization_details: &serde_json::Value,
    ) -> CredentialResult<OAuth2TokenResponse> {
        let assertion = self.identity.token().await?;
        let request = OAuth2TokenRequest {
            client: OAuth2ClientAuthentication {
                client_id: Some(self.client_id.clone()),
                ..Default::default()
            },
            grant: OAuth2Grant::JwtBearer {
                assertion: assertion.expose_secret().to_owned(),
                resource: Some(resource.to_owned()),
                authorization_details: Some(authorization_details.to_string()),
            },
        };
        request_token(&self.client, &self.token_endpoint, &request).await
    }
}

async fn metadata_identity_token(client: &Client, audience: &str) -> CredentialResult<Token> {
    fetch_identity_token(client, &identity_url(METADATA_IDENTITY_URL, audience)?).await
}

/// The metadata server's ID-token URL for `audience`.
fn identity_url(metadata_url: &str, audience: &str) -> CredentialResult<Url> {
    let mut url = Url::parse(metadata_url)
        .map_err(|error| invalid!("invalid metadata identity URL: {error}"))?;
    url.query_pairs_mut().append_pair("audience", audience);
    Ok(url)
}

/// An ID token from the metadata server at `url`, retrying transient failures.
async fn fetch_identity_token(client: &Client, url: &Url) -> CredentialResult<Token> {
    let mut delays = METADATA_RETRY_DELAYS.into_iter();
    loop {
        match request_identity_token(client, url).await {
            Err(error) if is_transient(&error) => match delays.next() {
                Some(delay) => {
                    debug!(%error, "retrying a transient metadata server failure");
                    sleep(delay).await;
                }
                None => return Err(error),
            },
            result => return result,
        }
    }
}

async fn request_identity_token(client: &Client, url: &Url) -> CredentialResult<Token> {
    let response = client
        .get(url.clone())
        .header("Metadata-Flavor", "Google")
        .send()
        .await?
        .accepted()
        .await?
        .text()
        .await?;
    Ok(Token::from(response))
}

/// Whether a metadata request may succeed when repeated: the server cancelled it (499), was
/// overloaded or failed, or could not be reached in time.
fn is_transient(error: &CredentialError) -> bool {
    match error {
        CredentialError::Rejected { status, .. } => {
            status.as_u16() == 499 || status.as_u16() == 429 || status.is_server_error()
        }
        CredentialError::Http(error) => error.is_timeout() || error.is_connect(),
        _ => false,
    }
}

#[cfg(test)]
mod tests {
    use std::time::{SystemTime, UNIX_EPOCH};

    use futures::future::join_all;
    use serde_json::json;
    use wiremock::matchers::{header, method, path, query_param};
    use wiremock::{Mock, MockServer, ResponseTemplate};

    use super::*;
    use crate::http::test_client;
    use crate::token::unsigned_jwt;

    const AUDIENCE: &str = "https://auth.spiraldb.test/oauth/token";

    fn id_token(lifetime: Duration) -> Token {
        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
        unsigned_jwt(json!({ "exp": (now + lifetime).as_secs() }))
    }

    async fn serve_one_id_token(metadata: &MockServer, token: &Token) {
        metadata.reset().await;
        Mock::given(method("GET"))
            .and(path("/identity"))
            .and(header("Metadata-Flavor", "Google"))
            .and(query_param("audience", AUDIENCE))
            .respond_with(
                ResponseTemplate::new(200)
                    .set_body_string(token.expose_secret())
                    .set_delay(Duration::from_millis(50)),
            )
            .expect(1)
            .mount(metadata)
            .await;
    }

    /// Concurrent grants share one metadata request, and a token within the renewal margin is
    /// replaced before it is presented again.
    #[tokio::test]
    async fn the_host_identity_is_fetched_once_and_renewed_before_it_expires() {
        let metadata = MockServer::start().await;
        let identity = HostIdentity::new(
            test_client(),
            &format!("{}/identity", metadata.uri()),
            AUDIENCE,
        )
        .unwrap();

        let fresh = id_token(Duration::from_secs(3600));
        serve_one_id_token(&metadata, &fresh).await;
        for token in join_all((0..20).map(|_| identity.token())).await {
            assert_eq!(token.unwrap().expose_secret(), fresh.expose_secret());
        }
        metadata.verify().await;

        *identity.current.lock().await = Some(id_token(IDENTITY_RENEWAL_MARGIN / 2));
        let renewed = id_token(Duration::from_secs(3600));
        serve_one_id_token(&metadata, &renewed).await;
        assert_eq!(
            identity.token().await.unwrap().expose_secret(),
            renewed.expose_secret()
        );
        metadata.verify().await;
    }

    /// A cancelled metadata request is retried, and a renewal that keeps failing presents the
    /// current token until it expires.
    #[tokio::test]
    async fn a_transient_metadata_failure_is_retried_and_renewal_falls_back() {
        let metadata = MockServer::start().await;
        let identity = HostIdentity::new(
            test_client(),
            &format!("{}/identity", metadata.uri()),
            AUDIENCE,
        )
        .unwrap();
        let cancelled = || {
            Mock::given(method("GET"))
                .and(path("/identity"))
                .respond_with(ResponseTemplate::new(499).set_body_string("context canceled"))
        };

        let fresh = id_token(Duration::from_secs(3600));
        cancelled().up_to_n_times(1).mount(&metadata).await;
        Mock::given(method("GET"))
            .and(path("/identity"))
            .respond_with(ResponseTemplate::new(200).set_body_string(fresh.expose_secret()))
            .mount(&metadata)
            .await;
        assert_eq!(
            identity.token().await.unwrap().expose_secret(),
            fresh.expose_secret()
        );

        metadata.reset().await;
        cancelled().mount(&metadata).await;
        let renewing = id_token(IDENTITY_RENEWAL_MARGIN / 2);
        *identity.current.lock().await = Some(renewing.clone());
        assert_eq!(
            identity.token().await.unwrap().expose_secret(),
            renewing.expose_secret()
        );

        *identity.current.lock().await = Some(id_token(Duration::from_secs(10)));
        assert!(identity.token().await.is_err());
    }
}