el-retired/ui/vessels/el-identity/src/oauth.rs

//! OAuth 2.0 flows implemented as Engram graph operations.
//!
//! Implements PKCE (RFC 7636) and Authorization Code flow without external
//! OAuth crates. Token exchange uses `reqwest` (already in the workspace).
//! All tokens are hashed before graph storage — the raw token never persists.
//!
//! Flow:
//! 1. `begin_auth_flow()` → PKCE verifier + challenge, redirect URL
//! 2. Provider redirects back with `code`
//! 3. `exchange_code()` → calls provider token endpoint, stores OAuthToken node
//! 4. `refresh_token()` → find OAuthToken node, call refresh endpoint, update node

use crate::{
    engram::EngramClient,
    error::{IdentityError, IdentityResult},
    nodes::{OAuthToken, User, EDGE_AUTHENTICATED_VIA, NODE_OAUTH_TOKEN},
    provider::OAuthProvider,
    session::SessionManager,
};
use chrono::{Duration, Utc};
use sha2::{Digest, Sha256};
use std::sync::Arc;

// ── PKCE helpers ──────────────────────────────────────────────────────────────

/// A PKCE verifier/challenge pair.
#[derive(Debug, Clone)]
pub struct PkceChallenge {
    /// The raw verifier — sent to the token endpoint.
    pub verifier: String,
    /// The challenge (BASE64URL(SHA256(verifier))) — sent in the auth request.
    pub challenge: String,
    /// Always "S256".
    pub method: &'static str,
}

impl PkceChallenge {
    /// Generate a new PKCE verifier and compute the S256 challenge.
    ///
    /// The verifier is a 43-character URL-safe random string derived from
    /// entropy collected from the system clock and a UUID.
    pub fn generate() -> Self {
        let verifier = generate_pkce_verifier();
        let challenge = pkce_s256_challenge(&verifier);
        Self {
            verifier,
            challenge,
            method: "S256",
        }
    }

    /// Verify that a verifier matches this challenge (used in tests and server-side).
    pub fn verify(&self, verifier: &str) -> bool {
        pkce_s256_challenge(verifier) == self.challenge
    }
}

fn generate_pkce_verifier() -> String {
    // 32 random bytes → base64url (43 chars, no padding)
    // We derive entropy from UUID (random in v4) + timestamp nanos.
    let id = uuid::Uuid::new_v4();
    let ts = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.subsec_nanos())
        .unwrap_or(0);
    let mut raw = [0u8; 32];
    let id_bytes = id.as_bytes();
    for i in 0..16 {
        raw[i] = id_bytes[i];
    }
    let ts_bytes = ts.to_le_bytes();
    for i in 0..4 {
        raw[16 + i] = ts_bytes[i];
    }
    // Fill remaining with XOR mix
    for i in 20..32 {
        raw[i] = id_bytes[i - 16] ^ ts_bytes[i % 4];
    }
    base64url_encode_no_pad(&raw)
}

fn pkce_s256_challenge(verifier: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(verifier.as_bytes());
    base64url_encode_no_pad(&hasher.finalize())
}

fn base64url_encode_no_pad(input: &[u8]) -> String {
    const CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
    let mut out = String::new();
    for chunk in input.chunks(3) {
        let b0 = chunk[0] as u32;
        let b1 = if chunk.len() > 1 { chunk[1] as u32 } else { 0 };
        let b2 = if chunk.len() > 2 { chunk[2] as u32 } else { 0 };
        let n = (b0 << 16) | (b1 << 8) | b2;
        out.push(CHARS[((n >> 18) & 63) as usize] as char);
        out.push(CHARS[((n >> 12) & 63) as usize] as char);
        if chunk.len() > 1 {
            out.push(CHARS[((n >> 6) & 63) as usize] as char);
        }
        if chunk.len() > 2 {
            out.push(CHARS[(n & 63) as usize] as char);
        }
    }
    out
}

/// Hash a token with SHA-256 for safe graph storage.
pub fn hash_token(token: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(token.as_bytes());
    hex::encode_lower_sha256(hasher.finalize().as_ref())
}

mod hex {
    pub fn encode_lower_sha256(bytes: &[u8]) -> String {
        bytes.iter().map(|b| format!("{:02x}", b)).collect()
    }
}

// ── Authorization Code Flow ───────────────────────────────────────────────────

/// Parameters for starting an OAuth authorization code flow.
#[derive(Debug, Clone)]
pub struct AuthFlowParams {
    /// The PKCE challenge (save the verifier for use at exchange time).
    pub pkce: PkceChallenge,
    /// The full redirect URL to send the user to.
    pub redirect_url: String,
    /// An opaque state value for CSRF protection.
    pub state: String,
}

/// The result of a successful token exchange.
#[derive(Debug, Clone)]
pub struct TokenSet {
    pub access_token: String,
    pub refresh_token: Option<String>,
    /// Expiry in seconds from now.
    pub expires_in: u64,
    pub scopes: Vec<String>,
}

/// OAuth flow coordinator — executes auth code + PKCE flows and writes
/// the resulting tokens to the Engram graph.
pub struct OAuthFlow {
    client: Arc<dyn EngramClient>,
    /// Kept for future use (e.g., session creation during OAuth callback).
    _session_manager: Arc<SessionManager>,
}

impl OAuthFlow {
    pub fn new(client: Arc<dyn EngramClient>, session_manager: Arc<SessionManager>) -> Self {
        Self { client, _session_manager: session_manager }
    }

    /// Step 1: Generate the redirect URL and PKCE parameters.
    ///
    /// The caller should:
    /// 1. Save `params.pkce.verifier` in the user's browser session (cookie/localStorage).
    /// 2. Redirect the user to `params.redirect_url`.
    pub fn begin_auth_flow(
        &self,
        provider: &dyn OAuthProvider,
        redirect_uri: &str,
        extra_scopes: &[&str],
    ) -> IdentityResult<AuthFlowParams> {
        let pkce = PkceChallenge::generate();
        let state = generate_pkce_verifier(); // reuse the verifier generator for state
        let url = provider.authorization_url(
            redirect_uri,
            &pkce.challenge,
            &state,
            extra_scopes,
        );
        Ok(AuthFlowParams { pkce, redirect_url: url, state })
    }

    /// Step 2: Exchange the authorization code for tokens.
    ///
    /// - `user` — the authenticated user to link the token to
    /// - `code` — the authorization code from the provider callback
    /// - `pkce_verifier` — the verifier saved in step 1
    /// - `session_id` — the session to attach the OAuthToken to
    ///
    /// Returns the session ID that now has an OAuthToken attached via graph edge.
    pub fn exchange_code(
        &self,
        provider: &dyn OAuthProvider,
        _user: &User,
        code: &str,
        pkce_verifier: &str,
        redirect_uri: &str,
        session_id: &str,
    ) -> IdentityResult<OAuthToken> {
        // Exchange code with provider (HTTP call inside OAuthProvider::exchange_code)
        let token_set = provider.exchange_code(code, pkce_verifier, redirect_uri)?;

        // Hash tokens before storing
        let access_hash = hash_token(&token_set.access_token);
        let refresh_hash = token_set.refresh_token.as_deref().map(hash_token);
        let expires_at = Utc::now() + Duration::seconds(token_set.expires_in as i64);

        let oauth_token = OAuthToken::new(
            provider.name(),
            access_hash,
            refresh_hash,
            expires_at,
            token_set.scopes,
        );

        // Store OAuthToken node in graph
        let token_id = self
            .client
            .create_node(NODE_OAUTH_TOKEN, oauth_token.to_value())
            .map_err(|e| IdentityError::GraphError(e.to_string()))?;

        // Edge: Session → OAuthToken (authenticated_via)
        self.client
            .create_edge(session_id, &token_id, EDGE_AUTHENTICATED_VIA)
            .map_err(|e| IdentityError::GraphError(e.to_string()))?;

        Ok(oauth_token)
    }

    /// Step 3: Refresh an expired access token.
    ///
    /// Finds the OAuthToken node for the given session, calls the provider's
    /// refresh endpoint, and updates the node in-place (delete old, create new).
    pub fn refresh_token(
        &self,
        provider: &dyn OAuthProvider,
        session_id: &str,
        refresh_token: &str,
    ) -> IdentityResult<OAuthToken> {
        // Exchange refresh token with provider
        let token_set = provider.refresh_token(refresh_token)?;

        // Find and delete old OAuthToken nodes attached to this session
        let old_tokens = self
            .client
            .find_connected(session_id, EDGE_AUTHENTICATED_VIA)
            .map_err(|e| IdentityError::GraphError(e.to_string()))?;

        for old in &old_tokens {
            if let Some(id) = old.get("id").and_then(|v| v.as_str()) {
                let _ = self.client.delete_node(id);
            }
        }

        // Store new token
        let access_hash = hash_token(&token_set.access_token);
        let refresh_hash = token_set.refresh_token.as_deref().map(hash_token);
        let expires_at = Utc::now() + Duration::seconds(token_set.expires_in as i64);

        let new_token = OAuthToken::new(
            provider.name(),
            access_hash,
            refresh_hash,
            expires_at,
            token_set.scopes,
        );

        let token_id = self
            .client
            .create_node(NODE_OAUTH_TOKEN, new_token.to_value())
            .map_err(|e| IdentityError::GraphError(e.to_string()))?;

        self.client
            .create_edge(session_id, &token_id, EDGE_AUTHENTICATED_VIA)
            .map_err(|e| IdentityError::GraphError(e.to_string()))?;

        Ok(new_token)
    }
}