maarcadetweet/crates/pds-server/src/routes/helpers.rs

//! Shared helpers for the PDS route handlers.
//!
//! These are used by both `repo.rs` (mutable repo operations) and `sync.rs`
//! (read-only sync endpoints). They handle the boilerplate of:
//!
//! * Loading every block belonging to a user from the `repo_blocks` table
//!   into an in-memory [`MemoryBlockstore`].
//! * Loading the user's secp256k1 signing key from `users.signing_key`.
//! * Detecting the all-zero placeholder we use for a fresh account that has
//!   no commits yet.
//! * Serialising a write path under a Postgres row lock so concurrent
//!   writers for the same DID can't trample each other's MST updates
//!   (Phase 5b review C1).

use crate::routes::types::ErrorBody;
use crate::state::AppState;
use at_crypto::cid::cid_from_multihash_bytes;
use at_repo::blockstore::{Blockstore, MemoryBlockstore};
use at_repo::repo::Repo;
use axum::http::StatusCode;
use axum::Json;
use bytes::Bytes;
use cid::Cid;
use k256::ecdsa::SigningKey;
use k256::SecretKey;
use sqlx::Postgres;
use std::sync::Arc;

/// Load every block belonging to `did` from the `repo_blocks` table into a
/// fresh in-memory blockstore. Used to reconstruct a [`crate::at_repo::Repo`]
/// for either mutation or read-only inspection.
pub async fn load_user_blockstore(
    state: &AppState,
    did: &str,
) -> Result<Arc<MemoryBlockstore>, (StatusCode, Json<ErrorBody>)> {
    let bs = MemoryBlockstore::new();
    let rows: Vec<(Vec<u8>, Vec<u8>)> = sqlx::query_as(
        "SELECT cid, block FROM repo_blocks WHERE did = $1",
    )
    .bind(did)
    .fetch_all(&state.db)
    .await
    .map_err(|e| {
        err(
            StatusCode::INTERNAL_SERVER_ERROR,
            "InternalServerError",
            format!("repo_blocks load: {e}"),
        )
    })?;
    for (cid_bytes, block) in rows {
        let cid = cid_from_multihash_bytes(&cid_bytes).map_err(|e| {
            err(
                StatusCode::INTERNAL_SERVER_ERROR,
                "InternalServerError",
                format!("invalid cid in repo_blocks: {e}"),
            )
        })?;
        bs.put(&cid, Bytes::from(block))
            .await
            .map_err(|e| {
                err(
                    StatusCode::INTERNAL_SERVER_ERROR,
                    "InternalServerError",
                    format!("blockstore put: {e}"),
                )
            })?;
    }
    Ok(Arc::new(bs))
}

/// Hex sentinel stored in `head_cid` / `head_commit` for a fresh account
/// (no commits yet).
pub fn is_zero_blob(b: &[u8]) -> bool {
    !b.is_empty() && b.iter().all(|x| *x == 0)
}

/// Construct the user's `SigningKey` from `users.signing_key` (raw k256
/// secret-bytes).
pub fn load_signing_key(
    bytes: &[u8],
) -> Result<SigningKey, (StatusCode, Json<ErrorBody>)> {
    let secret = SecretKey::from_slice(bytes).map_err(|e| {
        err(
            StatusCode::INTERNAL_SERVER_ERROR,
            "InternalServerError",
            format!("invalid signing key bytes: {e}"),
        )
    })?;
    Ok(SigningKey::from(secret))
}

/// Load the head commit CID + signed commit block for `did`. Returns
/// `Ok(None)` if the account has no commits yet.
pub async fn load_head_commit(
    state: &AppState,
    did: &str,
) -> Result<Option<(Cid, Vec<u8>)>, (StatusCode, Json<ErrorBody>)> {
    let row: Option<(Vec<u8>, Vec<u8>)> = sqlx::query_as(
        "SELECT head_cid, head_commit FROM repos WHERE did = $1",
    )
    .bind(did)
    .fetch_optional(&state.db)
    .await
    .map_err(|e| {
        err(
            StatusCode::INTERNAL_SERVER_ERROR,
            "InternalServerError",
            format!("repos read: {e}"),
        )
    })?;
    let (head_cid_blob, head_commit_blob) = match row {
        Some(r) => r,
        None => return Ok(None),
    };
    if is_zero_blob(&head_cid_blob) || is_zero_blob(&head_commit_blob) {
        return Ok(None);
    }
    let cid = cid_from_multihash_bytes(&head_cid_blob).map_err(|e| {
        err(
            StatusCode::INTERNAL_SERVER_ERROR,
            "InternalServerError",
            format!("invalid head_cid bytes: {e}"),
        )
    })?;
    Ok(Some((cid, head_commit_blob)))
}

/// Construct an XRPC-shaped error tuple used by all the route handlers.
pub fn err(
    code: StatusCode,
    name: &str,
    msg: impl Into<String>,
) -> (StatusCode, Json<ErrorBody>) {
    (
        code,
        Json(ErrorBody::new(name, Some(msg.into()))),
    )
}

/// Convert an `anyhow::Error` (the error type returned by `at_repo`'s
/// repo methods) into a `sqlx::Error` so the closure handed to
/// [`apply_repo_write`] can return its outcome via `Result<_, sqlx::Error>`.
///
/// `anyhow::Error` doesn't implement `sqlx::DatabaseError`, so we
/// can't use `?` directly — the conversion wraps the original error
/// into `sqlx::Error::Decode` which preserves the source via
/// `Box<dyn std::error::Error + Send + Sync>`. `anyhow::Error` doesn't
/// implement `std::error::Error` itself, so we downcast its source
/// chain to a `String` (losing fidelity but never panicking on the
/// unknown source type).
pub fn to_sqlx_error(e: anyhow::Error) -> sqlx::Error {
    // Walk the anyhow chain and surface the first source that
    // implements StdError; fall back to a string wrapper.
    let dyn_err: Box<dyn std::error::Error + Send + Sync> =
        match e.downcast::<Box<dyn std::error::Error + Send + Sync>>() {
            Ok(boxed) => boxed,
            Err(other) => {
                let s = format!("{other:#}");
                Box::<dyn std::error::Error + Send + Sync>::from(s)
            }
        };
    sqlx::Error::Decode(dyn_err)
}

// -- repo write helper (Phase 5b review C1) ---------------------------------
//
// The previous code did:
//   1. SELECT head_commit FROM repos WHERE did = $1     -- non-locking read
//   2. Build an in-memory MST + apply the operation
//   3. INSERT blocks into repo_blocks
//   4. UPDATE repos SET head_cid = ...
//
// Two concurrent writers could both read the same head_commit, both build a
// valid child commit, and the second UPDATE would silently overwrite the
// first. The first writer's MST changes would survive in `repo_blocks` but
// become unreachable from `head_commit`, so a follow-up load+save on the
// repo would still see them — and then `Mst::put` would either no-op
// (because the rkey already exists) or branch off into a stale tree,
// depending on which blocks landed.
//
// The fix is to take a row-level write lock on `repos` for the duration of
// the in-memory mutation + commit + persist. Postgres `SELECT … FOR UPDATE`
// inside a transaction does exactly that: the lock is released when the
// transaction commits or rolls back, so concurrent writers serialise
// behind the holder rather than racing on the head_commit column.

/// Result of a successful repo write: the new signed commit, the CID
/// pointing at the freshly-written head block, and the new revision
/// string. Callers use the commit for AppView ingest pushes.
#[derive(Debug, Clone)]
pub struct RepoWriteOutcome {
    pub commit: at_repo::commit::Commit,
    pub head_cid_bytes: Vec<u8>,
    pub head_commit_bytes: Vec<u8>,
}

/// Apply a write to the user's repo under a row-level lock on the
/// `repos` row, then commit. Concurrent writers for the same DID block
/// behind the holder and proceed serially.
///
/// The flow:
/// 1. `BEGIN`
/// 2. `SELECT head_commit FROM repos WHERE did = $1 FOR UPDATE`
/// 3. Hydrate the `Repo` from `repo_blocks` + the locked head commit.
/// 4. Run the user's closure (`put_record`, `delete_record`, …) with
///    a mutable reference to the repo. The closure returns a
///    `RepoWriteOutcome` once it's finished mutating the repo and
///    called `Repo::commit`.
/// 5. Persist every block the closure (and `Repo::commit`) wrote into
///    `repo_blocks`.
/// 6. `UPDATE repos SET head_* = …` with the new commit.
/// 7. `COMMIT` — releases the lock and makes the new head visible to
///    other writers, who will now re-load from the new head instead of
///    racing on the old one.
///
/// The closure's returned `RepoWriteOutcome` is built *before* the
/// `UPDATE` (so the new commit's `signed_bytes` and CID are known when we
/// write the row), but the transaction stays open until after the
/// `UPDATE`. If the closure or `UPDATE` fails, the transaction rolls
/// back and no head pointer or block row changes are visible.
pub async fn apply_repo_write<F>(
    state: &AppState,
    did: &str,
    f: F,
) -> Result<RepoWriteOutcome, (StatusCode, Json<ErrorBody>)>
where
    F: for<'b> FnOnce(
        &'b mut Repo<MemoryBlockstore>,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = Result<RepoWriteOutcome, sqlx::Error>> + Send + 'b>,
    >,
{
    let mut tx = state.db.begin().await.map_err(|e| {
        err(
            StatusCode::INTERNAL_SERVER_ERROR,
            "InternalServerError",
            format!("begin tx: {e}"),
        )
    })?;

    // 2. Take the row-level write lock. Postgres parks competing
    //    transactions here until we COMMIT/ROLLBACK.
    let head_row: Option<(Vec<u8>, Vec<u8>, Option<Vec<u8>>)> = sqlx::query_as(
        "SELECT head_cid, head_commit, prev_commit
           FROM repos
          WHERE did = $1
          FOR UPDATE",
    )
    .bind(did)
    .fetch_optional(&mut *tx)
    .await
    .map_err(|e| {
        err(
            StatusCode::INTERNAL_SERVER_ERROR,
            "InternalServerError",
            format!("repos FOR UPDATE: {e}"),
        )
    })?;

    let (head_cid_blob, head_commit_blob) = match head_row {
        Some(r) => (r.0, r.1),
        None => {
            return Err((
                StatusCode::NOT_FOUND,
                Json(ErrorBody::new(
                    "RepoNotFound",
                    Some(format!("no repo row for {did}")),
                )),
            ));
        }
    };

    // 3. Hydrate the user's signing key + blockstore. These reads are
    //    not lock-sensitive — the signing key doesn't change, and the
    //    blockstore reads are append-only from our perspective.
    //
    //    We grab the signing key from outside the transaction (it's
    //    a separate table) to keep the FOR UPDATE window as short as
    //    practical — long-running locks contend with other writers.
    //
    //    Note: a brand-new account may have a `repos` row but no
    //    signing key in `users`; in that case `load_signing_key` from
    //    the connection pool is fine because the transaction's
    //    isolation level (Postgres default READ COMMITTED) lets the
    //    second query see the committed row.
    let signing_key_bytes: Vec<u8> = sqlx::query_scalar(
        "SELECT signing_key FROM users WHERE did = $1",
    )
    .bind(did)
    .fetch_one(&state.db)
    .await
    .map_err(|e| {
        err(
            StatusCode::INTERNAL_SERVER_ERROR,
            "InternalServerError",
            format!("users.signing_key read: {e}"),
        )
    })?;
    let signing_key = load_signing_key(&signing_key_bytes)?;

    let blockstore = load_user_blockstore(state, did).await?;

    // 4. Build the in-memory Repo. Fresh accounts have the all-zero
    //    sentinel in head_cid / head_commit and start empty.
    let mut repo: Repo<MemoryBlockstore> = if is_zero_blob(&head_cid_blob)
        || is_zero_blob(&head_commit_blob)
    {
        Repo::new(did.to_string(), signing_key.clone(), blockstore.clone())
    } else {
        let head_cid = cid_from_multihash_bytes(&head_cid_blob).map_err(|e| {
            err(
                StatusCode::INTERNAL_SERVER_ERROR,
                "InternalServerError",
                format!("invalid head_cid bytes: {e}"),
            )
        })?;
        // Defensive: re-seed the head commit block in case it hasn't
        // been flushed into the user's blockstore. Without this, a
        // load immediately after a previous put_record could miss the
        // head block.
        blockstore
            .put(&head_cid, Bytes::from(head_commit_blob.clone()))
            .await
            .map_err(|e| {
                err(
                    StatusCode::INTERNAL_SERVER_ERROR,
                    "InternalServerError",
                    format!("seed head commit block: {e}"),
                )
            })?;
        Repo::load(
            did.to_string(),
            signing_key.clone(),
            blockstore.clone(),
            head_cid,
        )
        .await
        .map_err(|e| {
            err(
                StatusCode::INTERNAL_SERVER_ERROR,
                "InternalServerError",
                format!("repo load: {e:#}"),
            )
        })?
    };

    // 5. Run the caller's closure. The closure may add records, delete
    //    records, or do whatever else the repo supports. It receives a
    //    mutable reference to the repo and returns a future that
    //    completes once it's finished mutating + committing.
    //
    //    The transaction (`tx`) is *not* passed to the closure — none
    //    of the current write paths need it. If a future caller needs
    //    to run additional queries under the row lock, we'd extend
    //    this helper to also hand out a `&mut PgConnection` (which
    //    doesn't have the lifetime headache of `&mut Transaction`).
    let outcome: RepoWriteOutcome = f(&mut repo).await.map_err(|e| {
        err(
            StatusCode::INTERNAL_SERVER_ERROR,
            "InternalServerError",
            format!("repo mutation: {e:#}"),
        )
    })?;

    // 6. Persist every newly produced block (commit block + MST nodes +
    //    value blocks the closure added). We re-serialise the repo
    //    after the closure returns to make sure we capture everything
    //    `Repo::commit` produced — `Repo::commit` writes its commit
    //    block to the blockstore but `serialize_repo` is the canonical
    //    "what's in this repo right now" dump.
    let (_header, all_blocks) = repo.serialize_repo().await.map_err(|e| {
        err(
            StatusCode::INTERNAL_SERVER_ERROR,
            "InternalServerError",
            format!("repo.serialize_repo: {e:#}"),
        )
    })?;
    persist_user_blocks_in_tx(&mut tx, did, &all_blocks).await?;

    // 7. Update the head pointer. The prev_commit column carries the
    //    head CID we read under the lock — that's the CID the new
    //    commit's `prev` field also points at.
    let prev_param: Option<Vec<u8>> = if is_zero_blob(&head_cid_blob) {
        None
    } else {
        Some(head_cid_blob.clone())
    };
    sqlx::query(
        r#"UPDATE repos
            SET rev = $2,
                head_cid = $3,
                head_commit = $4,
                prev_commit = $5,
                indexed_at = now()
          WHERE did = $1"#,
    )
    .bind(did)
    .bind(&outcome.commit.rev)
    .bind(&outcome.head_cid_bytes)
    .bind(&outcome.head_commit_bytes)
    .bind(prev_param.as_deref())
    .execute(&mut *tx)
    .await
    .map_err(|e| {
        err(
            StatusCode::INTERNAL_SERVER_ERROR,
            "InternalServerError",
            format!("repos update: {e}"),
        )
    })?;

    tx.commit().await.map_err(|e| {
        err(
            StatusCode::INTERNAL_SERVER_ERROR,
            "InternalServerError",
            format!("tx commit: {e}"),
        )
    })?;

    Ok(outcome)
}

/// Persist every block in `blocks` into `repo_blocks` using the open
/// transaction. Mirrors the connection-pool version but uses the
/// transaction's connection so the writes are part of the same atomic
/// unit as the head pointer update.
async fn persist_user_blocks_in_tx(
    tx: &mut sqlx::Transaction<'_, Postgres>,
    did: &str,
    blocks: &std::collections::HashMap<Cid, Vec<u8>>,
) -> Result<(), (StatusCode, Json<ErrorBody>)> {
    for (cid, bytes) in blocks {
        if cid.to_bytes().iter().all(|b| *b == 0) {
            continue;
        }
        sqlx::query(
            r#"INSERT INTO repo_blocks (did, cid, block, size)
               VALUES ($1, $2, $3, $4)
               ON CONFLICT (did, cid) DO NOTHING"#,
        )
        .bind(did)
        .bind(cid.to_bytes().as_slice())
        .bind(bytes.as_slice())
        .bind(bytes.len() as i32)
        .execute(&mut **tx)
        .await
        .map_err(|e| {
            err(
                StatusCode::INTERNAL_SERVER_ERROR,
                "InternalServerError",
                format!("repo_blocks insert: {e}"),
            )
        })?;
    }
    Ok(())
}