crank/crates/crank-registry/src/migrations/schema_guard_v7.rs

use sqlx::{PgConnection, Row, query};

use super::{
    authority::MigrationError,
    schema_guard::{
        column_exists, constraint_expression, enum_constraint_matches, normalize_definition,
        relation_exists, schema_error,
    },
};

pub(super) async fn validate_v7_master_key_identity(
    connection: &mut PgConnection,
) -> Result<(), MigrationError> {
    for relation in ["master_key_identities", "master_key_rotations"] {
        if !relation_exists(connection, relation).await? {
            return Err(schema_error(7));
        }
    }
    for (column, data_type, nullable) in [
        ("master_key_epoch", "bigint", false),
        ("target_ciphertext", "text", true),
        ("target_key_version", "text", true),
        ("target_master_key_epoch", "bigint", true),
    ] {
        if !secret_version_column_matches(connection, column, data_type, nullable).await? {
            return Err(schema_error(7));
        }
    }
    validate_identity_constraints(connection).await?;
    validate_rotation_constraints(connection).await?;
    validate_secret_version_constraints(connection).await?;
    validate_active_identity_index(connection).await
}

async fn secret_version_column_matches(
    connection: &mut PgConnection,
    column: &str,
    data_type: &str,
    nullable: bool,
) -> Result<bool, MigrationError> {
    if !column_exists(connection, "secret_versions", column).await? {
        return Ok(false);
    }
    let row = query(
        "select data_type, is_nullable from information_schema.columns
         where table_schema = current_schema()
           and table_name = 'secret_versions'
           and column_name = $1",
    )
    .bind(column)
    .fetch_optional(connection)
    .await
    .map_err(|_| MigrationError::storage("preflight.schema"))?;
    Ok(row.is_some_and(|row| {
        row.try_get::<String, _>("data_type").ok().as_deref() == Some(data_type)
            && row.try_get::<String, _>("is_nullable").ok().as_deref()
                == Some(if nullable { "YES" } else { "NO" })
    }))
}

async fn validate_identity_constraints(
    connection: &mut PgConnection,
) -> Result<(), MigrationError> {
    let identity_epoch = constraint_expression(
        connection,
        "master_key_identities",
        "master_key_identities_epoch_check",
    )
    .await?;
    if normalize_definition(&identity_epoch) != "epoch>0" {
        return Err(schema_error(7));
    }
    let identity_fingerprint = constraint_expression(
        connection,
        "master_key_identities",
        "master_key_identities_fingerprint_check",
    )
    .await?;
    if !regex_constraint_matches(&identity_fingerprint, "fingerprint", "^[0-9a-f]{64}$") {
        return Err(schema_error(7));
    }
    let identity_cipher = constraint_expression(
        connection,
        "master_key_identities",
        "master_key_identities_cipher_contract_check",
    )
    .await?;
    if !enum_constraint_matches(
        &identity_cipher,
        "cipher_contract",
        &["secret-envelope-v2/aes-256-gcm-hkdf-sha256"],
    ) {
        return Err(schema_error(7));
    }
    let identity_status = constraint_expression(
        connection,
        "master_key_identities",
        "master_key_identities_status_check",
    )
    .await?;
    if enum_constraint_matches(
        &identity_status,
        "status",
        &["active", "pending", "retired", "revoked"],
    ) {
        Ok(())
    } else {
        Err(schema_error(7))
    }
}

async fn validate_rotation_constraints(
    connection: &mut PgConnection,
) -> Result<(), MigrationError> {
    let rotation_state = constraint_expression(
        connection,
        "master_key_rotations",
        "master_key_rotations_state_check",
    )
    .await?;
    if enum_constraint_matches(
        &rotation_state,
        "state",
        &[
            "preflighted",
            "running",
            "verifying",
            "verified",
            "promoted",
            "aborted",
            "failed",
        ],
    ) {
        Ok(())
    } else {
        Err(schema_error(7))
    }
}

async fn validate_secret_version_constraints(
    connection: &mut PgConnection,
) -> Result<(), MigrationError> {
    let secret_epoch = constraint_expression(
        connection,
        "secret_versions",
        "secret_versions_master_key_epoch_check",
    )
    .await?;
    if normalize_definition(&secret_epoch) != "master_key_epoch>0" {
        return Err(schema_error(7));
    }
    let secret_target_epoch = constraint_expression(
        connection,
        "secret_versions",
        "secret_versions_target_epoch_check",
    )
    .await?;
    if normalize_definition(&secret_target_epoch)
        != "target_master_key_epochisnullortarget_master_key_epoch>master_key_epoch"
    {
        return Err(schema_error(7));
    }
    let secret_all_or_none = constraint_expression(
        connection,
        "secret_versions",
        "secret_versions_target_all_or_none_check",
    )
    .await?;
    if target_ciphertext_all_or_none_matches(&secret_all_or_none) {
        Ok(())
    } else {
        Err(schema_error(7))
    }
}

async fn validate_active_identity_index(
    connection: &mut PgConnection,
) -> Result<(), MigrationError> {
    let active_index = query(
        "select t.relname as table_name, am.amname as access_method, i.indisvalid, i.indisready,
                i.indisunique, pg_get_indexdef(i.indexrelid, 1, true) as first_column,
                pg_get_expr(i.indpred, i.indrelid) as predicate
         from pg_catalog.pg_index i
         join pg_catalog.pg_class idx on idx.oid = i.indexrelid
         join pg_catalog.pg_class t on t.oid = i.indrelid
         join pg_catalog.pg_namespace n on n.oid = t.relnamespace
         join pg_catalog.pg_am am on am.oid = idx.relam
         where n.nspname = current_schema()
           and idx.relname = 'master_key_identities_active_idx'",
    )
    .fetch_optional(connection)
    .await
    .map_err(|_| MigrationError::storage("preflight.schema"))?;
    let valid = active_index.is_some_and(|row| {
        row.try_get::<String, _>("table_name").ok().as_deref() == Some("master_key_identities")
            && row.try_get::<String, _>("access_method").ok().as_deref() == Some("btree")
            && row.try_get::<bool, _>("indisvalid").ok() == Some(true)
            && row.try_get::<bool, _>("indisready").ok() == Some(true)
            && row.try_get::<bool, _>("indisunique").ok() == Some(true)
            && row.try_get::<String, _>("first_column").ok().as_deref() == Some("status")
            && row
                .try_get::<String, _>("predicate")
                .ok()
                .is_some_and(|value| normalize_definition(&value) == "status='active'::text")
    });
    if valid { Ok(()) } else { Err(schema_error(7)) }
}

fn regex_constraint_matches(expression: &str, column: &str, pattern: &str) -> bool {
    let normalized = normalize_definition(expression);
    if !normalized.contains(column) || normalized.contains("ortrue") {
        return false;
    }
    let values = expression
        .split('\'')
        .enumerate()
        .filter_map(|(index, value)| (index % 2 == 1).then_some(value))
        .collect::<Vec<_>>();
    values == [pattern]
}

fn target_ciphertext_all_or_none_matches(expression: &str) -> bool {
    let normalized = normalize_definition(expression);
    !normalized.contains("ortrue")
        && normalized.contains("target_ciphertextisnull")
        && normalized.contains("target_key_versionisnull")
        && normalized.contains("target_master_key_epochisnull")
        && normalized.contains("target_ciphertextisnotnull")
        && normalized.contains("target_key_versionisnotnull")
        && normalized.contains("target_master_key_epochisnotnull")
}