AkurAI-Notes / crates/notes-api/src/auth/idp_client.rs

Latest change 21fd83f4a6eb769a0ee34aa30f721b5cdc6881af - Fix IDP OAuth: split public issuer vs internal loopback URL for token exchange by Claude Code

pub struct UserInfo {
    pub sub: String,
    pub email: String,
}

pub struct IdpConfig {
    pub issuer: String,
    pub internal_url: String,
    pub client_id: String,
    pub client_secret: String,
    pub redirect_uri: String,
    pub allow: Vec<String>,
}

impl IdpConfig {
    pub fn from_env() -> Self {
        let env = std::env::var;
        Self {
            issuer: env("AKURAI_NOTES_IDP_ISSUER")
                .unwrap_or_else(|_| "https://auth.olibuijr.com".into()),
            internal_url: env("AKURAI_NOTES_IDP_INTERNAL_URL")
                .unwrap_or_else(|_| "http://127.0.0.1:3500".into()),
            client_id: env("AKURAI_NOTES_IDP_CLIENT_ID").unwrap_or_default(),
            client_secret: env("AKURAI_NOTES_IDP_CLIENT_SECRET").unwrap_or_default(),
            redirect_uri: env("AKURAI_NOTES_IDP_REDIRECT")
                .unwrap_or_else(|_| "https://akurai-notes.olibuijr.com/auth/callback".into()),
            allow: env("AKURAI_NOTES_ALLOW")
                .unwrap_or_else(|_| "olibuijr@olibuijr.com".into())
                .split(',')
                .map(str::trim)
                .map(String::from)
                .collect(),
        }
    }
}

pub fn exchange_code(code: &str, _state: &str, cfg: &IdpConfig) -> Result<UserInfo, String> {
    let basic = base64_encode(&format!("{}:{}", cfg.client_id, cfg.client_secret));
    let token_body = format!(
        "grant_type=authorization_code&code={}&redirect_uri={}",
        urlencode(code),
        urlencode(&cfg.redirect_uri),
    );

    let (host, port) = host_port(&cfg.internal_url);

    let response = http_post(
        &host,
        port,
        "/token",
        &[
            ("Authorization", &format!("Basic {}", basic)),
            ("Content-Type", "application/x-www-form-urlencoded"),
        ],
        &token_body,
    )?;

    let token_json = akurai_json::parse(&response.body).map_err(|e| format!("token parse: {e}"))?;

    let access_token = token_json
        .get("access_token")
        .and_then(|v| v.as_str())
        .ok_or_else(|| "no access_token in token response".to_string())?
        .to_string();

    let userinfo_response = http_get(
        &host,
        port,
        "/userinfo",
        &[("Authorization", &format!("Bearer {}", access_token))],
    )?;

    let info =
        akurai_json::parse(&userinfo_response.body).map_err(|e| format!("userinfo parse: {e}"))?;

    let sub = info
        .get("sub")
        .and_then(|v| v.as_str())
        .ok_or_else(|| "no sub in userinfo".to_string())?
        .to_string();

    let email = info
        .get("email")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .to_string();

    Ok(UserInfo { sub, email })
}

fn host_port(issuer: &str) -> (String, u16) {
    let rest = issuer
        .strip_prefix("https://")
        .or_else(|| issuer.strip_prefix("http://"))
        .unwrap_or(issuer);
    if let Some((host, port_str)) = rest.split_once(':') {
        let port: u16 = port_str.parse().unwrap_or(443);
        (host.to_string(), port)
    } else {
        let port = if issuer.starts_with("https://") {
            443
        } else {
            80
        };
        (rest.to_string(), port)
    }
}

#[allow(dead_code)]
struct HttpResponse {
    status: u16,
    body: String,
}

fn http_get(
    host: &str,
    port: u16,
    path: &str,
    headers: &[(&str, &str)],
) -> Result<HttpResponse, String> {
    let mut req = format!("GET {} HTTP/1.1\r\nHost: {}:{}\r\n", path, host, port);
    for (k, v) in headers {
        req.push_str(&format!("{}: {}\r\n", k, v));
    }
    req.push_str("Connection: close\r\n\r\n");
    send_http(host, port, &req)
}

fn http_post(
    host: &str,
    port: u16,
    path: &str,
    headers: &[(&str, &str)],
    body: &str,
) -> Result<HttpResponse, String> {
    let mut req = format!(
        "POST {} HTTP/1.1\r\nHost: {}:{}\r\nContent-Length: {}\r\n",
        path,
        host,
        port,
        body.len()
    );
    for (k, v) in headers {
        req.push_str(&format!("{}: {}\r\n", k, v));
    }
    req.push_str("Connection: close\r\n\r\n");
    req.push_str(body);
    send_http(host, port, &req)
}

fn send_http(host: &str, port: u16, request: &str) -> Result<HttpResponse, String> {
    use std::io::{Read, Write};
    use std::net::TcpStream;
    use std::time::Duration;

    let addr = format!("{}:{}", host, port);
    let stream = TcpStream::connect_timeout(
        &addr.parse().map_err(|e| format!("addr parse: {e}"))?,
        Duration::from_secs(5),
    )
    .map_err(|e| format!("connect: {e}"))?;

    stream
        .set_read_timeout(Some(Duration::from_secs(10)))
        .map_err(|e| format!("set timeout: {e}"))?;

    let mut wr = &stream;
    wr.write_all(request.as_bytes())
        .map_err(|e| format!("write: {e}"))?;

    let mut buf = Vec::new();
    let mut rd = &stream;
    rd.read_to_end(&mut buf).map_err(|e| format!("read: {e}"))?;

    let response = String::from_utf8_lossy(&buf).to_string();

    let status = response
        .lines()
        .next()
        .and_then(|line| line.split(' ').nth(1))
        .and_then(|s| s.parse::<u16>().ok())
        .unwrap_or(0);

    let body = if let Some(pos) = response.find("\r\n\r\n") {
        response[pos + 4..].to_string()
    } else {
        String::new()
    };

    if !(200..300).contains(&status) {
        return Err(format!(
            "HTTP {}: {}",
            status,
            body.chars().take(200).collect::<String>()
        ));
    }

    Ok(HttpResponse { status, body })
}

fn base64_encode(input: &str) -> String {
    const CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let bytes = input.as_bytes();
    let mut result = String::new();

    for chunk in bytes.chunks(3) {
        let b0 = chunk[0] as u32;
        let b1 = chunk.get(1).copied().unwrap_or(0) as u32;
        let b2 = chunk.get(2).copied().unwrap_or(0) as u32;
        let triple = (b0 << 16) | (b1 << 8) | b2;

        result.push(CHARS[((triple >> 18) & 0x3F) as usize] as char);
        result.push(CHARS[((triple >> 12) & 0x3F) as usize] as char);

        if chunk.len() > 1 {
            result.push(CHARS[((triple >> 6) & 0x3F) as usize] as char);
        } else {
            result.push('=');
        }

        if chunk.len() > 2 {
            result.push(CHARS[(triple & 0x3F) as usize] as char);
        } else {
            result.push('=');
        }
    }

    result
}

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