spiral-auth 0.1.10

use std::fmt::Debug;
use std::sync::Arc;
#[cfg(any(feature = "aws", feature = "gcp"))]
use std::time::Duration;

use async_trait::async_trait;
use parking_lot::RwLock;

#[cfg(feature = "aws")]
mod aws;
mod client;
#[cfg(feature = "device-auth")]
mod device;
#[cfg(feature = "gcp")]
mod gcp;
#[cfg(feature = "github")]
mod github;
#[cfg(feature = "kubernetes")]
mod kubernetes;
#[cfg(feature = "modal")]
mod modal;
mod oidc;
#[cfg(feature = "private-key-auth")]
mod private_key;
mod service;

#[cfg(feature = "aws")]
pub use aws::AwsAuthProvider;
pub use client::{ClientCredentials, ClientCredentialsProvider};
#[cfg(any(feature = "aws", feature = "gcp"))]
use cloud_detect::{ProviderId, detect};
#[cfg(feature = "device-auth")]
pub use device::{
    DEFAULT_DEVICE_AUTH, DeviceAuthBuilder, DeviceAuthn, DeviceCodeAuth, DeviceTokenStore,
    FileTokenStore, TokensModel,
};
#[cfg(feature = "gcp")]
pub use gcp::{GcpAuthProvider, GcpServiceAuthProvider};
#[cfg(feature = "github")]
pub use github::{GitHubAuthProvider, fetch_github_oidc_token};
#[cfg(feature = "kubernetes")]
pub use kubernetes::KubernetesAuthProvider;
#[cfg(feature = "modal")]
pub use modal::ModalAuthProvider;
pub use oidc::OidcAuthProvider;
#[cfg(feature = "private-key-auth")]
pub use private_key::{PrivateKeyCredentials, PrivateKeyJwtProvider};
use reqwest::Client;
pub use service::RefreshingAuthTokenProvider;
#[cfg(any(feature = "aws", feature = "gcp"))]
use tokio::{sync::OnceCell, time::timeout};
use url::Url;

use crate::http::Accepted as _;
use crate::{
    CredentialResult, OAuth2ClientAuthentication, OAuth2Grant, OAuth2TokenRequest,
    OAuth2TokenResponse, TOKEN_TYPE_ACCESS_TOKEN, TOKEN_TYPE_JWT, Token, TokenSource,
};

/// Spiral's authorization server, which signs users in for every stack.
///
/// `SPIRAL_AUTH_URL` overrides it; development mode uses the local server.
pub fn default_auth_url() -> Url {
    let configured = std::env::var("SPIRAL_AUTH_URL")
        .ok()
        .and_then(|url| Url::parse(&url).ok());
    configured.unwrap_or_else(|| {
        let url = if std::env::var("SPIRAL__DEV").is_ok_and(|value| !value.is_empty()) {
            "http://localhost:4280"
        } else {
            "https://auth.spiraldb.com"
        };
        Url::parse(url).expect("valid default authorization server URL")
    })
}

#[cfg(any(feature = "aws", feature = "gcp"))]
static CLOUD_PROVIDER_ID: OnceCell<Option<ProviderId>> = OnceCell::const_new();

#[cfg(any(feature = "aws", feature = "gcp"))]
async fn cloud_provider_id() -> &'static Option<ProviderId> {
    CLOUD_PROVIDER_ID
        .get_or_init(|| async {
            let duration = Duration::from_millis(
                std::env::var("SPIRAL__CLOUD_DETECT_TIMEOUT_MS")
                    .ok()
                    .and_then(|x| x.parse::<u64>().ok())
                    .unwrap_or(50),
            );
            timeout(duration, detect(Some(duration))).await.ok()
        })
        .await
}

/// Resolve the workload identity from the SPIRAL_WORKLOAD_ID environment variable.
/// Returns None if the variable is unset or empty.
pub fn resolve_workload_id() -> Option<String> {
    std::env::var("SPIRAL_WORKLOAD_ID")
        .ok()
        .filter(|s| !s.is_empty())
}

/// The token endpoint of the authorization server at `auth_url`.
pub fn token_endpoint(auth_url: &Url) -> CredentialResult<Url> {
    auth_url
        .join("/oauth/token")
        .map_err(|error| invalid!("invalid url {error}"))
}

/// RFC 7523 §2.1: obtain the token of the client `client_id` by presenting `assertion`, a JWT
/// from the client's platform whose audience is `token_endpoint`, as the JWT-bearer grant.
pub(crate) async fn jwt_bearer(
    client: &Client,
    token_endpoint: &Url,
    client_id: String,
    assertion: String,
    resource: Option<String>,
) -> CredentialResult<Token> {
    let response = request_token(
        client,
        token_endpoint,
        &OAuth2TokenRequest {
            client: OAuth2ClientAuthentication {
                client_id: Some(client_id),
                ..Default::default()
            },
            grant: OAuth2Grant::JwtBearer {
                assertion,
                resource,
                authorization_details: None,
            },
        },
    )
    .await?;
    Ok(response.into_access_token())
}

/// A registered client that obtains its own tokens with RFC 9396 authorization details: the
/// grant that the client authorized, for a resource server to enforce.
#[async_trait]
pub trait ClientGrant: Send + Sync {
    /// A token for `resource` that carries `authorization_details`, a JSON array.
    async fn grant(
        &self,
        resource: &str,
        authorization_details: &serde_json::Value,
    ) -> CredentialResult<OAuth2TokenResponse>;
}

/// The token endpoint's response to `request`.
pub(crate) async fn request_token(
    client: &Client,
    token_endpoint: &Url,
    request: &OAuth2TokenRequest,
) -> CredentialResult<OAuth2TokenResponse> {
    Ok(client
        .post(token_endpoint.clone())
        .form(request)
        .send()
        .await?
        .accepted()
        .await?
        .json::<OAuth2TokenResponse>()
        .await?)
}

/// RFC 8693 §2.1: narrow `subject`, an access token from the authorization server at
/// `token_endpoint`, to `resources` (RFC 8707 §2 lets a request name several). The narrowed
/// token keeps the subject's claims and expiry, and names only `resources`, which the subject
/// must name.
pub async fn exchange_token(
    client: &Client,
    token_endpoint: &Url,
    subject: &Token,
    resources: &[&str],
) -> CredentialResult<Token> {
    let mut form = vec![
        (
            "grant_type",
            "urn:ietf:params:oauth:grant-type:token-exchange",
        ),
        ("subject_token", subject.expose_secret()),
        ("subject_token_type", TOKEN_TYPE_ACCESS_TOKEN),
    ];
    form.extend(resources.iter().map(|resource| ("resource", *resource)));
    let response = client
        .post(token_endpoint.clone())
        .form(&form)
        .send()
        .await?
        .accepted()
        .await?
        .json::<OAuth2TokenResponse>()
        .await?;
    Ok(response.into_access_token())
}

/// RFC 8693 §2.1: exchange `subject`, an access token from the authorization server at
/// `token_endpoint`, for a JWT for `audience`, which is not a Spiral API, such as another
/// cloud's token service. The RFC 9396 `authorization_details` say what the JWT speaks for.
pub async fn exchange_for_jwt(
    client: &Client,
    token_endpoint: &Url,
    subject: &Token,
    audience: &str,
    authorization_details: &serde_json::Value,
) -> CredentialResult<Token> {
    let response = client
        .post(token_endpoint.clone())
        .form(&OAuth2TokenRequest {
            client: OAuth2ClientAuthentication::default(),
            grant: OAuth2Grant::TokenExchange {
                subject_token: subject.expose_secret().to_owned(),
                subject_token_type: TOKEN_TYPE_ACCESS_TOKEN.to_owned(),
                resource: None,
                audience: Some(audience.to_owned()),
                requested_token_type: Some(TOKEN_TYPE_JWT.to_owned()),
                authorization_details: Some(authorization_details.to_string()),
            },
        })
        .send()
        .await?
        .accepted()
        .await?
        .json::<OAuth2TokenResponse>()
        .await?;
    Ok(response.access_token)
}

/// Groups all workload-identity providers and warns if none of them detected a valid token.
///
/// This lets the outer FallbackAuthn continue to DeviceAuthn while still surfacing a diagnostic
/// when SPIRAL_WORKLOAD_ID is set but the environment lacks a usable platform token.
///
/// When `workload_id` is `None` (i.e. `SPIRAL_WORKLOAD_ID` is not set), the providers are
/// skipped entirely. If a workload ID is present but no provider returns a token, an error
/// is returned.
pub struct WorkloadAuthn {
    workload_id: Option<String>,
    providers: Vec<TokenSource>,
}

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

impl WorkloadAuthn {
    pub fn new(workload_id: Option<String>, providers: Vec<TokenSource>) -> Self {
        WorkloadAuthn {
            workload_id,
            providers,
        }
    }
}

impl WorkloadAuthn {
    /// A usable token, or `None` when this provider has no credentials to offer.
    pub async fn token(&self) -> CredentialResult<Option<Token>> {
        if self.workload_id.is_none() {
            // SPIRAL_OIDC_PROVIDER_TOKEN requires SPIRAL_WORKLOAD_ID to identify the principal
            // for token exchange. Catch this misconfiguration early.
            let has_oidc_token =
                std::env::var("SPIRAL_OIDC_PROVIDER_TOKEN").is_ok_and(|s| !s.is_empty());
            if has_oidc_token {
                bail!(
                    "SPIRAL_OIDC_PROVIDER_TOKEN is set but SPIRAL_WORKLOAD_ID is not. \
                     Both are required for custom OIDC authentication."
                );
            }
            return Ok(None);
        }

        for provider in &self.providers {
            if let Some(token) = provider.token().await? {
                return Ok(Some(token));
            }
        }

        bail!("SPIRAL_WORKLOAD_ID is set but no workload identity provider detected a valid token")
    }
}

impl_token_source!(WorkloadAuthn);

struct CachedToken {
    token: Token,
    provider: TokenSource,
}

pub struct FallbackAuthn {
    providers: Arc<[TokenSource]>,
    cached: RwLock<Option<CachedToken>>,
}

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

impl FallbackAuthn {
    pub fn new(providers: Arc<[TokenSource]>) -> Self {
        FallbackAuthn {
            providers,
            cached: RwLock::new(None),
        }
    }
}

impl FallbackAuthn {
    /// A usable token, or `None` when this provider has no credentials to offer.
    pub async fn token(&self) -> CredentialResult<Option<Token>> {
        {
            let cached = self.cached.read();

            // If token doesn't need refresh yet, return it
            if let Some(cached_token) = cached.as_ref()
                && !cached_token.token.needs_refresh()?
            {
                return Ok(Some(cached_token.token.clone()));
            }
        }

        // Always use previously successful provider, if it exists
        let cached_provider = self.cached.read().as_ref().map(|c| c.provider.clone());
        if let Some(provider) = cached_provider {
            if let Some(token) = provider.token().await? {
                let mut cached = self.cached.write();
                *cached = Some(CachedToken {
                    token: token.clone(),
                    provider,
                });

                return Ok(Some(token));
            }

            bail!("Previously resolved authn provider return no token");
        }

        // If we never succeeded, try all providers.
        for provider in self.providers.iter() {
            tracing::debug!(?provider, "trying authn provider");
            if let Some(token) = provider.token().await? {
                let mut cached = self.cached.write();
                *cached = Some(CachedToken {
                    token: token.clone(),
                    provider: provider.clone(),
                });
                return Ok(Some(token));
            }
        }

        Ok(None)
    }

    /// A token this provider already holds, without network requests or prompts.
    pub async fn cached_token(&self) -> CredentialResult<Option<Token>> {
        let cached = self.cached.read();
        if let Some(ct) = cached.as_ref()
            && !ct.token.is_expired()?
        {
            return Ok(Some(ct.token.clone()));
        }
        Ok(None)
    }
}

impl_token_source!(FallbackAuthn, cached);

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

    use serde_json::json;
    use serial_test::serial;

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

    fn create_test_token(expires_in_seconds: i64) -> Token {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;
        unsigned_jwt(json!({ "sub": "test_user", "exp": now + expires_in_seconds, "iat": now }))
    }

    struct MockAuthn {
        token: Option<Token>,
        call_count: AtomicUsize,
    }

    impl MockAuthn {
        fn new(token: Option<Token>) -> Self {
            Self {
                token,
                call_count: AtomicUsize::new(0),
            }
        }

        fn call_count(&self) -> usize {
            self.call_count.load(Ordering::SeqCst)
        }
    }

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

    impl MockAuthn {
        /// A usable token, or `None` when this provider has no credentials to offer.
        pub async fn token(&self) -> CredentialResult<Option<Token>> {
            self.call_count.fetch_add(1, Ordering::SeqCst);
            Ok(self.token.clone())
        }
    }

    impl_token_source!(Arc<MockAuthn>);

    #[tokio::test]
    async fn test_fallback_caches_token() {
        // Create a mock token that expires in 1 hour (won't trigger refresh)
        let token = create_test_token(3600);

        // Verify the token is not expired
        assert!(!token.is_expired().unwrap(), "Token should not be expired");

        let mock1 = Arc::new(MockAuthn::new(None));
        let mock2 = Arc::new(MockAuthn::new(Some(token.clone())));
        let mock3 = Arc::new(MockAuthn::new(None));

        let providers: Arc<[TokenSource]> =
            Arc::new([mock1.clone().into(), mock2.clone().into(), mock3.into()]);

        let fallback = FallbackAuthn::new(providers);

        // First call should try providers and cache result
        let result1 = fallback.token().await.unwrap();
        assert!(result1.is_some());
        assert_eq!(mock1.call_count(), 1);
        assert_eq!(mock2.call_count(), 1);

        // Second call should use cached token without calling providers
        let result2 = fallback.token().await.unwrap();
        assert!(result2.is_some());
        assert_eq!(mock1.call_count(), 1);
        assert_eq!(mock2.call_count(), 1);

        // Token values should match
        assert_eq!(
            result1.unwrap().expose_secret(),
            result2.unwrap().expose_secret()
        );
    }

    #[tokio::test]
    async fn test_fallback_remembers_successful_provider() {
        // Create an expired token (expired 1 hour ago) and a fresh token
        let expired_token = create_test_token(-3600);
        let fresh_token = create_test_token(3600);

        let mock1 = Arc::new(MockAuthn::new(None));
        let mock2 = Arc::new(MockAuthn::new(Some(fresh_token.clone())));

        let providers: Arc<[TokenSource]> = Arc::new([mock1.clone().into(), mock2.clone().into()]);

        let fallback = FallbackAuthn::new(providers);

        // First call
        fallback.token().await.unwrap();
        assert_eq!(mock1.call_count(), 1);
        assert_eq!(mock2.call_count(), 1);

        // Replace cached token with an expired one to simulate token expiration
        {
            let mut cached = fallback.cached.write();
            *cached = Some(CachedToken {
                token: expired_token,
                provider: mock2.clone().into(),
            });
        }

        // Second call should detect expired token and try mock2 first (the previously successful
        // provider)
        fallback.token().await.unwrap();
        // mock1 should not be called again since we remember mock2 succeeded
        assert_eq!(mock1.call_count(), 1);
        assert_eq!(mock2.call_count(), 2);
    }

    #[tokio::test]
    async fn test_fallback_no_providers_succeed() {
        let mock1 = Arc::new(MockAuthn::new(None));
        let mock2 = Arc::new(MockAuthn::new(None));

        let providers: Arc<[TokenSource]> = Arc::new([mock1.clone().into(), mock2.clone().into()]);

        let fallback = FallbackAuthn::new(providers);

        let result = fallback.token().await.unwrap();
        assert!(result.is_none());
        assert_eq!(mock1.call_count(), 1);
        assert_eq!(mock2.call_count(), 1);
    }

    #[tokio::test]
    async fn test_fail_if_cant_refresh() {
        #[derive(Debug)]
        struct SingleUseAuthn {
            counter: AtomicBool,
        }

        impl SingleUseAuthn {
            fn new() -> Self {
                Self {
                    counter: AtomicBool::new(true),
                }
            }
        }

        impl SingleUseAuthn {
            /// A usable token, or `None` when this provider has no credentials to offer.
            pub async fn token(&self) -> CredentialResult<Option<Token>> {
                if self.counter.fetch_and(false, Ordering::SeqCst) {
                    Ok(Some(create_test_token(1)))
                } else {
                    Ok(None)
                }
            }
        }

        impl_token_source!(SingleUseAuthn);

        let single_use = TokenSource::from(SingleUseAuthn::new());
        let mock = Arc::new(MockAuthn::new(None));
        let fallback = FallbackAuthn::new([mock.clone().into(), single_use].into());

        let result = fallback.token().await.unwrap();
        assert!(result.is_some());
        assert_eq!(mock.call_count(), 1);

        tokio::time::sleep(Duration::from_secs(1)).await;

        let second = fallback.token().await;
        assert_eq!(mock.call_count(), 1);
        assert!(second.is_err());
    }

    #[test]
    #[serial]
    fn test_resolve_workload_id() {
        // SAFETY: Test is marked #[serial] so no concurrent env var access.
        unsafe { std::env::remove_var("SPIRAL_WORKLOAD_ID") };

        // Unset env var returns None.
        assert_eq!(resolve_workload_id(), None);

        // Env var is returned when set.
        unsafe { std::env::set_var("SPIRAL_WORKLOAD_ID", "from-env") };
        assert_eq!(resolve_workload_id(), Some("from-env".to_string()));

        // Empty env var is treated as unset.
        unsafe { std::env::set_var("SPIRAL_WORKLOAD_ID", "") };
        assert_eq!(resolve_workload_id(), None);

        // Clean up.
        unsafe { std::env::remove_var("SPIRAL_WORKLOAD_ID") };
    }
}