fraiseql-auth 2.19.1

//! Postgres-backed [`OtpStore`] (#367).
//!
//! `InMemoryOtpStore` keeps codes and the per-email send budget in process
//! memory. Behind more than one replica that is a **silent security downgrade**,
//! not just an inconvenience: the send rate limit and the 3-attempt verify cap
//! are both per-process, so N replicas multiply both budgets by N and a
//! six-digit code becomes brute-forceable. This store makes the budgets shared,
//! which is what lets `[auth.local] otp = true` be served by a real deployment.
//!
//! At rest the code and its link token are stored as SHA-256 hashes — a database
//! read cannot replay a live magic link. The code comparison stays constant-time on
//! the hash; a token is found by its hash through a unique index.

use async_trait::async_trait;
use sha2::{Digest, Sha256};
use sqlx::{PgPool, Row};
use subtle::ConstantTimeEq as _;

use super::{
    MAX_VERIFY_ATTEMPTS, OTP_RATE_MAX, OTP_RATE_WINDOW_SECS, OTP_TTL_SECS, OtpIssue, OtpStore,
    link_token_hash,
};
use crate::{
    error::{AuthError, Result},
    session::unix_now,
};

/// DDL for the `OTP` tables. Executed at boot and by the live-PG suite.
pub const PG_OTP_SCHEMA_SQL: &str = concat!(
    claim_auth_schema!(),
    r"
CREATE TABLE IF NOT EXISTS auth.tb_otp_code (
    email      TEXT PRIMARY KEY,
    code_hash  BYTEA NOT NULL,
    expires_at BIGINT NOT NULL,
    attempts   INTEGER NOT NULL DEFAULT 0
);
CREATE TABLE IF NOT EXISTS auth.tb_otp_send_budget (
    email        TEXT PRIMARY KEY,
    count        INTEGER NOT NULL DEFAULT 0,
    window_start BIGINT NOT NULL
);
ALTER TABLE auth.tb_otp_send_budget ADD COLUMN IF NOT EXISTS guesses INTEGER NOT NULL DEFAULT 0;
ALTER TABLE auth.tb_otp_code ADD COLUMN IF NOT EXISTS link_token_hash BYTEA;
CREATE UNIQUE INDEX IF NOT EXISTS uq_otp_code_link_token_hash
    ON auth.tb_otp_code (link_token_hash) WHERE link_token_hash IS NOT NULL;
REVOKE ALL ON auth.tb_otp_code FROM PUBLIC;
REVOKE ALL ON auth.tb_otp_send_budget FROM PUBLIC;
"
);

fn code_hash(code: &str) -> Vec<u8> {
    Sha256::digest(code.as_bytes()).to_vec()
}

fn db_err(e: sqlx::Error) -> AuthError {
    AuthError::DatabaseError {
        message: e.to_string(),
    }
}

/// Postgres-backed `OTP` store.
#[derive(Debug, Clone)]
pub struct PgOtpStore {
    db: PgPool,
}

impl PgOtpStore {
    /// Create a store over an existing pool.
    #[must_use]
    pub const fn new(db: PgPool) -> Self {
        Self { db }
    }

    /// Execute the table DDL. Idempotent.
    ///
    /// # Errors
    ///
    /// Returns [`AuthError::DatabaseError`] if a statement fails.
    pub async fn init(&self) -> Result<()> {
        // raw_sql: a multi-statement batch whose schema claim is a `DO` block, which a
        // split on `;` would cut apart.
        sqlx::raw_sql(PG_OTP_SCHEMA_SQL).execute(&self.db).await.map_err(db_err)?;
        Ok(())
    }
}

// Reason: async_trait required for dyn-compatibility; remove when RTN + Send is stable
#[async_trait]
impl OtpStore for PgOtpStore {
    #[allow(clippy::cast_possible_wrap)] // Reason: expiry is only ever written from unix_now()
    async fn sweep_expired(&self) -> Result<u64> {
        let now = unix_now()? as i64;
        let codes = sqlx::query("DELETE FROM auth.tb_otp_code WHERE expires_at <= $1")
            .bind(now)
            .execute(&self.db)
            .await
            .map_err(db_err)?
            .rows_affected();

        // See the in-memory twin: a closed window is dead weight, and this table is the
        // one that otherwise keeps a row for every address that ever asked for a code.
        let window = i64::try_from(crate::otp::OTP_RATE_WINDOW_SECS).unwrap_or(i64::MAX);
        let budgets =
            sqlx::query("DELETE FROM auth.tb_otp_send_budget WHERE window_start + $1 <= $2")
                .bind(window)
                .bind(now)
                .execute(&self.db)
                .await
                .map_err(db_err)?
                .rows_affected();

        Ok(codes + budgets)
    }

    #[allow(clippy::cast_possible_wrap, clippy::cast_sign_loss)] // Reason: unix seconds and small counters fit i64/i32
    async fn create_otp(&self, email: &str) -> Result<OtpIssue> {
        let now = unix_now()?;

        // Reserve one send atomically: the window reset and the increment happen
        // in a single UPSERT, so two concurrent requests cannot both read a
        // stale count and each grant themselves a send. A new window also returns
        // the address's guesses; a send inside one does not.
        let reserved = sqlx::query(
            "INSERT INTO auth.tb_otp_send_budget (email, count, window_start)
             VALUES ($1, 1, $2)
             ON CONFLICT (email) DO UPDATE SET
                 count = CASE
                     WHEN $2 - auth.tb_otp_send_budget.window_start >= $3 THEN 1
                     ELSE auth.tb_otp_send_budget.count + 1 END,
                 guesses = CASE
                     WHEN $2 - auth.tb_otp_send_budget.window_start >= $3 THEN 0
                     ELSE auth.tb_otp_send_budget.guesses END,
                 window_start = CASE
                     WHEN $2 - auth.tb_otp_send_budget.window_start >= $3 THEN $2
                     ELSE auth.tb_otp_send_budget.window_start END
             RETURNING count, window_start",
        )
        .bind(email)
        .bind(now as i64)
        .bind(OTP_RATE_WINDOW_SECS as i64)
        .fetch_one(&self.db)
        .await
        .map_err(db_err)?;

        let count: i32 = reserved.get("count");
        let window_start: i64 = reserved.get("window_start");
        if count as u32 > OTP_RATE_MAX {
            return Err(AuthError::RateLimited {
                retry_after_secs: (window_start as u64 + OTP_RATE_WINDOW_SECS).saturating_sub(now),
            });
        }

        // One row holds the code and its link token, so a resend replaces both.
        let issue = OtpIssue::generate();
        sqlx::query(
            "INSERT INTO auth.tb_otp_code (email, code_hash, link_token_hash, expires_at, attempts)
             VALUES ($1, $2, $3, $4, 0)
             ON CONFLICT (email) DO UPDATE SET
                 code_hash = EXCLUDED.code_hash,
                 link_token_hash = EXCLUDED.link_token_hash,
                 expires_at = EXCLUDED.expires_at,
                 attempts = 0",
        )
        .bind(email)
        .bind(code_hash(&issue.code))
        .bind(link_token_hash(&issue.link_token))
        .bind((now + OTP_TTL_SECS) as i64)
        .execute(&self.db)
        .await
        .map_err(db_err)?;

        Ok(issue)
    }

    #[allow(clippy::cast_sign_loss, clippy::cast_possible_wrap)] // Reason: expiry/counters are written from unix_now() and small constants
    async fn verify_otp(&self, email: &str, code: &str) -> Result<()> {
        let now = unix_now()?;

        // Charge the attempt first and read the post-increment count, so a
        // concurrent flood cannot each read `attempts = 0` and get a free guess.
        let row = sqlx::query(
            "UPDATE auth.tb_otp_code SET attempts = attempts + 1 WHERE email = $1
             RETURNING code_hash, expires_at, attempts",
        )
        .bind(email)
        .fetch_optional(&self.db)
        .await
        .map_err(db_err)?
        .ok_or_else(|| AuthError::InvalidToken {
            reason: "no pending OTP for email".into(),
        })?;

        // Then charge the address. The code's own count restarts with every
        // send; this one restarts only with the send window, so a resend issues a new
        // code without issuing new guesses. Charged after the code lookup, so probing
        // an address with no pending code leaves no row behind.
        let guesses: i32 = sqlx::query(
            "INSERT INTO auth.tb_otp_send_budget (email, count, window_start, guesses)
             VALUES ($1, 0, $2, 1)
             ON CONFLICT (email) DO UPDATE SET
                 count = CASE
                     WHEN $2 - auth.tb_otp_send_budget.window_start >= $3 THEN 0
                     ELSE auth.tb_otp_send_budget.count END,
                 guesses = CASE
                     WHEN $2 - auth.tb_otp_send_budget.window_start >= $3 THEN 1
                     ELSE auth.tb_otp_send_budget.guesses + 1 END,
                 window_start = CASE
                     WHEN $2 - auth.tb_otp_send_budget.window_start >= $3 THEN $2
                     ELSE auth.tb_otp_send_budget.window_start END
             RETURNING guesses",
        )
        .bind(email)
        .bind(now as i64)
        .bind(OTP_RATE_WINDOW_SECS as i64)
        .fetch_one(&self.db)
        .await
        .map_err(db_err)?
        .get("guesses");

        let stored_hash: Vec<u8> = row.get("code_hash");
        let expires_at: i64 = row.get("expires_at");
        let attempts: i32 = row.get("attempts");

        // The DELETE is what makes the code single-use, so its failure must fail
        // the verification (#984): returning Ok after a failed consume reports a
        // guarantee that does not hold and leaves the code replayable until it
        // expires. Propagated on every path — a consume that silently does
        // nothing is the same defect wherever it sits.
        let consume = || async {
            sqlx::query("DELETE FROM auth.tb_otp_code WHERE email = $1")
                .bind(email)
                .execute(&self.db)
                .await
                .map_err(db_err)
                .map(|_| ())
        };

        if now >= expires_at as u64 {
            consume().await?;
            return Err(AuthError::InvalidToken {
                reason: "OTP has expired".into(),
            });
        }
        if attempts > MAX_VERIFY_ATTEMPTS as i32 || guesses > MAX_VERIFY_ATTEMPTS as i32 {
            consume().await?;
            return Err(AuthError::RateLimited {
                retry_after_secs: OTP_RATE_WINDOW_SECS,
            });
        }

        // Constant-time comparison on the hashes (#788): a data-dependent
        // comparison leaks how many leading bytes matched via its return time.
        let presented = code_hash(code);
        if !bool::from(presented.ct_eq(&stored_hash)) {
            return Err(AuthError::InvalidToken {
                reason: "invalid OTP code".into(),
            });
        }

        // Correct — consume it (single-use).
        consume().await?;
        Ok(())
    }

    #[allow(clippy::cast_sign_loss)] // Reason: expiry is only ever written from unix_now()
    async fn redeem_link(&self, link_token: &str) -> Result<String> {
        let now = unix_now()?;
        // The DELETE is the redemption: of two concurrent requests only one gets the row,
        // and the code mailed with the token goes with it.
        let row = sqlx::query(
            "DELETE FROM auth.tb_otp_code WHERE link_token_hash = $1
             RETURNING email, expires_at",
        )
        .bind(link_token_hash(link_token))
        .fetch_optional(&self.db)
        .await
        .map_err(db_err)?
        .ok_or_else(|| AuthError::InvalidToken {
            reason: "no pending link for token".into(),
        })?;
        let expires_at: i64 = row.get("expires_at");
        if now >= expires_at as u64 {
            return Err(AuthError::InvalidToken {
                reason: "link has expired".into(),
            });
        }
        Ok(row.get("email"))
    }
}