AkurAI-IDP / src/routes/auth_pages.rs

Latest change e1794828e5fae9a8545a66e77b2e8200cf8c090b - feat: add organization lifecycle telemetry by Ólafur Búi Ólafsson

use axum::{
    Extension, Form, Router,
    extract::Query,
    http::{HeaderMap, StatusCode, header},
    response::{Html, IntoResponse, Response},
    routing::get,
};
use serde::Deserialize;
use std::time::{SystemTime, UNIX_EPOCH};

use crate::config;
use crate::db::with_db;
use crate::lib::audit::{
    SESSION_CREATED, USER_LOCKED, USER_LOGIN, USER_LOGIN_FAILED, log_audit_event,
};
use crate::lib::crypto::{generate_secure_token, sha256};
use crate::lib::html::{Locale, auth_page, esc_html, t};
use crate::lib::organizations::{first_admin_login_pending, mark_first_admin_login};
use crate::lib::password::verify_password;
use crate::lib::telemetry::OrganizationTelemetry;
use crate::lib::totp::verify_totp;

pub fn router() -> Router {
    Router::new()
        .route("/login", get(login_page).post(login_submit))
        .route("/mfa", get(mfa_page).post(mfa_submit))
        .route("/logout", get(logout).post(logout))
}

// ---------------------------------------------------------------------------
// Query / form structs
// ---------------------------------------------------------------------------

#[derive(Deserialize, Default)]
struct LoginQuery {
    error: Option<String>,
    return_to: Option<String>,
}

#[derive(Deserialize)]
struct LoginForm {
    email: Option<String>,
    password: Option<String>,
    return_to: Option<String>,
    _csrf: Option<String>,
}

#[derive(Deserialize)]
struct MfaForm {
    code: Option<String>,
    _csrf: Option<String>,
}

// ---------------------------------------------------------------------------
// GET /login
// ---------------------------------------------------------------------------

async fn login_page(Query(q): Query<LoginQuery>, headers: HeaderMap) -> impl IntoResponse {
    let locale =
        Locale::from_cookie_header(headers.get(header::COOKIE).and_then(|v| v.to_str().ok()));
    let csrf = get_csrf_cookie(&headers).unwrap_or_default();
    let return_to = q.return_to.as_deref().unwrap_or("");
    let error_html = match &q.error {
        Some(msg) => format!(r#"<div class="error">{}</div>"#, esc_html(msg)),
        None => String::new(),
    };

    let lbl_email = t(locale, "Netfang", "Email");
    let lbl_password = t(locale, "Lykilorð", "Password");
    let lbl_submit = t(locale, "Skrá inn", "Sign in");
    let page_title = t(locale, "Skrá inn", "Sign in");

    let body = format!(
        r#"{error_html}
<form method="post" action="/login">
  <input type="hidden" name="_csrf" value="{csrf}">
  <input type="hidden" name="return_to" value="{return_to}">
  <label for="email">{lbl_email}</label>
  <input type="email" id="email" name="email" required autofocus>
  <label for="password">{lbl_password}</label>
  <input type="password" id="password" name="password" required>
  <button type="submit">{lbl_submit}</button>
</form>"#,
        csrf = esc_html(&csrf),
        return_to = esc_html(return_to),
        lbl_email = lbl_email,
        lbl_password = lbl_password,
        lbl_submit = lbl_submit,
    );

    Html(auth_page(locale, page_title, &body))
}

// ---------------------------------------------------------------------------
// POST /login
// ---------------------------------------------------------------------------

async fn login_submit(
    telemetry: Option<Extension<OrganizationTelemetry>>,
    headers: HeaderMap,
    Form(form): Form<LoginForm>,
) -> Response {
    let email = form.email.as_deref().unwrap_or("").trim().to_lowercase();
    let telemetry = telemetry.map(|Extension(value)| value);
    let password = form.password.as_deref().unwrap_or("");
    let return_to = safe_return_to(form.return_to.as_deref());
    let ip = extract_ip(&headers);

    if email.is_empty() || password.is_empty() {
        return login_redirect("Email and password are required", &return_to);
    }

    // Global lookup across tenants (email not globally unique, but we pick first match)
    let user = with_db(|conn| {
        let mut stmt = conn
            .prepare(
                "SELECT id, tenant_id, password_hash, mfa_enabled,
                        locked_until, failed_attempts
                 FROM users WHERE LOWER(email) = ?1 LIMIT 1",
            )
            .ok()?;
        stmt.query_row(rusqlite::params![email], |row| {
            Ok(UserRow {
                id: row.get(0)?,
                tenant_id: row.get(1)?,
                password_hash: row.get(2)?,
                mfa_enabled: row.get::<_, i64>(3)? != 0,
                locked_until: row.get(4)?,
                failed_attempts: row.get::<_, i64>(5)?,
            })
        })
        .ok()
    });

    let user = match user {
        Some(u) => u,
        None => {
            // Keep unknown-user and wrong-password paths on the same expensive
            // password-processing path without revealing account existence.
            let _ = crate::lib::password::hash_password(password);
            return login_redirect("Invalid email or password", &return_to);
        }
    };

    let now = now_epoch();

    // Check lockout
    if let Some(locked_until) = user.locked_until
        && locked_until > now
    {
        return login_redirect("Account temporarily locked. Try again later.", &return_to);
    }

    // Verify password
    let pw_ok = verify_password(&user.password_hash, password).unwrap_or(false);
    if !pw_ok {
        let new_attempts = user.failed_attempts + 1;
        with_db(|conn| {
            if new_attempts >= 5 {
                let lock_until = now + 15 * 60; // 15 minutes
                conn.execute(
                    "UPDATE users SET failed_attempts = ?1, locked_until = ?2, updated_at = ?3 WHERE id = ?4",
                    rusqlite::params![new_attempts, lock_until, now, user.id],
                )
                .ok();
                log_audit_event(
                    conn,
                    Some(&user.tenant_id),
                    Some(&user.id),
                    USER_LOCKED,
                    Some(&ip),
                    None,
                );
            } else {
                conn.execute(
                    "UPDATE users SET failed_attempts = ?1, updated_at = ?2 WHERE id = ?3",
                    rusqlite::params![new_attempts, now, user.id],
                )
                .ok();
            }
            log_audit_event(
                conn,
                Some(&user.tenant_id),
                Some(&user.id),
                USER_LOGIN_FAILED,
                Some(&ip),
                None,
            );
        });
        return login_redirect("Invalid email or password", &return_to);
    }

    // Password correct — reset failed attempts
    with_db(|conn| {
        conn.execute(
            "UPDATE users SET failed_attempts = 0, locked_until = NULL, updated_at = ?1 WHERE id = ?2",
            rusqlite::params![now, user.id],
        )
        .ok();
    });

    if user.mfa_enabled {
        let token = generate_secure_token(32);
        let token_hash = sha256(&token);
        let pending_saved = with_db(|conn| {
            conn.execute(
                "INSERT INTO mfa_pending (token_hash, user_id, created_at, attempts)
                 VALUES (?1, ?2, ?3, 0)",
                rusqlite::params![token_hash, user.id, now],
            )
        });
        if pending_saved.is_err() {
            return (StatusCode::INTERNAL_SERVER_ERROR, "internal server error").into_response();
        }
        let cookie = format!(
            "idp_mfa_pending={token}; Path=/; HttpOnly; SameSite=Strict; Secure; Max-Age=300"
        );
        let redirect_url = format!("/mfa?return_to={}", urlencoded(&return_to));
        return (
            StatusCode::SEE_OTHER,
            [
                (header::LOCATION, redirect_url),
                (header::SET_COOKIE, cookie),
            ],
        )
            .into_response();
    }

    // No MFA — create session directly
    create_session_response(
        &user.id,
        &user.tenant_id,
        &ip,
        &extract_ua(&headers),
        &return_to,
        telemetry.as_ref(),
    )
}

// ---------------------------------------------------------------------------
// GET /mfa
// ---------------------------------------------------------------------------

async fn mfa_page(Query(q): Query<LoginQuery>, headers: HeaderMap) -> impl IntoResponse {
    let locale =
        Locale::from_cookie_header(headers.get(header::COOKIE).and_then(|v| v.to_str().ok()));
    let csrf = get_csrf_cookie(&headers).unwrap_or_default();
    let error_html = match &q.error {
        Some(msg) => format!(r#"<div class="error">{}</div>"#, esc_html(msg)),
        None => String::new(),
    };

    let lbl_code = t(locale, "Auðkenniskóði", "Authentication code");
    let lbl_hint = t(
        locale,
        "Sláðu inn 6 stafa auðkenniskóðann þinn eða varakóða.",
        "Enter your 6-digit authenticator code or a backup code.",
    );
    let lbl_submit = t(locale, "Staðfesta", "Verify");
    let page_title = t(locale, "Tvíþátta auðkenning", "Two-factor authentication");

    let body = format!(
        r#"{error_html}
<form method="post" action="/mfa">
  <input type="hidden" name="_csrf" value="{csrf}">
  <label for="code">{lbl_code}</label>
  <input type="text" id="code" name="code" inputmode="numeric" pattern="[0-9]*" autocomplete="one-time-code" required autofocus>
  <p class="hint">{lbl_hint}</p>
  <button type="submit">{lbl_submit}</button>
</form>"#,
        csrf = esc_html(&csrf),
        lbl_code = lbl_code,
        lbl_hint = lbl_hint,
        lbl_submit = lbl_submit,
    );

    Html(auth_page(locale, page_title, &body))
}

// ---------------------------------------------------------------------------
// POST /mfa
// ---------------------------------------------------------------------------

async fn mfa_submit(
    telemetry: Option<Extension<OrganizationTelemetry>>,
    headers: HeaderMap,
    Form(form): Form<MfaForm>,
) -> Response {
    let code = form.code.as_deref().unwrap_or("").trim().to_string();
    let ip = extract_ip(&headers);
    let telemetry = telemetry.map(|Extension(value)| value);

    let pending = match get_cookie(&headers, "idp_mfa_pending") {
        Some(v) if !v.is_empty() => v,
        _ => return login_redirect("MFA session expired. Please log in again.", ""),
    };
    let pending_hash = sha256(&pending);
    let (uid, tenant_id, mfa_secret) = match with_db(|conn| {
        conn.query_row(
            "SELECT u.id, u.tenant_id, u.mfa_secret
             FROM mfa_pending p JOIN users u ON u.id = p.user_id
             WHERE p.token_hash = ?1 AND p.created_at > ?2 AND p.attempts < 5",
            rusqlite::params![pending_hash, now_epoch() - 300],
            |row| {
                Ok((
                    row.get::<_, String>(0)?,
                    row.get::<_, String>(1)?,
                    row.get::<_, Option<String>>(2)?,
                ))
            },
        )
    }) {
        Ok(value) => value,
        Err(_) => return login_redirect("MFA session expired. Please log in again.", ""),
    };
    let secret = match mfa_secret {
        Some(s) if !s.is_empty() => s,
        _ => return login_redirect("MFA not configured.", ""),
    };
    let attempt = with_db(|conn| {
        conn.execute(
            "UPDATE mfa_pending SET attempts = attempts + 1 WHERE token_hash = ?1",
            [&pending_hash],
        )
    });
    if attempt.is_err() {
        return (StatusCode::INTERNAL_SERVER_ERROR, "internal server error").into_response();
    }
    let totp_ok = verify_totp(&secret, &code);
    if !totp_ok {
        let code_hash = sha256(&code);
        let backup_ok = with_db(|conn| {
            let changed = conn
                .execute(
                    "UPDATE backup_codes SET used = 1 WHERE user_id = ?1 AND code_hash = ?2 AND used = 0",
                    rusqlite::params![uid, code_hash],
                )
                .unwrap_or(0);
            changed > 0
        });
        if !backup_ok {
            return (
                StatusCode::SEE_OTHER,
                [(
                    header::LOCATION,
                    format!(
                        "/mfa?error={}",
                        urlencoded("Invalid code. Please try again.")
                    ),
                )],
            )
                .into_response();
        }
    }
    let _ = with_db(|conn| {
        conn.execute(
            "DELETE FROM mfa_pending WHERE token_hash = ?1",
            [&pending_hash],
        )
    });
    let clear_pending =
        "idp_mfa_pending=; Path=/; HttpOnly; SameSite=Strict; Secure; Max-Age=0".to_string();
    let return_to = safe_return_to(None);
    let mut resp = create_session_response(
        &uid,
        &tenant_id,
        &ip,
        &extract_ua(&headers),
        &return_to,
        telemetry.as_ref(),
    );
    resp.headers_mut()
        .append(header::SET_COOKIE, clear_pending.parse().unwrap());
    resp
}

// ---------------------------------------------------------------------------
// GET|POST /logout
// ---------------------------------------------------------------------------

async fn logout(headers: HeaderMap) -> Response {
    let ip = extract_ip(&headers);

    if let Some(session_id) = get_cookie(&headers, "idp_session") {
        with_db(|conn| {
            // Get user info for audit before deleting
            let user_info: Option<(String, String)> = conn
                .prepare("SELECT user_id, (SELECT tenant_id FROM users WHERE id = sessions.user_id) FROM sessions WHERE id = ?1")
                .ok()
                .and_then(|mut stmt| {
                    stmt.query_row(rusqlite::params![session_id], |row| {
                        Ok((row.get(0)?, row.get(1)?))
                    })
                    .ok()
                });

            conn.execute(
                "DELETE FROM sessions WHERE id = ?1",
                rusqlite::params![session_id],
            )
            .ok();

            if let Some((user_id, tenant_id)) = user_info {
                log_audit_event(
                    conn,
                    Some(&tenant_id),
                    Some(&user_id),
                    crate::lib::audit::SESSION_REVOKED,
                    Some(&ip),
                    None,
                );
            }
        });
    }

    let clear = "idp_session=; Path=/; HttpOnly; SameSite=Lax; Secure; Max-Age=0";
    (
        StatusCode::SEE_OTHER,
        [
            (header::LOCATION, "/login".to_string()),
            (header::SET_COOKIE, clear.to_string()),
        ],
    )
        .into_response()
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

struct UserRow {
    id: String,
    tenant_id: String,
    password_hash: String,
    mfa_enabled: bool,
    locked_until: Option<i64>,
    failed_attempts: i64,
}

fn create_session_response(
    user_id: &str,
    tenant_id: &str,
    ip: &str,
    user_agent: &str,
    return_to: &str,
    telemetry: Option<&OrganizationTelemetry>,
) -> Response {
    let session_id = generate_secure_token(32);
    let now = now_epoch();
    let expires = now + 24 * 3600; // 24 hours
    let telemetry_enabled = telemetry.is_some();
    let created = with_db(|conn| -> rusqlite::Result<bool> {
        let authorization_version: i64 = conn.query_row(
            "SELECT authorization_version FROM users WHERE id = ?1",
            [user_id],
            |row| row.get(0),
        )?;
        conn.execute(
            "INSERT INTO sessions
             (id, user_id, ip, user_agent, authorization_version,
              active_organization_id, created_at, expires_at)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
            rusqlite::params![
                session_id,
                user_id,
                ip,
                user_agent,
                authorization_version,
                tenant_id,
                now,
                expires
            ],
        )?;
        log_audit_event(
            conn,
            Some(tenant_id),
            Some(user_id),
            SESSION_CREATED,
            Some(ip),
            None,
        );
        log_audit_event(
            conn,
            Some(tenant_id),
            Some(user_id),
            USER_LOGIN,
            Some(ip),
            None,
        );
        let first_admin_login = first_admin_login_pending(conn, tenant_id, user_id)?;
        if first_admin_login && !telemetry_enabled {
            mark_first_admin_login(conn, tenant_id, user_id, now)?;
        }
        Ok(first_admin_login)
    });
    if let Err(error) = created {
        tracing::error!(%error, "failed to create session");
        return (StatusCode::INTERNAL_SERVER_ERROR, "internal server error").into_response();
    }
    if created.as_ref().is_ok_and(|first| *first)
        && let Some(telemetry) = telemetry
    {
        let outcome = telemetry.first_admin_login(tenant_id, now);
        if matches!(
            outcome,
            akurai_telemetry::EmitOutcome::Queued | akurai_telemetry::EmitOutcome::Duplicate
        ) && with_db(|conn| mark_first_admin_login(conn, tenant_id, user_id, now)).is_err()
        {
            tracing::warn!("first-admin login marker persistence failed; retrying later");
        }
    }

    let cookie =
        format!("idp_session={session_id}; Path=/; HttpOnly; SameSite=Lax; Secure; Max-Age=86400");
    let dest = if return_to.is_empty() { "/" } else { return_to };
    let escaped_dest = esc_html(dest);
    let html = format!(
        r#"<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta http-equiv="refresh" content="0; url={escaped_dest}">
  <title>Signing in - AkurAI ID</title>
</head>
<body>
  <p>Signing in... <a href="{escaped_dest}">Continue</a></p>
</body>
</html>"#
    );

    (
        StatusCode::OK,
        [
            (header::SET_COOKIE, cookie),
            (header::CONTENT_TYPE, "text/html; charset=utf-8".to_string()),
        ],
        html,
    )
        .into_response()
}

/// Validate return_to: must start with "/" and not "//", or match base_url origin.
fn safe_return_to(input: Option<&str>) -> String {
    let val = input.unwrap_or("/").trim();
    if val.is_empty() {
        return "/".to_string();
    }

    // Relative path: must start with "/" but not "//" (protocol-relative)
    if val.starts_with('/') && !val.starts_with("//") {
        return val.to_string();
    }

    // Absolute URL: must match our base_url origin
    let base = &config::get().base_url;
    if let Some(origin) = extract_origin(base)
        && val.starts_with(&origin)
    {
        return val.to_string();
    }

    "/".to_string()
}

fn extract_origin(url: &str) -> Option<String> {
    // Extract scheme + host from a URL like "https://auth.example.com/path"
    let after_scheme = url
        .strip_prefix("https://")
        .map(|rest| ("https", rest))
        .or_else(|| url.strip_prefix("http://").map(|rest| ("http", rest)))?;
    let (scheme, rest) = after_scheme;
    let host = rest.split('/').next()?;
    Some(format!("{scheme}://{host}"))
}

fn now_epoch() -> i64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs() as i64
}

fn get_cookie(headers: &HeaderMap, name: &str) -> Option<String> {
    let cookies = headers.get(header::COOKIE)?.to_str().ok()?;
    let prefix = format!("{name}=");
    for part in cookies.split(';') {
        let part = part.trim();
        if let Some(value) = part.strip_prefix(&prefix) {
            return Some(value.to_string());
        }
    }
    None
}

fn get_csrf_cookie(headers: &HeaderMap) -> Option<String> {
    get_cookie(headers, "_csrf")
}

fn extract_ip(headers: &HeaderMap) -> String {
    // Check X-Forwarded-For, X-Real-IP, then fall back to empty
    headers
        .get("x-forwarded-for")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.split(',').next())
        .map(|s| s.trim().to_string())
        .or_else(|| {
            headers
                .get("x-real-ip")
                .and_then(|v| v.to_str().ok())
                .map(|s| s.trim().to_string())
        })
        .unwrap_or_default()
}

fn extract_ua(headers: &HeaderMap) -> String {
    headers
        .get(header::USER_AGENT)
        .and_then(|v| v.to_str().ok())
        .unwrap_or("")
        .to_string()
}

fn login_redirect(error: &str, return_to: &str) -> Response {
    let mut url = format!("/login?error={}", urlencoded(error));
    if !return_to.is_empty() && return_to != "/" {
        url.push_str(&format!("&return_to={}", urlencoded(return_to)));
    }
    (StatusCode::SEE_OTHER, [(header::LOCATION, url)]).into_response()
}

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

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

    #[tokio::test]
    async fn login_request_succeeds_without_telemetry_extension() {
        let response = login_page(Query(LoginQuery::default()), HeaderMap::new())
            .await
            .into_response();
        assert_eq!(response.status(), StatusCode::OK);
    }
}