AkurAI-IDP / src/routes/organizations.rs

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

use axum::{
    Extension, Json, Router,
    extract::Path,
    http::StatusCode,
    response::{IntoResponse, Response},
    routing::{get, post},
};
use rusqlite::OptionalExtension;
use serde::Deserialize;
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use uuid::Uuid;

use crate::db::with_db;
use crate::lib::crypto::generate_secure_token;
use crate::lib::organizations::{
    active_membership, bump_authorization_version, can_administer,
    move_primary_organization_if_needed, normalize_email, valid_role,
};
use crate::lib::telemetry::OrganizationTelemetry;
use crate::middleware::auth::AuthUser;

#[derive(Deserialize)]
struct CreateOrganization {
    display_name: String,
    slug: String,
    domain: Option<String>,
}

#[derive(Deserialize)]
struct InviteMember {
    email: String,
    role: String,
}

#[derive(Deserialize)]
struct AcceptInvitation {
    token: String,
}

#[derive(Deserialize)]
struct TransferAdmin {
    user_id: String,
}

#[derive(Deserialize)]
struct LinkWorkspace {
    workspace_id: String,
    installation_id: Option<String>,
    app_id: Option<String>,
}

pub fn router() -> Router {
    Router::new()
        .route("/", get(list_organizations).post(create_organization))
        .route("/invitations/accept", post(accept_invitation))
        .route("/{organization_id}", get(get_organization))
        .route("/{organization_id}/invitations", post(invite_member))
        .route(
            "/{organization_id}/invitations/{invitation_id}/revoke",
            post(revoke_invitation),
        )
        .route(
            "/{organization_id}/members/{user_id}/revoke",
            post(revoke_member),
        )
        .route("/{organization_id}/leave", post(leave_organization))
        .route("/{organization_id}/admin-transfer", post(transfer_admin))
        .route(
            "/{organization_id}/workspaces",
            get(list_workspaces).post(link_workspace),
        )
        .route(
            "/{organization_id}/workspaces/{workspace_id}",
            post(unlink_workspace),
        )
}

async fn list_organizations(Extension(user): Extension<AuthUser>) -> Response {
    let result: Result<Vec<Value>, rusqlite::Error> = with_db(|conn| {
        let mut stmt = conn.prepare(
            "SELECT o.id, o.display_name, o.slug, o.domain, m.role, o.created_at, o.updated_at
             FROM organizations o
             JOIN organization_memberships m ON m.organization_id = o.id
             WHERE m.user_id = ?1 AND m.status = 'active'
             ORDER BY o.display_name, o.id",
        )?;
        let rows = stmt.query_map([&user.id], |row| {
            Ok(json!({
                "organization_id": row.get::<_, String>(0)?,
                "display_name": row.get::<_, String>(1)?,
                "slug": row.get::<_, String>(2)?,
                "domain": row.get::<_, Option<String>>(3)?,
                "role": row.get::<_, String>(4)?,
                "created_at": row.get::<_, i64>(5)?,
                "updated_at": row.get::<_, i64>(6)?,
            }))
        })?;
        rows.collect::<Result<Vec<_>, _>>()
    });

    match result {
        Ok(organizations) => Json(json!({"organizations": organizations})).into_response(),
        Err(_) => server_error(),
    }
}
async fn create_organization(
    Extension(user): Extension<AuthUser>,
    telemetry: Option<Extension<OrganizationTelemetry>>,
    Json(body): Json<CreateOrganization>,
) -> Response {
    let display_name = body.display_name.trim().to_string();
    let slug = body.slug.trim().to_ascii_lowercase();
    if display_name.is_empty() || slug.is_empty() {
        return error(
            StatusCode::BAD_REQUEST,
            "display_name and slug are required",
        );
    }

    let organization_id = Uuid::new_v4().to_string();
    let now = chrono::Utc::now().timestamp();
    let result: Result<(), rusqlite::Error> = with_db(|conn| {
        let tx = conn.unchecked_transaction()?;
        tx.execute(
            "INSERT INTO tenants (id, name, slug, domain, created_at)
             VALUES (?1, ?2, ?3, ?4, ?5)",
            rusqlite::params![
                organization_id,
                display_name,
                slug,
                body.domain.clone(),
                now
            ],
        )?;
        tx.execute(
            "INSERT INTO organizations
             (id, display_name, slug, domain, created_at, updated_at)
             VALUES (?1, ?2, ?3, ?4, ?5, ?5)",
            rusqlite::params![organization_id, display_name, slug, body.domain, now],
        )?;
        tx.execute(
            "INSERT INTO organization_memberships
             (organization_id, user_id, role, status, created_at, updated_at)
             VALUES (?1, ?2, 'owner', 'active', ?3, ?3)",
            rusqlite::params![organization_id, user.id, now],
        )?;
        bump_authorization_version(&tx, &user.id)?;
        tx.commit()
    });

    match result {
        Ok(()) => {
            if let Some(Extension(telemetry)) = telemetry {
                telemetry.organization_created(&organization_id, now);
                telemetry.active_member_count_after_organization_created(&organization_id, now);
            }
            (
                StatusCode::CREATED,
                Json(json!({
                    "organization_id": organization_id,
                    "display_name": display_name,
                    "slug": slug,
                    "role": "owner"
                })),
            )
                .into_response()
        }
        Err(e) if e.to_string().contains("UNIQUE") => error(
            StatusCode::CONFLICT,
            "organization slug or membership already exists",
        ),
        Err(_) => server_error(),
    }
}

async fn get_organization(
    Extension(user): Extension<AuthUser>,
    Path(organization_id): Path<String>,
) -> Response {
    let result: Result<Option<Value>, rusqlite::Error> = with_db(|conn| {
        let membership = active_membership(conn, &user.id, &organization_id)?;
        let Some(membership) = membership else {
            return Ok(None);
        };
        conn.query_row(
            "SELECT id, display_name, slug, domain, created_at, updated_at
             FROM organizations WHERE id = ?1",
            [&organization_id],
            |row| {
                Ok(json!({
                    "organization_id": row.get::<_, String>(0)?,
                    "display_name": row.get::<_, String>(1)?,
                    "slug": row.get::<_, String>(2)?,
                    "domain": row.get::<_, Option<String>>(3)?,
                    "created_at": row.get::<_, i64>(4)?,
                    "updated_at": row.get::<_, i64>(5)?,
                    "role": membership.role,
                }))
            },
        )
        .optional()
    });

    match result {
        Ok(Some(organization)) => Json(organization).into_response(),
        Ok(None) => error(StatusCode::NOT_FOUND, "organization not found"),
        Err(_) => server_error(),
    }
}

async fn invite_member(
    Extension(user): Extension<AuthUser>,
    telemetry: Option<Extension<OrganizationTelemetry>>,
    Path(organization_id): Path<String>,
    Json(body): Json<InviteMember>,
) -> Response {
    let email = normalize_email(&body.email);
    if email.is_empty() || !valid_role(&body.role) || body.role == "owner" {
        return error(
            StatusCode::BAD_REQUEST,
            "valid email and member/admin role are required",
        );
    }

    let token = generate_secure_token(32);
    let token_hash = hash_token(&token);
    let invitation_id = Uuid::new_v4().to_string();
    let now = chrono::Utc::now().timestamp();
    let expires_at = now + 7 * 24 * 3600;
    let result: Result<Result<(), &str>, rusqlite::Error> = with_db(|conn| {
        let tx = conn.unchecked_transaction()?;
        let membership = active_membership(&tx, &user.id, &organization_id)?;
        if !membership.as_ref().is_some_and(|m| can_administer(&m.role)) {
            return Ok(Err("forbidden"));
        }
        // Rotate any existing pending invitation for this normalized address.
        tx.execute(
            "UPDATE organization_invitations SET revoked_at = ?1
             WHERE organization_id = ?2 AND lower(email) = ?3
               AND accepted_at IS NULL AND revoked_at IS NULL",
            rusqlite::params![now, organization_id, email],
        )?;
        tx.execute(
            "INSERT INTO organization_invitations
             (id, organization_id, email, role, invited_by, token_hash, expires_at, created_at)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
            rusqlite::params![
                invitation_id,
                organization_id,
                email,
                body.role,
                user.id,
                token_hash,
                expires_at,
                now
            ],
        )?;
        // Keep the one-time secret in an internal delivery draft. This route
        // never transmits it; a separately authorized delivery worker may.
        tx.execute(
            "INSERT INTO organization_invitation_delivery_drafts
             (invitation_id, token, created_at, status)
             VALUES (?1, ?2, ?3, 'draft')",
            rusqlite::params![invitation_id, token, now],
        )?;
        tx.commit()?;
        Ok(Ok(()))
    });

    match result {
        Ok(Ok(())) => {
            if let Some(Extension(telemetry)) = telemetry {
                telemetry.invitation_draft_created(
                    &organization_id,
                    &invitation_id,
                    &body.role,
                    now,
                );
            }
            (
                StatusCode::CREATED,
                Json(json!({
                    "invitation_id": invitation_id,
                    "organization_id": organization_id,
                    "email": email,
                    "role": body.role,
                    "expires_at": expires_at,
                    "delivery": {"status": "draft"}
                })),
            )
                .into_response()
        }
        Ok(Err("forbidden")) => error(StatusCode::FORBIDDEN, "organization admin required"),
        Ok(Err(_)) => error(StatusCode::BAD_REQUEST, "invalid organization invitation"),
        Err(_) => server_error(),
    }
}

async fn accept_invitation(
    Extension(user): Extension<AuthUser>,
    telemetry: Option<Extension<OrganizationTelemetry>>,
    Json(body): Json<AcceptInvitation>,
) -> Response {
    if body.token.trim().is_empty() {
        return error(StatusCode::BAD_REQUEST, "invalid invitation");
    }
    let token_hash = hash_token(body.token.trim());
    let now = chrono::Utc::now().timestamp();
    type AcceptanceResult =
        Result<Result<(String, String, String, i64), &'static str>, rusqlite::Error>;
    let result: AcceptanceResult = with_db(|conn| {
        let tx = conn.unchecked_transaction()?;
        let invitation = tx
            .query_row(
                "SELECT id, organization_id, email, role, expires_at, accepted_at, revoked_at
                 FROM organization_invitations WHERE token_hash = ?1",
                [&token_hash],
                |row| {
                    Ok((
                        row.get::<_, String>(0)?,
                        row.get::<_, String>(1)?,
                        row.get::<_, String>(2)?,
                        row.get::<_, String>(3)?,
                        row.get::<_, i64>(4)?,
                        row.get::<_, Option<i64>>(5)?,
                        row.get::<_, Option<i64>>(6)?,
                    ))
                },
            )
            .optional()?;
        let Some((id, organization_id, email, role, expires_at, accepted_at, revoked_at)) =
            invitation
        else {
            return Ok(Err("invalid"));
        };
        if accepted_at.is_some()
            || revoked_at.is_some()
            || expires_at <= now
            || normalize_email(&email) != normalize_email(&user.email)
        {
            return Ok(Err("invalid"));
        }
        let prior_status: Option<String> = tx
            .query_row(
                "SELECT status FROM organization_memberships
                 WHERE organization_id = ?1 AND user_id = ?2",
                rusqlite::params![organization_id, user.id],
                |row| row.get(0),
            )
            .optional()?;
        if prior_status.as_deref() == Some("revoked") {
            return Ok(Err("invalid"));
        }
        tx.execute(
            "INSERT INTO organization_memberships
             (organization_id, user_id, role, status, created_at, updated_at)
             VALUES (?1, ?2, ?3, 'active', ?4, ?4)
             ON CONFLICT(organization_id, user_id) DO UPDATE SET
               role = excluded.role, updated_at = excluded.updated_at",
            rusqlite::params![organization_id, user.id, role, now],
        )?;
        tx.execute(
            "UPDATE organization_invitations SET accepted_at = ?1 WHERE id = ?2",
            rusqlite::params![now, id],
        )?;
        bump_authorization_version(&tx, &user.id)?;
        let active_count: i64 = tx.query_row(
            "SELECT COUNT(*) FROM organization_memberships
             WHERE organization_id = ?1 AND status = 'active'",
            [&organization_id],
            |row| row.get(0),
        )?;
        tx.commit()?;
        Ok(Ok((id, organization_id, role, active_count)))
    });

    match result {
        Ok(Ok((invitation_id, organization_id, role, active_count))) => {
            if let Some(Extension(telemetry)) = telemetry {
                telemetry.invitation_accepted(&organization_id, &invitation_id, &role, now);
                telemetry.active_member_count_after_invitation(
                    &organization_id,
                    &invitation_id,
                    active_count,
                    now,
                );
            }
            (
                StatusCode::OK,
                Json(json!({"organization_id": organization_id, "role": role})),
            )
                .into_response()
        }
        Ok(Err(_)) => error(StatusCode::BAD_REQUEST, "invalid invitation"),
        Err(_) => server_error(),
    }
}

async fn revoke_invitation(
    Extension(user): Extension<AuthUser>,
    Path((organization_id, invitation_id)): Path<(String, String)>,
) -> Response {
    let result: Result<Result<(), &str>, rusqlite::Error> = with_db(|conn| {
        let membership = active_membership(conn, &user.id, &organization_id)?;
        if !membership.as_ref().is_some_and(|m| can_administer(&m.role)) {
            return Ok(Err("forbidden"));
        }
        let changed = conn.execute(
            "UPDATE organization_invitations SET revoked_at = ?1
             WHERE id = ?2 AND organization_id = ?3 AND accepted_at IS NULL AND revoked_at IS NULL",
            rusqlite::params![
                chrono::Utc::now().timestamp(),
                invitation_id,
                organization_id
            ],
        )?;
        Ok(if changed == 0 {
            Err("not_found")
        } else {
            Ok(())
        })
    });
    match result {
        Ok(Ok(())) => Json(json!({"ok": true})).into_response(),
        Ok(Err("forbidden")) => error(StatusCode::FORBIDDEN, "organization admin required"),
        Ok(Err("not_found")) => error(StatusCode::NOT_FOUND, "invitation not found"),
        _ => server_error(),
    }
}

async fn revoke_member(
    Extension(user): Extension<AuthUser>,
    telemetry: Option<Extension<OrganizationTelemetry>>,
    Path((organization_id, member_id)): Path<(String, String)>,
) -> Response {
    let now = chrono::Utc::now().timestamp();
    let result: Result<Result<(), &str>, rusqlite::Error> = with_db(|conn| {
        let actor = active_membership(conn, &user.id, &organization_id)?;
        if !actor.as_ref().is_some_and(|m| can_administer(&m.role)) {
            return Ok(Err("forbidden"));
        }
        let target = active_membership(conn, &member_id, &organization_id)?;
        let Some(target) = target else {
            return Ok(Err("not_found"));
        };
        if target.role == "owner" || member_id == user.id {
            return Ok(Err("owner"));
        }
        conn.execute(
            "UPDATE organization_memberships SET status = 'revoked', updated_at = ?1
             WHERE organization_id = ?2 AND user_id = ?3 AND status = 'active'",
            rusqlite::params![now, organization_id, member_id],
        )?;
        move_primary_organization_if_needed(conn, &member_id, &organization_id)?;
        bump_authorization_version(conn, &member_id)?;
        Ok(Ok(()))
    });
    match result {
        Ok(Ok(())) => {
            if let Some(Extension(telemetry)) = telemetry {
                telemetry.active_member_count_after_membership_revoked(&organization_id, now);
            }
            Json(json!({"ok": true})).into_response()
        }
        Ok(Err("forbidden")) => error(StatusCode::FORBIDDEN, "organization admin required"),
        Ok(Err("owner")) => error(
            StatusCode::CONFLICT,
            "owner must transfer administration first",
        ),
        Ok(Err("not_found")) => error(StatusCode::NOT_FOUND, "active member not found"),
        _ => server_error(),
    }
}

async fn leave_organization(
    Extension(user): Extension<AuthUser>,
    telemetry: Option<Extension<OrganizationTelemetry>>,
    Path(organization_id): Path<String>,
) -> Response {
    let now = chrono::Utc::now().timestamp();
    let result: Result<Result<(), &str>, rusqlite::Error> = with_db(|conn| {
        let membership = active_membership(conn, &user.id, &organization_id)?;
        let Some(membership) = membership else {
            return Ok(Err("not_found"));
        };
        if membership.role == "owner" {
            return Ok(Err("owner"));
        }
        let primary: String = conn.query_row(
            "SELECT tenant_id FROM users WHERE id = ?1",
            [&user.id],
            |row| row.get(0),
        )?;
        if primary == organization_id {
            let other_memberships: i64 = conn.query_row(
                "SELECT COUNT(*) FROM organization_memberships
                 WHERE user_id = ?1 AND status = 'active' AND organization_id <> ?2",
                rusqlite::params![user.id, organization_id],
                |row| row.get(0),
            )?;
            if other_memberships == 0 {
                return Ok(Err("last"));
            }
        }
        conn.execute(
            "UPDATE organization_memberships SET status = 'revoked', updated_at = ?1
             WHERE organization_id = ?2 AND user_id = ?3 AND status = 'active'",
            rusqlite::params![now, organization_id, user.id],
        )?;
        move_primary_organization_if_needed(conn, &user.id, &organization_id)?;
        bump_authorization_version(conn, &user.id)?;
        Ok(Ok(()))
    });
    match result {
        Ok(Ok(())) => {
            if let Some(Extension(telemetry)) = telemetry {
                telemetry.active_member_count_after_membership_left(&organization_id, now);
            }
            Json(json!({"ok": true})).into_response()
        }
        Ok(Err("not_found")) => error(StatusCode::NOT_FOUND, "active membership not found"),
        Ok(Err("last")) => error(StatusCode::CONFLICT, "cannot leave the only organization"),
        _ => server_error(),
    }
}

async fn transfer_admin(
    Extension(user): Extension<AuthUser>,
    Path(organization_id): Path<String>,
    Json(body): Json<TransferAdmin>,
) -> Response {
    let result: Result<Result<(), &str>, rusqlite::Error> = with_db(|conn| {
        let actor = active_membership(conn, &user.id, &organization_id)?;
        if !actor.as_ref().is_some_and(|m| m.role == "owner") {
            return Ok(Err("forbidden"));
        }
        let target = active_membership(conn, &body.user_id, &organization_id)?;
        if target.is_none() {
            return Ok(Err("not_found"));
        }
        if body.user_id == user.id {
            return Ok(Err("same_user"));
        }
        let now = chrono::Utc::now().timestamp();
        conn.execute(
            "UPDATE organization_memberships SET role = 'admin', updated_at = ?1
             WHERE organization_id = ?2 AND user_id = ?3",
            rusqlite::params![now, organization_id, user.id],
        )?;
        conn.execute(
            "UPDATE organization_memberships SET role = 'owner', updated_at = ?1
             WHERE organization_id = ?2 AND user_id = ?3",
            rusqlite::params![now, organization_id, body.user_id],
        )?;
        bump_authorization_version(conn, &user.id)?;
        bump_authorization_version(conn, &body.user_id)?;
        Ok(Ok(()))
    });
    match result {
        Ok(Ok(())) => Json(json!({"ok": true, "owner_id": body.user_id})).into_response(),
        Ok(Err("forbidden")) => error(StatusCode::FORBIDDEN, "organization owner required"),
        Ok(Err("same_user")) => error(
            StatusCode::BAD_REQUEST,
            "new owner must be another active member",
        ),
        Ok(Err("not_found")) => error(StatusCode::NOT_FOUND, "active member not found"),
        _ => server_error(),
    }
}

async fn list_workspaces(
    Extension(user): Extension<AuthUser>,
    Path(organization_id): Path<String>,
) -> Response {
    let result: Result<Option<Vec<Value>>, rusqlite::Error> = with_db(|conn| {
        if active_membership(conn, &user.id, &organization_id)?.is_none() {
            return Ok(None);
        }
        let mut stmt = conn.prepare(
            "SELECT id, workspace_id, installation_id, app_id, created_at
             FROM organization_workspaces WHERE organization_id = ?1 ORDER BY created_at, id",
        )?;
        let rows = stmt.query_map([&organization_id], |row| {
            Ok(json!({
                "link_id": row.get::<_, String>(0)?,
                "workspace_id": row.get::<_, String>(1)?,
                "installation_id": row.get::<_, Option<String>>(2)?,
                "app_id": row.get::<_, Option<String>>(3)?,
                "created_at": row.get::<_, i64>(4)?,
            }))
        })?;
        Ok(Some(rows.collect::<Result<Vec<_>, _>>()?))
    });
    match result {
        Ok(Some(workspaces)) => {
            Json(json!({"organization_id": organization_id, "workspaces": workspaces}))
                .into_response()
        }
        Ok(None) => error(StatusCode::NOT_FOUND, "organization not found"),
        Err(_) => server_error(),
    }
}

async fn link_workspace(
    Extension(user): Extension<AuthUser>,
    Path(organization_id): Path<String>,
    Json(body): Json<LinkWorkspace>,
) -> Response {
    if body.workspace_id.trim().is_empty() {
        return error(StatusCode::BAD_REQUEST, "workspace_id is required");
    }
    let link_id = Uuid::new_v4().to_string();
    let result: Result<Result<(), &str>, rusqlite::Error> = with_db(|conn| {
        let membership = active_membership(conn, &user.id, &organization_id)?;
        if !membership.as_ref().is_some_and(|m| can_administer(&m.role)) {
            return Ok(Err("forbidden"));
        }
        conn.execute(
            "INSERT INTO organization_workspaces
             (id, organization_id, workspace_id, installation_id, app_id, created_at)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
            rusqlite::params![
                link_id,
                organization_id,
                body.workspace_id,
                body.installation_id,
                body.app_id,
                chrono::Utc::now().timestamp()
            ],
        )?;
        Ok(Ok(()))
    });
    match result {
        Ok(Ok(())) => (
            StatusCode::CREATED,
            Json(json!({"link_id": link_id, "organization_id": organization_id, "workspace_id": body.workspace_id})),
        )
            .into_response(),
        Ok(Err("forbidden")) => error(StatusCode::FORBIDDEN, "organization admin required"),
        Err(e) if e.to_string().contains("UNIQUE") => {
            error(StatusCode::CONFLICT, "workspace is already linked")
        }
        _ => server_error(),
    }
}

async fn unlink_workspace(
    Extension(user): Extension<AuthUser>,
    Path((organization_id, workspace_id)): Path<(String, String)>,
) -> Response {
    let result: Result<Result<(), &str>, rusqlite::Error> = with_db(|conn| {
        let membership = active_membership(conn, &user.id, &organization_id)?;
        if !membership.as_ref().is_some_and(|m| can_administer(&m.role)) {
            return Ok(Err("forbidden"));
        }
        let changed = conn.execute(
            "DELETE FROM organization_workspaces WHERE organization_id = ?1 AND workspace_id = ?2",
            rusqlite::params![organization_id, workspace_id],
        )?;
        Ok(if changed == 0 {
            Err("not_found")
        } else {
            Ok(())
        })
    });
    match result {
        Ok(Ok(())) => Json(json!({"ok": true})).into_response(),
        Ok(Err("forbidden")) => error(StatusCode::FORBIDDEN, "organization admin required"),
        Ok(Err("not_found")) => error(StatusCode::NOT_FOUND, "workspace link not found"),
        _ => server_error(),
    }
}

fn hash_token(token: &str) -> String {
    hex::encode(Sha256::digest(token.as_bytes()))
}

fn error(status: StatusCode, message: &str) -> Response {
    (status, Json(json!({"error": message}))).into_response()
}

fn server_error() -> Response {
    error(StatusCode::INTERNAL_SERVER_ERROR, "internal server error")
}