clawmates/crates/cm-auth/src/service.rs

use argon2::password_hash::rand_core::OsRng;
use argon2::password_hash::{PasswordHash, PasswordHasher, PasswordVerifier, SaltString};
use argon2::Argon2;
use cm_domain::{Role, UserId, WorkspaceId};
use sqlx::{PgPool, Row};
use time::{Duration, OffsetDateTime};

use crate::token::{hash_token, SessionToken};

/// How long a login session stays valid.
pub const SESSION_TTL: Duration = Duration::days(7);

/// The authenticated caller attached to every API request: everything RBAC
/// decisions need, nothing more.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AuthedUser {
    pub user_id: UserId,
    pub workspace_id: WorkspaceId,
    pub role: Role,
}

#[derive(Debug, thiserror::Error)]
pub enum AuthError {
    /// Wrong password and unknown email are deliberately the same error so
    /// login cannot be used to enumerate accounts.
    #[error("invalid credentials")]
    InvalidCredentials,
    #[error("unauthenticated")]
    Unauthenticated,
    #[error("password hashing failed: {0}")]
    Hashing(String),
    #[error(transparent)]
    Db(#[from] sqlx::Error),
}

#[derive(Clone)]
pub struct AuthService {
    pool: PgPool,
    verifier: Option<std::sync::Arc<crate::JwtVerifier>>,
    /// SaaS mode: a brand-new external identity provisions its own workspace
    /// (owner) rather than joining the instance's single workspace.
    per_signup_workspace: bool,
}

impl AuthService {
    pub fn new(pool: PgPool) -> AuthService {
        AuthService {
            pool,
            verifier: None,
            per_signup_workspace: false,
        }
    }

    /// Enables hosted-identity session JWTs (Clerk / OIDC) alongside
    /// local sessions.
    pub fn with_verifier(mut self, verifier: std::sync::Arc<crate::JwtVerifier>) -> AuthService {
        self.verifier = Some(verifier);
        self
    }

    /// In SaaS deployments, give each new hosted-identity sign-in its own
    /// workspace (they become its owner) instead of joining the first one.
    pub fn with_per_signup_workspace(mut self, enabled: bool) -> AuthService {
        self.per_signup_workspace = enabled;
        self
    }

    /// Verifies an issuer session token and JIT-provisions the user on
    /// first sight. Role tracks the issuer claim on every login (Clerk's
    /// org role is the source of truth for SSO users).
    async fn authenticate_external(
        &self,
        token: &str,
        verifier: &crate::JwtVerifier,
    ) -> Result<AuthedUser, AuthError> {
        let claims = verifier
            .verify(token)
            .await
            .map_err(|_| AuthError::Unauthenticated)?;
        let role = match claims.role.as_deref() {
            Some(role) if role.contains("admin") => Role::Owner,
            _ => Role::Member,
        };
        let role_str = if role == Role::Owner {
            "owner"
        } else {
            "member"
        };

        if let Some(row) = sqlx::query!(
            "UPDATE users SET role = $2 WHERE auth_subject = $1
             RETURNING id, workspace_id",
            claims.sub,
            role_str,
        )
        .fetch_optional(&self.pool)
        .await?
        {
            return Ok(AuthedUser {
                user_id: UserId::from(row.id),
                workspace_id: WorkspaceId::from(row.workspace_id),
                role,
            });
        }

        // Link to an existing local account by email, if one exists.
        if let Some(email) = claims.email.as_deref() {
            if let Some(row) = sqlx::query!(
                "UPDATE users SET auth_subject = $2, role = $3 WHERE email = $1
                 RETURNING id, workspace_id",
                email,
                claims.sub,
                role_str,
            )
            .fetch_optional(&self.pool)
            .await?
            {
                return Ok(AuthedUser {
                    user_id: UserId::from(row.id),
                    workspace_id: WorkspaceId::from(row.workspace_id),
                    role,
                });
            }
        }

        // First sight. Serialize concurrent first-logins for the same subject
        // (the authed shell fires several API calls at once) with a per-subject
        // advisory lock, so SaaS mode can't create duplicate workspaces and the
        // appliance path can't double-insert the user.
        let email = claims
            .email
            .clone()
            .unwrap_or_else(|| format!("{}@sso.local", claims.sub));
        let display_name = email.split('@').next().unwrap_or("teammate").to_owned();

        let mut tx = self.pool.begin().await?;
        sqlx::query("SELECT pg_advisory_xact_lock(hashtext($1))")
            .bind(&claims.sub)
            .execute(&mut *tx)
            .await?;

        // A concurrent winner may have already provisioned this identity.
        if let Some(row) =
            sqlx::query("SELECT id, workspace_id, role FROM users WHERE auth_subject = $1")
                .bind(&claims.sub)
                .fetch_optional(&mut *tx)
                .await?
        {
            tx.commit().await?;
            let role_db: String = row.get("role");
            return Ok(AuthedUser {
                user_id: UserId::from(row.get::<uuid::Uuid, _>("id")),
                workspace_id: WorkspaceId::from(row.get::<uuid::Uuid, _>("workspace_id")),
                role: if role_db == "owner" {
                    Role::Owner
                } else {
                    Role::Member
                },
            });
        }

        // Pick the workspace: SaaS provisions a fresh one this user owns;
        // appliance joins the instance's single workspace.
        let (workspace, effective_role): (uuid::Uuid, &str) = if self.per_signup_workspace {
            let ws = WorkspaceId::new();
            sqlx::query("INSERT INTO workspaces (id, name, plan) VALUES ($1, $2, 'free')")
                .bind(ws.as_uuid())
                .bind(format!("{display_name}'s workspace"))
                .execute(&mut *tx)
                .await?;
            (ws.as_uuid(), "owner")
        } else {
            let ws: uuid::Uuid =
                sqlx::query_scalar("SELECT id FROM workspaces ORDER BY created_at, id LIMIT 1")
                    .fetch_optional(&mut *tx)
                    .await?
                    .ok_or(AuthError::Unauthenticated)?;
            (ws, role_str)
        };
        let user_id = UserId::new();
        let row = sqlx::query(
            "INSERT INTO users (id, workspace_id, email, role, display_name, auth_subject)
             VALUES ($1, $2, $3, $4, $5, $6)
             RETURNING id, workspace_id",
        )
        .bind(user_id.as_uuid())
        .bind(workspace)
        .bind(&email)
        .bind(effective_role)
        .bind(&display_name)
        .bind(&claims.sub)
        .fetch_one(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(AuthedUser {
            user_id: UserId::from(row.get::<uuid::Uuid, _>("id")),
            workspace_id: WorkspaceId::from(row.get::<uuid::Uuid, _>("workspace_id")),
            role: if effective_role == "owner" {
                Role::Owner
            } else {
                Role::Member
            },
        })
    }

    /// Sets (or resets) a user's local password.
    pub async fn set_password(&self, user_id: UserId, password: &str) -> Result<(), AuthError> {
        let salt = SaltString::generate(&mut OsRng);
        let hash = Argon2::default()
            .hash_password(password.as_bytes(), &salt)
            .map_err(|e| AuthError::Hashing(e.to_string()))?
            .to_string();
        sqlx::query!(
            "UPDATE users SET password_hash = $2 WHERE id = $1",
            user_id.as_uuid(),
            hash,
        )
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    /// Verifies email + password and opens a new session.
    pub async fn login_local(
        &self,
        email: &str,
        password: &str,
    ) -> Result<SessionToken, AuthError> {
        let row = sqlx::query!(
            "SELECT id, password_hash FROM users WHERE email = $1",
            email,
        )
        .fetch_optional(&self.pool)
        .await?;
        let (user_id, stored_hash) = match row {
            Some(row) => match row.password_hash {
                Some(hash) => (row.id, hash),
                None => return Err(AuthError::InvalidCredentials),
            },
            None => return Err(AuthError::InvalidCredentials),
        };

        let parsed =
            PasswordHash::new(&stored_hash).map_err(|e| AuthError::Hashing(e.to_string()))?;
        Argon2::default()
            .verify_password(password.as_bytes(), &parsed)
            .map_err(|_| AuthError::InvalidCredentials)?;

        let token = SessionToken::generate();
        sqlx::query!(
            "INSERT INTO auth_sessions (token_hash, user_id, expires_at)
             VALUES ($1, $2, $3)",
            hash_token(token.secret()),
            user_id,
            OffsetDateTime::now_utc() + SESSION_TTL,
        )
        .execute(&self.pool)
        .await?;
        Ok(token)
    }

    /// Resolves a bearer token to the authenticated user. Issuer session
    /// JWTs (three dot-separated segments) take the hosted-identity path;
    /// everything else is a local opaque session token.
    pub async fn authenticate(&self, token_secret: &str) -> Result<AuthedUser, AuthError> {
        if let Some(verifier) = self.verifier.clone() {
            if token_secret.matches('.').count() == 2 {
                return self.authenticate_external(token_secret, &verifier).await;
            }
        }
        let row = sqlx::query!(
            "SELECT u.id, u.workspace_id, u.role
             FROM auth_sessions s
             JOIN users u ON u.id = s.user_id
             WHERE s.token_hash = $1 AND s.expires_at > now()",
            hash_token(token_secret),
        )
        .fetch_optional(&self.pool)
        .await?
        .ok_or(AuthError::Unauthenticated)?;
        Ok(AuthedUser {
            user_id: UserId::from(row.id),
            workspace_id: WorkspaceId::from(row.workspace_id),
            role: if row.role == "owner" {
                Role::Owner
            } else {
                Role::Member
            },
        })
    }

    /// Mint a long-lived opaque session for an internal service caller
    /// (e.g. the per-team ZeroClaw runtime calling back into the MCP door).
    /// Returns the plaintext token — the caller is responsible for handing
    /// it to the exact process that needs it and not persisting it broadly.
    pub async fn mint_service_session(
        &self,
        user_id: UserId,
        ttl: Duration,
    ) -> Result<SessionToken, AuthError> {
        let token = SessionToken::generate();
        sqlx::query!(
            "INSERT INTO auth_sessions (token_hash, user_id, expires_at)
             VALUES ($1, $2, $3)",
            hash_token(token.secret()),
            user_id.as_uuid(),
            OffsetDateTime::now_utc() + ttl,
        )
        .execute(&self.pool)
        .await?;
        Ok(token)
    }

    /// Ends the session for this token.
    pub async fn logout(&self, token_secret: &str) -> Result<(), AuthError> {
        sqlx::query!(
            "DELETE FROM auth_sessions WHERE token_hash = $1",
            hash_token(token_secret),
        )
        .execute(&self.pool)
        .await?;
        Ok(())
    }
}