AkurAI-IDP / src/lib/organizations.rs

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

use rusqlite::{Connection, OptionalExtension};

/// Version of the organization claims emitted by this IDP.
pub const CLAIMS_VERSION: i64 = 1;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Membership {
    pub organization_id: String,
    pub user_id: String,
    pub role: String,
    pub authorization_version: i64,
}

pub fn active_membership(
    conn: &Connection,
    user_id: &str,
    organization_id: &str,
) -> rusqlite::Result<Option<Membership>> {
    conn.query_row(
        "SELECT m.organization_id, m.user_id, m.role, u.authorization_version
         FROM organization_memberships m
         JOIN users u ON u.id = m.user_id
         WHERE m.organization_id = ?1 AND m.user_id = ?2 AND m.status = 'active'",
        rusqlite::params![organization_id, user_id],
        |row| {
            Ok(Membership {
                organization_id: row.get(0)?,
                user_id: row.get(1)?,
                role: row.get(2)?,
                authorization_version: row.get(3)?,
            })
        },
    )
    .optional()
}

pub fn user_authorization_version(conn: &Connection, user_id: &str) -> rusqlite::Result<i64> {
    conn.query_row(
        "SELECT authorization_version FROM users WHERE id = ?1",
        [user_id],
        |row| row.get(0),
    )
}

/// Invalidate every existing session and token for a user.  JWTs carry this
/// version and are checked against the authoritative row on every request.
pub fn bump_authorization_version(conn: &Connection, user_id: &str) -> rusqlite::Result<()> {
    conn.execute(
        "UPDATE users SET authorization_version = authorization_version + 1, updated_at = ?1
         WHERE id = ?2",
        rusqlite::params![chrono::Utc::now().timestamp(), user_id],
    )?;
    Ok(())
}

pub fn valid_role(role: &str) -> bool {
    matches!(role, "owner" | "admin" | "member")
}

pub fn can_administer(role: &str) -> bool {
    matches!(role, "owner" | "admin")
}

pub fn normalize_email(email: &str) -> String {
    email.trim().to_ascii_lowercase()
}
/// Keep the legacy primary-organization alias usable when a user leaves or
/// loses access to that organization.  Secondary memberships remain
/// authoritative and do not get copied into tokens.
pub fn move_primary_organization_if_needed(
    conn: &Connection,
    user_id: &str,
    organization_id: &str,
) -> rusqlite::Result<bool> {
    let primary: Option<String> = conn
        .query_row(
            "SELECT tenant_id FROM users WHERE id = ?1",
            [user_id],
            |row| row.get(0),
        )
        .optional()?;
    if primary.as_deref() != Some(organization_id) {
        return Ok(false);
    }
    let replacement: Option<String> = conn
        .query_row(
            "SELECT organization_id FROM organization_memberships
             WHERE user_id = ?1 AND status = 'active' AND organization_id <> ?2
             ORDER BY created_at, organization_id LIMIT 1",
            rusqlite::params![user_id, organization_id],
            |row| row.get(0),
        )
        .optional()?;
    let Some(replacement) = replacement else {
        return Ok(false);
    };
    conn.execute(
        "UPDATE users SET tenant_id = ?, updated_at = ? WHERE id = ?",
        rusqlite::params![replacement, chrono::Utc::now().timestamp(), user_id],
    )?;
    Ok(true)
}

/// Return whether an administrator's first login still needs its durable marker.
///
/// The check is intentionally separate from [`mark_first_admin_login`]: callers
/// enqueue the deterministic event first and persist the marker only after a
/// successful or duplicate-safe enqueue. A queue outage therefore leaves the
/// marker pending for a later login retry.
pub fn first_admin_login_pending(
    conn: &Connection,
    organization_id: &str,
    user_id: &str,
) -> rusqlite::Result<bool> {
    let is_admin: Option<String> = conn
        .query_row(
            "SELECT role FROM organization_memberships
             WHERE organization_id = ?1 AND user_id = ?2 AND status = 'active'",
            rusqlite::params![organization_id, user_id],
            |row| row.get(0),
        )
        .optional()?;
    if !is_admin.as_deref().is_some_and(can_administer) {
        return Ok(false);
    }
    let marker_exists: Option<i64> = conn
        .query_row(
            "SELECT 1 FROM organization_lifecycle_markers
             WHERE organization_id = ?1 AND marker = 'first-admin-login'",
            [organization_id],
            |row| row.get(0),
        )
        .optional()?;
    Ok(marker_exists.is_none())
}

/// Persist the first-admin marker after the deterministic event was accepted.
pub fn mark_first_admin_login(
    conn: &Connection,
    organization_id: &str,
    user_id: &str,
    occurred_at: i64,
) -> rusqlite::Result<bool> {
    let changed = conn.execute(
        "INSERT OR IGNORE INTO organization_lifecycle_markers
         (organization_id, marker, user_id, occurred_at)
         VALUES (?1, 'first-admin-login', ?2, ?3)",
        rusqlite::params![organization_id, user_id, occurred_at],
    )?;
    Ok(changed == 1)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::db::schema::create_tables;

    fn fixture() -> Connection {
        let conn = Connection::open_in_memory().expect("in-memory db");
        create_tables(&conn);
        conn.execute(
            "INSERT INTO tenants (id, name, slug, created_at) VALUES ('org-a', 'A', 'a', 1)",
            [],
        )
        .expect("tenant");
        conn.execute(
            "INSERT INTO users
             (id, tenant_id, email, password_hash, created_at, updated_at)
             VALUES ('user-a', 'org-a', 'a@example.test', 'hash', 1, 1)",
            [],
        )
        .expect("user");
        create_tables(&conn);
        conn
    }

    #[test]
    fn display_name_and_email_changes_preserve_canonical_ids() {
        let conn = fixture();
        let before = active_membership(&conn, "user-a", "org-a")
            .expect("lookup")
            .expect("membership");
        conn.execute(
            "UPDATE organizations SET display_name = 'Renamed' WHERE id = 'org-a'",
            [],
        )
        .expect("organization rename");
        conn.execute(
            "UPDATE users SET email = 'new@example.test' WHERE id = 'user-a'",
            [],
        )
        .expect("email rename");
        let after = active_membership(&conn, "user-a", "org-a")
            .expect("lookup")
            .expect("membership");
        assert_eq!(before.organization_id, after.organization_id);
        assert_eq!(before.user_id, after.user_id);
    }

    #[test]
    fn membership_lookup_is_tenant_bound() {
        let conn = fixture();
        conn.execute(
            "INSERT INTO organizations (id, display_name, slug, created_at, updated_at)
             VALUES ('org-b', 'B', 'b', 1, 1)",
            [],
        )
        .expect("second organization");
        assert!(
            active_membership(&conn, "user-a", "org-b")
                .expect("lookup")
                .is_none()
        );
    }

    #[test]
    fn authorization_version_invalidates_stale_membership() {
        let conn = fixture();
        let before = active_membership(&conn, "user-a", "org-a")
            .expect("lookup")
            .expect("membership");
        bump_authorization_version(&conn, "user-a").expect("bump");
        let after = active_membership(&conn, "user-a", "org-a")
            .expect("lookup")
            .expect("membership");
        assert_ne!(before.authorization_version, after.authorization_version);
    }
    #[test]
    fn revoked_membership_is_not_authorized() {
        let conn = fixture();
        conn.execute(
            "UPDATE organization_memberships SET status = 'revoked'
             WHERE organization_id = 'org-a' AND user_id = 'user-a'",
            [],
        )
        .expect("revoke membership");
        assert!(
            active_membership(&conn, "user-a", "org-a")
                .expect("lookup")
                .is_none()
        );
    }
    #[test]
    fn first_admin_login_marker_is_idempotent() {
        let conn = fixture();
        assert!(first_admin_login_pending(&conn, "org-a", "user-a").expect("first marker pending"));
        assert!(mark_first_admin_login(&conn, "org-a", "user-a", 10).expect("first marker"));
        assert!(!first_admin_login_pending(&conn, "org-a", "user-a").expect("duplicate marker"));
        assert!(!mark_first_admin_login(&conn, "org-a", "user-a", 11).expect("duplicate marker"));
        let count: i64 = conn
            .query_row(
                "SELECT COUNT(*) FROM organization_lifecycle_markers
                 WHERE organization_id = 'org-a' AND marker = 'first-admin-login'",
                [],
                |row| row.get(0),
            )
            .expect("marker count");
        assert_eq!(count, 1);
    }

    #[test]
    fn first_admin_enqueue_then_marker_is_race_safe() {
        let conn = fixture();
        assert!(first_admin_login_pending(&conn, "org-a", "user-a").expect("pending"));
        assert!(first_admin_login_pending(&conn, "org-a", "user-a").expect("pending race"));
        assert!(mark_first_admin_login(&conn, "org-a", "user-a", 10).expect("first marker"));
        assert!(!mark_first_admin_login(&conn, "org-a", "user-a", 11).expect("duplicate marker"));
        assert!(!first_admin_login_pending(&conn, "org-a", "user-a").expect("marker persisted"));
    }
}