AkurAI-PassVault / crates/passvault-mcp/src/auth.rs

Latest change f0de6c29b6eff8151dc0f8eebcffcf409026c988 - feat: expose PassVault UI through AkurAI ID by Ólafur Búi Ólafsson

use akurai_json::Value;
use std::env;
use std::io::{Read, Write};
use std::net::TcpStream;
use std::time::Duration;

pub struct UserInfo {
    pub email: String,
}

#[derive(Clone)]
pub struct IdpConfig {
    pub issuer: String,
    pub internal_url: String,
    pub client_id: String,
    pub client_secret: String,
    pub redirect_uri: String,
    pub allowed_emails: Vec<String>,
}

impl IdpConfig {
    pub fn from_env() -> Self {
        Self {
            issuer: env::var("AKURAI_PASSVAULT_IDP_ISSUER")
                .unwrap_or_else(|_| "https://auth.olibuijr.com".into()),
            internal_url: env::var("AKURAI_PASSVAULT_IDP_INTERNAL_URL")
                .unwrap_or_else(|_| "http://127.0.0.1:3500".into()),
            client_id: env::var("AKURAI_PASSVAULT_IDP_CLIENT_ID").unwrap_or_default(),
            client_secret: env::var("AKURAI_PASSVAULT_IDP_CLIENT_SECRET").unwrap_or_default(),
            redirect_uri: env::var("AKURAI_PASSVAULT_IDP_REDIRECT")
                .unwrap_or_else(|_| "https://akurai-passvault.olibuijr.com/auth/callback".into()),
            allowed_emails: env::var("AKURAI_PASSVAULT_IDP_ALLOW")
                .unwrap_or_else(|_| "olibuijr@olibuijr.com".into())
                .split(',')
                .map(str::trim)
                .filter(|email| !email.is_empty())
                .map(String::from)
                .collect(),
        }
    }

    pub fn is_configured(&self) -> bool {
        !self.client_id.is_empty() && !self.client_secret.is_empty()
    }

    pub fn authorize_url(&self, state: &str) -> String {
        format!(
            "{}/authorize?client_id={}&redirect_uri={}&response_type=code&scope=openid%20profile%20email&state={}",
            self.issuer.trim_end_matches('/'),
            urlencode(&self.client_id),
            urlencode(&self.redirect_uri),
            urlencode(state),
        )
    }

    pub fn allows(&self, email: &str) -> bool {
        self.allowed_emails
            .iter()
            .any(|allowed| allowed.eq_ignore_ascii_case(email))
    }
}

pub fn exchange_code(code: &str, config: &IdpConfig) -> Result<UserInfo, String> {
    let basic = base64_encode(&format!("{}:{}", config.client_id, config.client_secret));
    let body = format!(
        "grant_type=authorization_code&code={}&redirect_uri={}",
        urlencode(code),
        urlencode(&config.redirect_uri),
    );
    let (host, port) = host_port(&config.internal_url)?;
    let token = request(
        &host,
        port,
        "POST",
        "/token",
        &[
            ("Authorization", format!("Basic {basic}")),
            ("Content-Type", "application/x-www-form-urlencoded".into()),
        ],
        &body,
    )?;
    let token: Value =
        akurai_json::parse(&token).map_err(|error| format!("token parse: {error}"))?;
    let access_token = token
        .get("access_token")
        .and_then(Value::as_str)
        .ok_or_else(|| "token response omitted access_token".to_string())?;
    let userinfo = request(
        &host,
        port,
        "GET",
        "/userinfo",
        &[("Authorization", format!("Bearer {access_token}"))],
        "",
    )?;
    let userinfo: Value =
        akurai_json::parse(&userinfo).map_err(|error| format!("userinfo parse: {error}"))?;
    let email = userinfo
        .get("email")
        .and_then(Value::as_str)
        .filter(|email| !email.is_empty())
        .ok_or_else(|| "userinfo omitted email".to_string())?
        .to_string();
    Ok(UserInfo { email })
}

fn host_port(url: &str) -> Result<(String, u16), String> {
    let (scheme, rest) = url
        .split_once("://")
        .ok_or_else(|| "IDP internal URL must include a scheme".to_string())?;
    if scheme != "http" {
        return Err("IDP internal URL must use loopback HTTP".into());
    }
    let (host, port) = match rest.rsplit_once(':') {
        Some((host, port)) => (
            host,
            port.parse::<u16>()
                .map_err(|_| "invalid IDP internal port".to_string())?,
        ),
        None => (rest, 80),
    };
    if host != "127.0.0.1" && host != "localhost" {
        return Err("IDP internal URL must be loopback".into());
    }
    Ok((host.to_string(), port))
}

fn request(
    host: &str,
    port: u16,
    method: &str,
    path: &str,
    headers: &[(impl AsRef<str>, String)],
    body: &str,
) -> Result<String, String> {
    let mut stream = TcpStream::connect_timeout(
        &format!("{host}:{port}")
            .parse()
            .map_err(|error| format!("IDP address: {error}"))?,
        Duration::from_secs(5),
    )
    .map_err(|error| format!("IDP connect: {error}"))?;
    stream
        .set_read_timeout(Some(Duration::from_secs(10)))
        .map_err(|error| format!("IDP timeout: {error}"))?;
    let mut message = format!(
        "{method} {path} HTTP/1.1\r\nHost: {host}:{port}\r\nConnection: close\r\nContent-Length: {}\r\n",
        body.len()
    );
    for (name, value) in headers {
        message.push_str(name.as_ref());
        message.push_str(": ");
        message.push_str(value);
        message.push_str("\r\n");
    }
    message.push_str("\r\n");
    message.push_str(body);
    stream
        .write_all(message.as_bytes())
        .map_err(|error| format!("IDP write: {error}"))?;
    let mut response = String::new();
    stream
        .read_to_string(&mut response)
        .map_err(|error| format!("IDP read: {error}"))?;
    let status = response
        .lines()
        .next()
        .and_then(|line| line.split_whitespace().nth(1))
        .and_then(|status| status.parse::<u16>().ok())
        .unwrap_or_default();
    let body = response
        .split_once("\r\n\r\n")
        .map(|(_, body)| body)
        .unwrap_or_default();
    if !(200..300).contains(&status) {
        return Err(format!("IDP returned HTTP {status}"));
    }
    Ok(body.to_string())
}

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

fn base64_encode(value: &str) -> String {
    const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut encoded = String::with_capacity(value.len().div_ceil(3) * 4);
    for chunk in value.as_bytes().chunks(3) {
        let first = chunk[0] as u32;
        let second = chunk.get(1).copied().unwrap_or_default() as u32;
        let third = chunk.get(2).copied().unwrap_or_default() as u32;
        let value = first << 16 | second << 8 | third;
        encoded.push(ALPHABET[((value >> 18) & 63) as usize] as char);
        encoded.push(ALPHABET[((value >> 12) & 63) as usize] as char);
        encoded.push(if chunk.len() > 1 {
            ALPHABET[((value >> 6) & 63) as usize] as char
        } else {
            '='
        });
        encoded.push(if chunk.len() > 2 {
            ALPHABET[(value & 63) as usize] as char
        } else {
            '='
        });
    }
    encoded
}

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

    fn config() -> IdpConfig {
        IdpConfig {
            issuer: "https://auth.example/".into(),
            internal_url: "http://127.0.0.1:3500".into(),
            client_id: "client id".into(),
            client_secret: "secret".into(),
            redirect_uri: "https://vault.example/auth/callback".into(),
            allowed_emails: vec!["Owner@Example.com".into()],
        }
    }

    #[test]
    fn authorize_url_encodes_values() {
        let url = config().authorize_url("state value");
        assert!(url.starts_with("https://auth.example/authorize?"));
        assert!(url.contains("client_id=client%20id"));
        assert!(url.contains("state=state%20value"));
    }

    #[test]
    fn allowlist_is_case_insensitive() {
        assert!(config().allows("owner@example.com"));
        assert!(!config().allows("other@example.com"));
    }

    #[test]
    fn internal_idp_must_be_loopback_http() {
        assert_eq!(host_port("http://127.0.0.1:3500").unwrap().1, 3500);
        assert!(host_port("https://auth.example").is_err());
        assert!(host_port("http://10.0.0.1:3500").is_err());
    }

    #[test]
    fn basic_auth_encoding_is_standard() {
        assert_eq!(base64_encode("client:secret"), "Y2xpZW50OnNlY3JldA==");
    }
}