OpenShell/crates/openshell-bootstrap/src/oidc_token.rs

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! OIDC token storage.
//!
//! Stores OIDC token bundles (access token, refresh token, metadata) at
//! `$XDG_CONFIG_HOME/openshell/gateways/<name>/oidc_token.json`.
//! File permissions are `0600` (owner-only).

use crate::paths::user_gateway_dir;
use miette::{IntoDiagnostic, Result, WrapErr};
use openshell_core::paths::{ensure_parent_dir_restricted, set_file_owner_only};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;

/// OIDC token bundle persisted to disk.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OidcTokenBundle {
    /// `OAuth2` access token (JWT).
    pub access_token: String,

    /// `OAuth2` refresh token. `None` for `client_credentials` grants.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub refresh_token: Option<String>,

    /// Unix timestamp when the access token expires.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub expires_at: Option<u64>,

    /// OIDC issuer URL.
    pub issuer: String,

    /// OIDC client ID used to obtain the token.
    pub client_id: String,
}

/// Path to the stored OIDC token bundle for a gateway.
pub fn oidc_token_path(gateway_name: &str) -> Result<PathBuf> {
    Ok(user_gateway_dir(gateway_name)?.join("oidc_token.json"))
}

/// Path to the one-shot marker that asks the next browser login to prompt.
pub fn oidc_login_prompt_required_path(gateway_name: &str) -> Result<PathBuf> {
    Ok(user_gateway_dir(gateway_name)?.join("oidc_login_prompt_required"))
}

/// Store an OIDC token bundle for a gateway.
pub fn store_oidc_token(gateway_name: &str, bundle: &OidcTokenBundle) -> Result<()> {
    let path = oidc_token_path(gateway_name)?;
    ensure_parent_dir_restricted(&path)?;
    let json = serde_json::to_string_pretty(bundle)
        .into_diagnostic()
        .wrap_err("failed to serialize OIDC token bundle")?;
    std::fs::write(&path, json)
        .into_diagnostic()
        .wrap_err_with(|| format!("failed to write OIDC token to {}", path.display()))?;
    set_file_owner_only(&path)?;
    clear_oidc_login_prompt(gateway_name)?;
    Ok(())
}

/// Load a stored OIDC token bundle for a gateway.
///
/// Returns `None` if the token file does not exist or cannot be parsed.
pub fn load_oidc_token(gateway_name: &str) -> Option<OidcTokenBundle> {
    let path = oidc_token_path(gateway_name).ok()?;
    if !path.exists() {
        return None;
    }
    let contents = std::fs::read_to_string(&path).ok()?;
    serde_json::from_str(&contents).ok()
}

/// Remove a stored OIDC token.
pub fn remove_oidc_token(gateway_name: &str) -> Result<()> {
    let path = oidc_token_path(gateway_name)?;
    if path.exists() {
        std::fs::remove_file(&path)
            .into_diagnostic()
            .wrap_err_with(|| format!("failed to remove {}", path.display()))?;
    }
    Ok(())
}

/// Mark the next interactive OIDC login as requiring a fresh `IdP` prompt.
pub fn request_oidc_login_prompt(gateway_name: &str) -> Result<()> {
    let path = oidc_login_prompt_required_path(gateway_name)?;
    ensure_parent_dir_restricted(&path)?;
    std::fs::write(&path, b"1\n")
        .into_diagnostic()
        .wrap_err_with(|| format!("failed to write {}", path.display()))?;
    set_file_owner_only(&path)?;
    Ok(())
}

/// Return whether the next interactive OIDC login should request a fresh prompt.
pub fn oidc_login_prompt_required(gateway_name: &str) -> bool {
    oidc_login_prompt_required_path(gateway_name).is_ok_and(|path| path.exists())
}

/// Clear the one-shot fresh-login marker for a gateway.
pub fn clear_oidc_login_prompt(gateway_name: &str) -> Result<()> {
    let path = oidc_login_prompt_required_path(gateway_name)?;
    if path.exists() {
        std::fs::remove_file(&path)
            .into_diagnostic()
            .wrap_err_with(|| format!("failed to remove {}", path.display()))?;
    }
    Ok(())
}

/// Check if the stored access token is expired or near expiry.
///
/// Returns `true` if the token expires within the next 30 seconds.
pub fn is_token_expired(bundle: &OidcTokenBundle) -> bool {
    let Some(expires_at) = bundle.expires_at else {
        // No expiry info — assume valid.
        return false;
    };
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs();
    now + 30 >= expires_at
}

/// Check whether the stored access token has actually expired.
///
/// Unlike [`is_token_expired`], this does not include the 30-second refresh
/// window. Tokens without expiry metadata are treated as valid.
pub fn is_token_actually_expired(bundle: &OidcTokenBundle) -> bool {
    let Some(expires_at) = bundle.expires_at else {
        return false;
    };
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs();
    now >= expires_at
}

#[cfg(test)]
mod tests {
    use super::*;

    #[allow(unsafe_code)]
    fn with_tmp_xdg<F: FnOnce()>(tmp: &std::path::Path, f: F) {
        let _guard = crate::XDG_TEST_LOCK
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let orig = std::env::var("XDG_CONFIG_HOME").ok();
        unsafe {
            std::env::set_var("XDG_CONFIG_HOME", tmp);
        }
        f();
        unsafe {
            match orig {
                Some(v) => std::env::set_var("XDG_CONFIG_HOME", v),
                None => std::env::remove_var("XDG_CONFIG_HOME"),
            }
        }
    }

    #[test]
    fn oidc_token_paths_reject_multi_component_gateway_names() {
        let tmp = tempfile::tempdir().unwrap();
        with_tmp_xdg(tmp.path(), || {
            let bundle = OidcTokenBundle {
                access_token: "token".to_string(),
                refresh_token: None,
                expires_at: None,
                issuer: "https://issuer.example.com".to_string(),
                client_id: "openshell-cli".to_string(),
            };
            assert!(store_oidc_token("../escape", &bundle).is_err());
            assert!(load_oidc_token("../escape").is_none());
            assert!(remove_oidc_token("../escape").is_err());
        });
    }

    #[test]
    fn expiry_checks_distinguish_refresh_window_from_actual_expiry() {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs();
        let bundle = OidcTokenBundle {
            access_token: "token".to_string(),
            refresh_token: None,
            expires_at: Some(now + 10),
            issuer: "https://issuer.example.com".to_string(),
            client_id: "openshell-cli".to_string(),
        };

        assert!(is_token_expired(&bundle));
        assert!(!is_token_actually_expired(&bundle));

        let expired = OidcTokenBundle {
            expires_at: Some(now.saturating_sub(1)),
            ..bundle
        };
        assert!(is_token_actually_expired(&expired));
    }

    #[test]
    fn oidc_login_prompt_marker_is_per_gateway() {
        let tmp = tempfile::tempdir().unwrap();
        with_tmp_xdg(tmp.path(), || {
            assert!(!oidc_login_prompt_required("alpha"));
            assert!(!oidc_login_prompt_required("beta"));

            request_oidc_login_prompt("alpha").unwrap();

            assert!(oidc_login_prompt_required("alpha"));
            assert!(!oidc_login_prompt_required("beta"));

            let bundle = OidcTokenBundle {
                access_token: "token".to_string(),
                refresh_token: None,
                expires_at: None,
                issuer: "https://issuer.example.com".to_string(),
                client_id: "openshell-cli".to_string(),
            };
            store_oidc_token("alpha", &bundle).unwrap();

            assert!(!oidc_login_prompt_required("alpha"));

            request_oidc_login_prompt("alpha").unwrap();
            assert!(oidc_login_prompt_required("alpha"));

            clear_oidc_login_prompt("alpha").unwrap();

            assert!(!oidc_login_prompt_required("alpha"));
        });
    }
}