Genarrative/server-rs/crates/shared-contracts/src/editor_generation.rs

use sha2::{Digest, Sha256};

pub const EDITOR_GENERATION_RESOURCE_ID_PREFIX: &str = "editor-resource-";
pub const EDITOR_GENERATION_ASSET_ID_PREFIX: &str = "editor-asset-";
pub const EDITOR_GENERATION_ASSET_OBJECT_ID_PREFIX: &str = "assetobj_";
pub const EDITOR_GENERATION_ASSET_BINDING_ID_PREFIX: &str = "assetbind_";

pub fn editor_generation_request_fingerprint(
    operation_kind: &str,
    request_payload_json: &str,
) -> String {
    sha256_hex(
        format!(
            "editor-generation-request-v1\0{}\0{}",
            operation_kind.trim(),
            request_payload_json
        )
        .as_bytes(),
    )
}

pub fn editor_generation_stable_asset_object_id(
    owner_user_id: &str,
    operation_kind: &str,
    operation_id: &str,
    slot: &str,
) -> String {
    editor_generation_stable_record_id(
        EDITOR_GENERATION_ASSET_OBJECT_ID_PREFIX,
        owner_user_id,
        operation_kind,
        operation_id,
        slot,
        "asset-object",
    )
}

pub fn editor_generation_stable_asset_binding_id(
    owner_user_id: &str,
    operation_kind: &str,
    operation_id: &str,
    slot: &str,
) -> String {
    editor_generation_stable_record_id(
        EDITOR_GENERATION_ASSET_BINDING_ID_PREFIX,
        owner_user_id,
        operation_kind,
        operation_id,
        slot,
        "asset-binding",
    )
}

pub fn editor_generation_stable_resource_id(
    owner_user_id: &str,
    operation_kind: &str,
    operation_id: &str,
    slot: &str,
) -> String {
    editor_generation_stable_record_id(
        EDITOR_GENERATION_RESOURCE_ID_PREFIX,
        owner_user_id,
        operation_kind,
        operation_id,
        slot,
        "project-resource",
    )
}

pub fn editor_generation_stable_asset_id(
    owner_user_id: &str,
    operation_kind: &str,
    operation_id: &str,
    slot: &str,
) -> String {
    editor_generation_stable_record_id(
        EDITOR_GENERATION_ASSET_ID_PREFIX,
        owner_user_id,
        operation_kind,
        operation_id,
        slot,
        "account-asset",
    )
}

fn editor_generation_stable_record_id(
    prefix: &str,
    owner_user_id: &str,
    operation_kind: &str,
    operation_id: &str,
    slot: &str,
    record_kind: &str,
) -> String {
    let digest = sha256_hex(
        format!(
            "editor-generation-result-v1\0{}\0{}\0{}\0{}\0{}",
            owner_user_id.trim(),
            operation_kind.trim(),
            operation_id.trim(),
            slot.trim(),
            record_kind,
        )
        .as_bytes(),
    );
    format!("{prefix}{}", &digest[..32])
}

fn sha256_hex(bytes: &[u8]) -> String {
    format!("{:x}", Sha256::digest(bytes))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn stable_editor_generation_ids_are_replay_safe_and_domain_separated() {
        let first = editor_generation_stable_resource_id(
            "owner-1",
            "editor_image_generation",
            "task-1",
            "primary",
        );
        assert_eq!(
            first,
            editor_generation_stable_resource_id(
                "owner-1",
                "editor_image_generation",
                "task-1",
                "primary",
            )
        );
        assert_ne!(
            first,
            editor_generation_stable_resource_id(
                "owner-1",
                "editor_image_generation",
                "task-1",
                "provider-source",
            )
        );
        assert_ne!(
            first,
            editor_generation_stable_asset_id(
                "owner-1",
                "editor_image_generation",
                "task-1",
                "primary",
            )
        );
        assert_ne!(
            first,
            editor_generation_stable_asset_object_id(
                "owner-1",
                "editor_image_generation",
                "task-1",
                "primary",
            )
        );
        assert_eq!(
            editor_generation_stable_asset_object_id(
                "owner-1",
                "editor_image_generation",
                "task-1",
                "primary",
            )
            .len(),
            EDITOR_GENERATION_ASSET_OBJECT_ID_PREFIX.len() + 32
        );
        assert_ne!(
            editor_generation_stable_asset_object_id(
                "owner-1",
                "editor_image_generation",
                "task-1",
                "primary",
            ),
            editor_generation_stable_asset_binding_id(
                "owner-1",
                "editor_image_generation",
                "task-1",
                "primary",
            )
        );
    }

    #[test]
    fn request_fingerprint_binds_kind_and_exact_payload() {
        let first = editor_generation_request_fingerprint("image", r#"{"prompt":"a"}"#);
        assert_eq!(
            first,
            editor_generation_request_fingerprint("image", r#"{"prompt":"a"}"#)
        );
        assert_ne!(
            first,
            editor_generation_request_fingerprint("image", r#"{"prompt":"b"}"#)
        );
        assert_ne!(
            first,
            editor_generation_request_fingerprint("video", r#"{"prompt":"a"}"#)
        );
    }
}