secrets/crates/secrets-core/src/service/rollback.rs

use anyhow::Result;
use serde_json::Value;
use sqlx::PgPool;
use uuid::Uuid;

use crate::crypto;
use crate::db;

#[derive(Debug, serde::Serialize)]
pub struct RollbackResult {
    pub name: String,
    pub folder: String,
    #[serde(rename = "type")]
    pub entry_type: String,
    pub restored_version: i64,
}

/// Roll back entry `name` to `to_version` (or the most recent snapshot if None).
/// `folder` is optional; if omitted and multiple entries share the name, an error is returned.
pub async fn run(
    pool: &PgPool,
    name: &str,
    folder: Option<&str>,
    to_version: Option<i64>,
    master_key: &[u8; 32],
    user_id: Option<Uuid>,
) -> Result<RollbackResult> {
    #[derive(sqlx::FromRow)]
    struct EntryHistoryRow {
        entry_id: Uuid,
        folder: String,
        #[sqlx(rename = "type")]
        entry_type: String,
        version: i64,
        action: String,
        tags: Vec<String>,
        metadata: Value,
    }

    // Disambiguate: find the unique entry_id for (name, folder).
    // Query entries_history by entry_id once we know it; first resolve via name + optional folder.
    let entry_id: Option<Uuid> = if let Some(uid) = user_id {
        if let Some(f) = folder {
            sqlx::query_scalar(
                "SELECT DISTINCT entry_id FROM entries_history \
                 WHERE name = $1 AND folder = $2 AND user_id = $3 LIMIT 1",
            )
            .bind(name)
            .bind(f)
            .bind(uid)
            .fetch_optional(pool)
            .await?
        } else {
            let ids: Vec<Uuid> = sqlx::query_scalar(
                "SELECT DISTINCT entry_id FROM entries_history \
                 WHERE name = $1 AND user_id = $2",
            )
            .bind(name)
            .bind(uid)
            .fetch_all(pool)
            .await?;
            match ids.len() {
                0 => None,
                1 => Some(ids[0]),
                _ => {
                    let folders: Vec<String> = sqlx::query_scalar(
                        "SELECT DISTINCT folder FROM entries_history \
                         WHERE name = $1 AND user_id = $2",
                    )
                    .bind(name)
                    .bind(uid)
                    .fetch_all(pool)
                    .await?;
                    anyhow::bail!(
                        "Ambiguous: entries named '{}' exist in folders: [{}]. \
                         Specify 'folder' to disambiguate.",
                        name,
                        folders.join(", ")
                    )
                }
            }
        }
    } else if let Some(f) = folder {
        sqlx::query_scalar(
            "SELECT DISTINCT entry_id FROM entries_history \
             WHERE name = $1 AND folder = $2 AND user_id IS NULL LIMIT 1",
        )
        .bind(name)
        .bind(f)
        .fetch_optional(pool)
        .await?
    } else {
        let ids: Vec<Uuid> = sqlx::query_scalar(
            "SELECT DISTINCT entry_id FROM entries_history \
             WHERE name = $1 AND user_id IS NULL",
        )
        .bind(name)
        .fetch_all(pool)
        .await?;
        match ids.len() {
            0 => None,
            1 => Some(ids[0]),
            _ => {
                let folders: Vec<String> = sqlx::query_scalar(
                    "SELECT DISTINCT folder FROM entries_history \
                     WHERE name = $1 AND user_id IS NULL",
                )
                .bind(name)
                .fetch_all(pool)
                .await?;
                anyhow::bail!(
                    "Ambiguous: entries named '{}' exist in folders: [{}]. \
                     Specify 'folder' to disambiguate.",
                    name,
                    folders.join(", ")
                )
            }
        }
    };

    let entry_id = entry_id.ok_or_else(|| anyhow::anyhow!("No history found for '{}'", name))?;

    let snap: Option<EntryHistoryRow> = if let Some(ver) = to_version {
        sqlx::query_as(
            "SELECT entry_id, folder, type, version, action, tags, metadata \
             FROM entries_history \
             WHERE entry_id = $1 AND version = $2 ORDER BY id DESC LIMIT 1",
        )
        .bind(entry_id)
        .bind(ver)
        .fetch_optional(pool)
        .await?
    } else {
        sqlx::query_as(
            "SELECT entry_id, folder, type, version, action, tags, metadata \
             FROM entries_history \
             WHERE entry_id = $1 ORDER BY id DESC LIMIT 1",
        )
        .bind(entry_id)
        .fetch_optional(pool)
        .await?
    };

    let snap = snap.ok_or_else(|| {
        anyhow::anyhow!(
            "No history found for '{}'{}.",
            name,
            to_version
                .map(|v| format!(" at version {}", v))
                .unwrap_or_default()
        )
    })?;

    #[derive(sqlx::FromRow)]
    struct SecretHistoryRow {
        field_name: String,
        encrypted: Vec<u8>,
        action: String,
    }

    let field_snaps: Vec<SecretHistoryRow> = sqlx::query_as(
        "SELECT field_name, encrypted, action FROM secrets_history \
         WHERE entry_id = $1 AND entry_version = $2 ORDER BY field_name",
    )
    .bind(snap.entry_id)
    .bind(snap.version)
    .fetch_all(pool)
    .await?;

    for f in &field_snaps {
        if f.action != "delete" && !f.encrypted.is_empty() {
            crypto::decrypt_json(master_key, &f.encrypted).map_err(|e| {
                anyhow::anyhow!(
                    "Cannot decrypt snapshot for field '{}': {}",
                    f.field_name,
                    e
                )
            })?;
        }
    }

    let mut tx = pool.begin().await?;

    #[derive(sqlx::FromRow)]
    struct LiveEntry {
        id: Uuid,
        version: i64,
        folder: String,
        #[sqlx(rename = "type")]
        entry_type: String,
        tags: Vec<String>,
        metadata: Value,
        #[allow(dead_code)]
        notes: String,
    }

    // Lock the live entry if it exists (matched by entry_id for precision).
    let live: Option<LiveEntry> = sqlx::query_as(
        "SELECT id, version, folder, type, tags, metadata, notes FROM entries \
         WHERE id = $1 FOR UPDATE",
    )
    .bind(entry_id)
    .fetch_optional(&mut *tx)
    .await?;

    let live_entry_id = if let Some(ref lr) = live {
        if let Err(e) = db::snapshot_entry_history(
            &mut tx,
            db::EntrySnapshotParams {
                entry_id: lr.id,
                user_id,
                folder: &lr.folder,
                entry_type: &lr.entry_type,
                name,
                version: lr.version,
                action: "rollback",
                tags: &lr.tags,
                metadata: &lr.metadata,
            },
        )
        .await
        {
            tracing::warn!(error = %e, "failed to snapshot entry before rollback");
        }

        #[derive(sqlx::FromRow)]
        struct LiveField {
            id: Uuid,
            field_name: String,
            encrypted: Vec<u8>,
        }
        let live_fields: Vec<LiveField> =
            sqlx::query_as("SELECT id, field_name, encrypted FROM secrets WHERE entry_id = $1")
                .bind(lr.id)
                .fetch_all(&mut *tx)
                .await?;

        for f in &live_fields {
            if let Err(e) = db::snapshot_secret_history(
                &mut tx,
                db::SecretSnapshotParams {
                    entry_id: lr.id,
                    secret_id: f.id,
                    entry_version: lr.version,
                    field_name: &f.field_name,
                    encrypted: &f.encrypted,
                    action: "rollback",
                },
            )
            .await
            {
                tracing::warn!(error = %e, "failed to snapshot secret field before rollback");
            }
        }

        sqlx::query(
            "UPDATE entries SET tags = $1, metadata = $2, version = version + 1, \
             updated_at = NOW() WHERE id = $3",
        )
        .bind(&snap.tags)
        .bind(&snap.metadata)
        .bind(lr.id)
        .execute(&mut *tx)
        .await?;

        lr.id
    } else {
        if let Some(uid) = user_id {
            sqlx::query_scalar(
                "INSERT INTO entries \
                 (user_id, folder, type, name, notes, tags, metadata, version, updated_at) \
                 VALUES ($1, $2, $3, $4, '', $5, $6, $7, NOW()) RETURNING id",
            )
            .bind(uid)
            .bind(&snap.folder)
            .bind(&snap.entry_type)
            .bind(name)
            .bind(&snap.tags)
            .bind(&snap.metadata)
            .bind(snap.version)
            .fetch_one(&mut *tx)
            .await?
        } else {
            sqlx::query_scalar(
                "INSERT INTO entries \
                 (folder, type, name, notes, tags, metadata, version, updated_at) \
                 VALUES ($1, $2, $3, '', $4, $5, $6, NOW()) RETURNING id",
            )
            .bind(&snap.folder)
            .bind(&snap.entry_type)
            .bind(name)
            .bind(&snap.tags)
            .bind(&snap.metadata)
            .bind(snap.version)
            .fetch_one(&mut *tx)
            .await?
        }
    };

    sqlx::query("DELETE FROM secrets WHERE entry_id = $1")
        .bind(live_entry_id)
        .execute(&mut *tx)
        .await?;

    for f in &field_snaps {
        if f.action == "delete" {
            continue;
        }
        sqlx::query("INSERT INTO secrets (entry_id, field_name, encrypted) VALUES ($1, $2, $3)")
            .bind(live_entry_id)
            .bind(&f.field_name)
            .bind(&f.encrypted)
            .execute(&mut *tx)
            .await?;
    }

    crate::audit::log_tx(
        &mut tx,
        user_id,
        "rollback",
        &snap.folder,
        &snap.entry_type,
        name,
        serde_json::json!({
            "restored_version": snap.version,
            "original_action": snap.action,
        }),
    )
    .await;

    tx.commit().await?;

    Ok(RollbackResult {
        name: name.to_string(),
        folder: snap.folder,
        entry_type: snap.entry_type,
        restored_version: snap.version,
    })
}