robinpath_modules/modules/

oauth_mod.rs

use robinpath::{RobinPath, Value};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

struct TokenStore {
    access_token: String,
    refresh_token: Option<String>,
    token_type: String,
    expires_at: Option<u64>,
    scope: Option<String>,
}

pub fn register(rp: &mut RobinPath) {
    let state = Arc::new(Mutex::new(HashMap::<String, TokenStore>::new()));

    // oauth.authUrl baseUrl options → URL string
    rp.register_builtin("oauth.authUrl", |args, _| {
        let base_url = args.first().map(|v| v.to_display_string()).unwrap_or_default();
        let opts = args.get(1).cloned().unwrap_or(Value::Null);
        if base_url.is_empty() {
            return Err("Authorization URL is required".to_string());
        }
        let mut params = Vec::new();
        if let Value::Object(obj) = &opts {
            let response_type = obj.get("responseType")
                .map(|v| v.to_display_string())
                .unwrap_or_else(|| "code".to_string());
            params.push(format!("response_type={}", url_encode(&response_type)));
            if let Some(v) = obj.get("clientId") {
                params.push(format!("client_id={}", url_encode(&v.to_display_string())));
            }
            if let Some(v) = obj.get("redirectUri") {
                params.push(format!("redirect_uri={}", url_encode(&v.to_display_string())));
            }
            if let Some(v) = obj.get("scope") {
                params.push(format!("scope={}", url_encode(&v.to_display_string())));
            }
            if let Some(v) = obj.get("state") {
                params.push(format!("state={}", url_encode(&v.to_display_string())));
            }
            if let Some(v) = obj.get("codeChallenge") {
                params.push(format!("code_challenge={}", url_encode(&v.to_display_string())));
            }
            if let Some(v) = obj.get("codeChallengeMethod") {
                params.push(format!("code_challenge_method={}", url_encode(&v.to_display_string())));
            }
            if let Some(v) = obj.get("accessType") {
                params.push(format!("access_type={}", url_encode(&v.to_display_string())));
            }
            if let Some(v) = obj.get("prompt") {
                params.push(format!("prompt={}", url_encode(&v.to_display_string())));
            }
        }
        let sep = if base_url.contains('?') { "&" } else { "?" };
        Ok(Value::String(format!("{}{}{}", base_url, sep, params.join("&"))))
    });

    // oauth.pkceVerifier length? → random base64url string
    rp.register_builtin("oauth.pkceVerifier", |args, _| {
        let length = args.first().map(|v| v.to_number() as usize).unwrap_or(64);
        let clamped = length.max(43).min(128);
        let bytes: Vec<u8> = random_bytes(clamped);
        let encoded = base64url_encode(&bytes);
        Ok(Value::String(encoded[..clamped.min(encoded.len())].to_string()))
    });

    // oauth.pkceChallenge verifier method? → {challenge, method}
    rp.register_builtin("oauth.pkceChallenge", |args, _| {
        let verifier = args.first().map(|v| v.to_display_string()).unwrap_or_default();
        let method = args.get(1).map(|v| v.to_display_string()).unwrap_or_else(|| "S256".to_string());
        let mut obj = indexmap::IndexMap::new();
        if method == "S256" {
            use sha2::Digest;
            let hash = sha2::Sha256::digest(verifier.as_bytes());
            let challenge = base64url_encode(&hash);
            obj.insert("challenge".to_string(), Value::String(challenge));
            obj.insert("method".to_string(), Value::String("S256".to_string()));
        } else if method == "plain" {
            obj.insert("challenge".to_string(), Value::String(verifier));
            obj.insert("method".to_string(), Value::String("plain".to_string()));
        } else {
            return Err(format!("Unsupported PKCE method: {}. Use \"S256\" or \"plain\".", method));
        }
        Ok(Value::Object(obj))
    });

    // oauth.generateState length? → random hex string
    rp.register_builtin("oauth.generateState", |args, _| {
        let length = args.first().map(|v| v.to_number() as usize).unwrap_or(32);
        let bytes = random_bytes(length);
        let hex: String = bytes.iter().map(|b| format!("{:02x}", b)).collect();
        Ok(Value::String(hex))
    });

    // oauth.storeToken name tokenData → stores a token
    let s = state.clone();
    rp.register_builtin("oauth.storeToken", move |args, _| {
        let name = args.first().map(|v| v.to_display_string()).unwrap_or_default();
        let data = args.get(1).cloned().unwrap_or(Value::Null);
        if let Value::Object(obj) = &data {
            let store = TokenStore {
                access_token: obj.get("accessToken").map(|v| v.to_display_string()).unwrap_or_default(),
                refresh_token: obj.get("refreshToken").map(|v| v.to_display_string()),
                token_type: obj.get("tokenType").map(|v| v.to_display_string()).unwrap_or_else(|| "Bearer".to_string()),
                expires_at: obj.get("expiresAt").map(|v| v.to_number() as u64),
                scope: obj.get("scope").map(|v| v.to_display_string()),
            };
            s.lock().unwrap().insert(name, store);
            Ok(Value::Bool(true))
        } else {
            Err("Token data must be an object".to_string())
        }
    });

    // oauth.getToken name → token object or null
    let s = state.clone();
    rp.register_builtin("oauth.getToken", move |args, _| {
        let name = args.first().map(|v| v.to_display_string()).unwrap_or_default();
        let tokens = s.lock().unwrap();
        match tokens.get(&name) {
            Some(store) => {
                let mut obj = indexmap::IndexMap::new();
                obj.insert("accessToken".to_string(), Value::String(store.access_token.clone()));
                if let Some(ref rt) = store.refresh_token {
                    obj.insert("refreshToken".to_string(), Value::String(rt.clone()));
                }
                obj.insert("tokenType".to_string(), Value::String(store.token_type.clone()));
                let now_ms = std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_millis() as u64;
                let expired = store.expires_at.map(|ea| now_ms > ea).unwrap_or(false);
                obj.insert("expired".to_string(), Value::Bool(expired));
                if let Some(ea) = store.expires_at {
                    obj.insert("expiresAt".to_string(), Value::Number(ea as f64));
                }
                if let Some(ref sc) = store.scope {
                    obj.insert("scope".to_string(), Value::String(sc.clone()));
                }
                Ok(Value::Object(obj))
            }
            None => Ok(Value::Null),
        }
    });

    // oauth.isExpired name bufferMs? → bool
    let s = state.clone();
    rp.register_builtin("oauth.isExpired", move |args, _| {
        let name = args.first().map(|v| v.to_display_string()).unwrap_or_default();
        let buffer_ms = args.get(1).map(|v| v.to_number() as u64).unwrap_or(60000);
        let tokens = s.lock().unwrap();
        match tokens.get(&name) {
            Some(store) => {
                if let Some(ea) = store.expires_at {
                    let now_ms = std::time::SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .unwrap_or_default()
                        .as_millis() as u64;
                    Ok(Value::Bool(now_ms + buffer_ms > ea))
                } else {
                    Ok(Value::Bool(false))
                }
            }
            None => Ok(Value::Bool(true)),
        }
    });

    // oauth.clearTokens name? → bool
    let s = state.clone();
    rp.register_builtin("oauth.clearTokens", move |args, _| {
        let name = args.first().map(|v| v.to_display_string());
        let mut tokens = s.lock().unwrap();
        if let Some(n) = name {
            Ok(Value::Bool(tokens.remove(&n).is_some()))
        } else {
            tokens.clear();
            Ok(Value::Bool(true))
        }
    });

    // oauth.exchangeCode tokenUrl options → {accessToken, ...} (requires reqwest)
    #[cfg(feature = "api")]
    {
        let s = state.clone();
        rp.register_builtin("oauth.exchangeCode", move |args, _| {
            let token_url = args.first().map(|v| v.to_display_string()).unwrap_or_default();
            let opts = args.get(1).cloned().unwrap_or(Value::Null);
            if token_url.is_empty() {
                return Err("Token URL is required".to_string());
            }
            let mut params = Vec::new();
            params.push(("grant_type".to_string(), "authorization_code".to_string()));
            if let Value::Object(obj) = &opts {
                if let Some(v) = obj.get("code") { params.push(("code".to_string(), v.to_display_string())); }
                if let Some(v) = obj.get("clientId") { params.push(("client_id".to_string(), v.to_display_string())); }
                if let Some(v) = obj.get("clientSecret") { params.push(("client_secret".to_string(), v.to_display_string())); }
                if let Some(v) = obj.get("redirectUri") { params.push(("redirect_uri".to_string(), v.to_display_string())); }
                if let Some(v) = obj.get("codeVerifier") { params.push(("code_verifier".to_string(), v.to_display_string())); }
            }
            do_token_request(&token_url, &params, &opts, &s)
        });
    }

    // oauth.refreshToken tokenUrl options → {accessToken, ...}
    #[cfg(feature = "api")]
    {
        let s = state.clone();
        rp.register_builtin("oauth.refreshToken", move |args, _| {
            let token_url = args.first().map(|v| v.to_display_string()).unwrap_or_default();
            let opts = args.get(1).cloned().unwrap_or(Value::Null);
            if token_url.is_empty() {
                return Err("Token URL is required".to_string());
            }
            let refresh = if let Value::Object(obj) = &opts {
                obj.get("refreshToken").map(|v| v.to_display_string())
                    .or_else(|| {
                        obj.get("name").and_then(|n| {
                            s.lock().unwrap().get(&n.to_display_string())
                                .and_then(|st| st.refresh_token.clone())
                        })
                    })
            } else {
                None
            };
            let refresh = refresh.ok_or("Refresh token is required")?;
            let mut params = Vec::new();
            params.push(("grant_type".to_string(), "refresh_token".to_string()));
            params.push(("refresh_token".to_string(), refresh));
            if let Value::Object(obj) = &opts {
                if let Some(v) = obj.get("clientId") { params.push(("client_id".to_string(), v.to_display_string())); }
                if let Some(v) = obj.get("clientSecret") { params.push(("client_secret".to_string(), v.to_display_string())); }
                if let Some(v) = obj.get("scope") { params.push(("scope".to_string(), v.to_display_string())); }
            }
            do_token_request(&token_url, &params, &opts, &s)
        });
    }

    // oauth.clientCredentials tokenUrl options → {accessToken, ...}
    #[cfg(feature = "api")]
    {
        let s = state.clone();
        rp.register_builtin("oauth.clientCredentials", move |args, _| {
            let token_url = args.first().map(|v| v.to_display_string()).unwrap_or_default();
            let opts = args.get(1).cloned().unwrap_or(Value::Null);
            if token_url.is_empty() {
                return Err("Token URL is required".to_string());
            }
            let mut params = Vec::new();
            params.push(("grant_type".to_string(), "client_credentials".to_string()));
            if let Value::Object(obj) = &opts {
                if let Some(v) = obj.get("clientId") { params.push(("client_id".to_string(), v.to_display_string())); }
                if let Some(v) = obj.get("clientSecret") { params.push(("client_secret".to_string(), v.to_display_string())); }
                if let Some(v) = obj.get("scope") { params.push(("scope".to_string(), v.to_display_string())); }
            }
            do_token_request(&token_url, &params, &opts, &s)
        });
    }
}

#[cfg(feature = "api")]
fn do_token_request(
    url: &str,
    params: &[(String, String)],
    opts: &Value,
    state: &Arc<Mutex<HashMap<String, TokenStore>>>,
) -> Result<Value, String> {
    let client = reqwest::blocking::Client::new();
    let resp = client.post(url)
        .header("Content-Type", "application/x-www-form-urlencoded")
        .form(params)
        .send()
        .map_err(|e| format!("OAuth request error: {}", e))?;

    let status = resp.status().as_u16();
    let body_text = resp.text().map_err(|e| format!("body error: {}", e))?;
    let json: serde_json::Value = serde_json::from_str(&body_text)
        .map_err(|_| format!("Invalid JSON response: {}", body_text))?;

    if status >= 400 {
        let error = json.get("error").and_then(|v| v.as_str()).unwrap_or("unknown");
        let desc = json.get("error_description").and_then(|v| v.as_str()).unwrap_or("");
        return Err(format!("OAuth error: {} - {}", error, desc));
    }

    let access_token = json.get("access_token").and_then(|v| v.as_str()).unwrap_or("").to_string();
    let refresh_token = json.get("refresh_token").and_then(|v| v.as_str()).map(|s| s.to_string());
    let token_type = json.get("token_type").and_then(|v| v.as_str()).unwrap_or("Bearer").to_string();
    let expires_in = json.get("expires_in").and_then(|v| v.as_u64());
    let scope = json.get("scope").and_then(|v| v.as_str()).map(|s| s.to_string());

    // Store if name provided
    if let Value::Object(obj) = opts {
        if let Some(name_val) = obj.get("name") {
            let name = name_val.to_display_string();
            let now_ms = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_millis() as u64;
            let store = TokenStore {
                access_token: access_token.clone(),
                refresh_token: refresh_token.clone(),
                token_type: token_type.clone(),
                expires_at: expires_in.map(|ei| now_ms + ei * 1000),
                scope: scope.clone(),
            };
            state.lock().unwrap().insert(name, store);
        }
    }

    let mut result = indexmap::IndexMap::new();
    result.insert("accessToken".to_string(), Value::String(access_token));
    if let Some(rt) = refresh_token {
        result.insert("refreshToken".to_string(), Value::String(rt));
    }
    result.insert("tokenType".to_string(), Value::String(token_type));
    if let Some(ei) = expires_in {
        result.insert("expiresIn".to_string(), Value::Number(ei as f64));
    }
    if let Some(s) = scope {
        result.insert("scope".to_string(), Value::String(s));
    }
    Ok(Value::Object(result))
}

fn random_bytes(n: usize) -> Vec<u8> {
    use std::collections::hash_map::RandomState;
    use std::hash::{BuildHasher, Hasher};
    let mut bytes = Vec::with_capacity(n);
    for i in 0..n {
        let state = RandomState::new();
        let mut hasher = state.build_hasher();
        hasher.write_usize(i);
        bytes.push((hasher.finish() & 0xFF) as u8);
    }
    bytes
}

fn base64url_encode(data: &[u8]) -> String {
    use base64::Engine;
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data)
}

fn url_encode(s: &str) -> String {
    let mut result = String::new();
    for b in s.bytes() {
        match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
                result.push(b as char);
            }
            _ => result.push_str(&format!("%{:02X}", b)),
        }
    }
    result
}