AkurAI-Drive / backend/src/idp.rs

Latest change 3c20b23ae95e841dff36ec4a0d843bde1767ca77 - fix: cargo fmt backend/src/idp.rs by Ólafur Búi Ólafsson

use base64::Engine;

use crate::config;

/// The subset of OIDC userinfo the Drive needs.
pub struct UserInfo {
    pub sub: String,
    pub email: String,
    pub groups: Vec<String>,
    pub tenant_id: Option<String>,
}

/// Back-channel: exchange the authorization `code` for an access token, then
/// fetch userinfo. Blocking (uses `ureq`) — call from `spawn_blocking`.
///
/// All traffic here is plain HTTP over localhost (`IDP_INTERNAL_URL`), so no TLS
/// is involved. The token endpoint uses HTTP Basic auth with the OIDC client
/// credentials; userinfo uses the resulting Bearer token.
pub fn exchange_code_for_userinfo(code: &str) -> Result<UserInfo, String> {
    let cfg = config::get();

    // 1. Token exchange (HTTP Basic client auth, form body).
    let basic = base64::engine::general_purpose::STANDARD
        .encode(format!("{}:{}", cfg.oidc_client_id, cfg.oidc_client_secret));
    let token_url = format!("{}/token", cfg.idp_internal_url);

    let token_body = ureq::post(&token_url)
        .set("Authorization", &format!("Basic {basic}"))
        .send_form(&[
            ("grant_type", "authorization_code"),
            ("code", code),
            ("redirect_uri", &cfg.oidc_redirect_uri),
        ])
        .map_err(|e| format!("token request failed: {e}"))?
        .into_string()
        .map_err(|e| format!("token body read failed: {e}"))?;

    let token_json: serde_json::Value =
        serde_json::from_str(&token_body).map_err(|e| format!("token json parse failed: {e}"))?;

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

    // 2. Userinfo (Bearer auth).
    let userinfo_url = format!("{}/userinfo", cfg.idp_internal_url);
    let userinfo_body = ureq::get(&userinfo_url)
        .set("Authorization", &format!("Bearer {access_token}"))
        .call()
        .map_err(|e| format!("userinfo request failed: {e}"))?
        .into_string()
        .map_err(|e| format!("userinfo body read failed: {e}"))?;

    let info: serde_json::Value = serde_json::from_str(&userinfo_body)
        .map_err(|e| format!("userinfo json parse failed: {e}"))?;

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

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

    let groups = info
        .get("groups")
        .and_then(|v| v.as_array())
        .map(|arr| {
            arr.iter()
                .filter_map(|g| g.as_str().map(String::from))
                .collect()
        })
        .unwrap_or_default();

    let tenant_id = info
        .get("tenant_id")
        .and_then(|v| v.as_str())
        .map(String::from);

    Ok(UserInfo {
        sub,
        email,
        groups,
        tenant_id,
    })
}

/// Outcome of a headless email+password login against the IDP.
pub enum PasswordLogin {
    /// Credentials accepted and the OIDC code was exchanged into userinfo.
    Ok(UserInfo),
    /// The IDP rejected the email/password pair.
    InvalidCredentials,
    /// The account has MFA enabled — the browser flow must be used instead.
    MfaRequired,
}

/// Headless email+password login. The IDP exposes no password grant, so we drive
/// the same browser login flow it serves, but server-side over localhost:
///
/// 1. `GET /login` — read the `_csrf` form token and the `_csrf` cookie.
/// 2. `POST /login` (form-encoded, `Cookie: _csrf=…`) — success is signalled by a
///    `Set-Cookie: idp_session=…`; a redirect to `/mfa` means MFA is required;
///    anything else (the login page re-rendered with an error) is bad credentials.
/// 3. `GET /authorize` (`Cookie: idp_session=…`, redirects NOT followed) — read the
///    `code` out of the `Location` redirect to our `redirect_uri`.
/// 4. Exchange the code for userinfo via [`exchange_code_for_userinfo`].
///
/// Blocking (uses `ureq`) — call from `spawn_blocking`. All traffic is plain HTTP
/// over `IDP_INTERNAL_URL` (localhost), so there is no TLS. `ureq` keeps no cookie
/// jar, so cookies are set/read by hand. The password is never logged.
pub fn password_login(email: &str, password: &str) -> Result<PasswordLogin, String> {
    let cfg = config::get();
    // A dedicated agent that does NOT follow redirects: every step needs to read
    // the raw `Set-Cookie`/`Location` of the 3xx itself, not chase it.
    let agent = ureq::builder().redirects(0).build();
    let login_url = format!("{}/login", cfg.idp_internal_url);

    // 1. GET /login → grab the CSRF token (form field + cookie).
    let login_page = agent
        .get(&login_url)
        .call()
        .map_err(|e| format!("GET /login failed: {e}"))?;
    let login_cookies: Vec<String> = login_page
        .all("set-cookie")
        .into_iter()
        .map(String::from)
        .collect();
    let csrf_cookie = set_cookie_value(&login_cookies, "_csrf");
    let html = login_page
        .into_string()
        .map_err(|e| format!("read /login body failed: {e}"))?;
    let csrf_form = parse_csrf_input(&html);

    // The IDP renders the cookie token into the form, so the two are equal; fall
    // back to whichever we managed to read. The CSRF check is a double-submit —
    // all that matters is that the cookie and the form field we send back match.
    let csrf_for_cookie = csrf_cookie
        .clone()
        .or_else(|| csrf_form.clone())
        .ok_or_else(|| "no _csrf token on /login".to_string())?;
    let csrf_for_form = csrf_form
        .or(csrf_cookie)
        .ok_or_else(|| "no _csrf token on /login".to_string())?;

    // 2. POST /login (form-encoded) with the CSRF cookie + field.
    let submit = agent
        .post(&login_url)
        .set("Cookie", &format!("_csrf={csrf_for_cookie}"))
        .send_form(&[
            ("email", email),
            ("password", password),
            ("_csrf", &csrf_for_form),
        ]);
    let submit = match submit {
        Ok(r) => r,
        // >=400 (e.g. a 403 CSRF rejection) surfaces here — treat as a hard error.
        Err(ureq::Error::Status(code, _)) => return Err(format!("POST /login returned {code}")),
        Err(e) => return Err(format!("POST /login failed: {e}")),
    };
    let submit_cookies: Vec<String> = submit
        .all("set-cookie")
        .into_iter()
        .map(String::from)
        .collect();
    let submit_location = submit.header("location").map(String::from);

    let idp_session = match set_cookie_value(&submit_cookies, "idp_session") {
        Some(s) => s,
        None => {
            // No session cookie: either MFA (redirect to /mfa) or bad credentials
            // (the login page re-rendered, redirect back to /login).
            if submit_location.as_deref().is_some_and(is_mfa_redirect) {
                return Ok(PasswordLogin::MfaRequired);
            }
            return Ok(PasswordLogin::InvalidCredentials);
        }
    };

    // 3. GET /authorize with the session cookie; read the code from the redirect.
    let state = crate::util::token();
    let query = form_urlencoded::Serializer::new(String::new())
        .append_pair("client_id", &cfg.oidc_client_id)
        .append_pair("redirect_uri", &cfg.oidc_redirect_uri)
        .append_pair("response_type", "code")
        .append_pair("scope", "openid profile email groups")
        .append_pair("state", &state)
        .finish()
        .replace('+', "%20");
    let authorize_url = format!("{}/authorize?{query}", cfg.idp_internal_url);

    let authorized = agent
        .get(&authorize_url)
        .set("Cookie", &format!("idp_session={idp_session}"))
        .call();
    let authorized = match authorized {
        Ok(r) => r,
        Err(ureq::Error::Status(code, _)) => {
            return Err(format!("GET /authorize returned {code}"));
        }
        Err(e) => return Err(format!("GET /authorize failed: {e}")),
    };
    let location = authorized
        .header("location")
        .map(String::from)
        .ok_or_else(|| "no Location on /authorize response".to_string())?;
    let code = code_from_redirect(&location)
        .ok_or_else(|| "no code in /authorize redirect".to_string())?;

    // 4. Same token + userinfo exchange the browser callback uses.
    let info = exchange_code_for_userinfo(&code)?;
    Ok(PasswordLogin::Ok(info))
}

/// Pull the `value` of the first `<input ... name="_csrf" ... value="…">` in the
/// IDP login page. The IDP renders `name` before `value` in the same tag.
fn parse_csrf_input(html: &str) -> Option<String> {
    let name_at = html.find("name=\"_csrf\"")?;
    let after = &html[name_at..];
    let val_at = after.find("value=\"")?;
    let rest = &after[val_at + "value=\"".len()..];
    let end = rest.find('"')?;
    let value = &rest[..end];
    (!value.is_empty()).then(|| value.to_string())
}

/// Read a cookie's value from a list of `Set-Cookie` header lines. Each line is
/// `name=value; Attr; Attr` — we match the leading `name=` and take everything up
/// to the first `;`. Empty values (cookie-clearing) are treated as absent.
fn set_cookie_value(set_cookies: &[String], name: &str) -> Option<String> {
    let prefix = format!("{name}=");
    for line in set_cookies {
        let pair = line.split(';').next().unwrap_or("").trim();
        if let Some(value) = pair.strip_prefix(&prefix)
            && !value.is_empty()
        {
            return Some(value.to_string());
        }
    }
    None
}

/// True if a `Location` points at the IDP's `/mfa` page (path component is `/mfa`),
/// whether it is a bare path (`/mfa?…`) or an absolute URL.
fn is_mfa_redirect(location: &str) -> bool {
    let path = location.split(['?', '#']).next().unwrap_or("");
    path == "/mfa" || path.ends_with("/mfa")
}

/// Extract the `code` query parameter from an OAuth redirect `Location`.
fn code_from_redirect(location: &str) -> Option<String> {
    let query = location.split_once('?').map(|(_, q)| q)?;
    form_urlencoded::parse(query.as_bytes())
        .find(|(k, _)| k == "code")
        .map(|(_, v)| v.into_owned())
        .filter(|c| !c.is_empty())
}

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

    #[test]
    fn parses_csrf_from_login_form() {
        // Mirrors the IDP's rendered hidden input.
        let html = r#"<form method="post" action="/login">
  <input type="hidden" name="_csrf" value="abc123DEF">
  <input type="hidden" name="return_to" value="">
  <button type="submit">Sign in</button>
</form>"#;
        assert_eq!(parse_csrf_input(html), Some("abc123DEF".to_string()));
    }

    #[test]
    fn missing_or_empty_csrf_is_none() {
        assert_eq!(parse_csrf_input("<form></form>"), None);
        assert_eq!(parse_csrf_input(r#"<input name="_csrf" value="">"#), None);
    }

    #[test]
    fn reads_named_cookie_value() {
        let cookies = vec![
            "_csrf=tok-xyz; Path=/; HttpOnly; SameSite=Strict; Secure".to_string(),
            "idp_session=sess-789; Path=/; HttpOnly; SameSite=Lax; Secure; Max-Age=86400"
                .to_string(),
        ];
        assert_eq!(
            set_cookie_value(&cookies, "_csrf"),
            Some("tok-xyz".to_string())
        );
        assert_eq!(
            set_cookie_value(&cookies, "idp_session"),
            Some("sess-789".to_string())
        );
        assert_eq!(set_cookie_value(&cookies, "missing"), None);
    }

    #[test]
    fn cleared_cookie_reads_as_absent() {
        let cookies = vec!["idp_session=; Path=/; Max-Age=0".to_string()];
        assert_eq!(set_cookie_value(&cookies, "idp_session"), None);
    }

    #[test]
    fn detects_mfa_redirect() {
        assert!(is_mfa_redirect("/mfa?return_to=%2Fauthorize"));
        assert!(is_mfa_redirect("https://auth.olibuijr.com/mfa"));
        assert!(!is_mfa_redirect(
            "/login?error=Invalid%20email%20or%20password"
        ));
        assert!(!is_mfa_redirect(
            "https://akurai-drive.olibuijr.com/auth/callback?code=x&state=y"
        ));
    }

    #[test]
    fn extracts_code_from_redirect() {
        assert_eq!(
            code_from_redirect(
                "https://akurai-drive.olibuijr.com/auth/callback?code=AUTHCODE&state=st"
            ),
            Some("AUTHCODE".to_string())
        );
        assert_eq!(
            code_from_redirect("/login?error=Invalid%20email%20or%20password"),
            None
        );
        assert_eq!(code_from_redirect("https://x/cb?code=&state=y"), None);
    }
}