ImageApi/src/versions/xmp.rs

//! Reading the version link out of a PNG: the `XML:com.adobe.xmp` iTXt
//! chunk the ComfyUI producer writes (`comfyui-imageapi-versions`), and the
//! ComfyUI `prompt` tEXt chunk the backfill uses for files without XMP.
//!
//! The chunk stream is walked by hand: no pixel decode, and the `png`
//! crate's text API would still inflate every IDAT's neighbours for no gain.
//! IDAT payloads are seeked over, so a multi-megabyte restoration costs a
//! handful of small reads on a network share.

use std::fs::File;
use std::io::{self, BufReader, Read, Seek, SeekFrom};
use std::path::Path;
use std::sync::LazyLock;

use regex::Regex;

/// iTXt keyword exiftool, Lightroom and the producer put XMP under.
pub const XMP_KEYWORD: &str = "XML:com.adobe.xmp";

/// Text chunks bigger than this are skipped rather than read: ComfyUI's
/// `prompt` / `workflow` JSON is tens of KiB, an XMP packet around 1 KiB.
const MAX_TEXT_CHUNK: u32 = 4 * 1024 * 1024;

const PNG_SIGNATURE: [u8; 8] = [0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A];

/// What a version PNG's XMP says about its provenance.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VersionXmp {
    /// `xmpMM:DerivedFrom/stRef:documentID`: the original's `content_hash`.
    pub original_hash: String,
    /// `xmpMM:Ingredients[]/stRef:documentID`: reference photos.
    pub references: Vec<String>,
    /// `imageapi:versionKind`, as written (not yet validated).
    pub kind: Option<String>,
}

/// Value of the first text chunk (`tEXt` or uncompressed `iTXt`) whose
/// keyword is in `keywords`, for each keyword found. Compressed chunks are
/// skipped (the producer never compresses). `Ok(vec![])` for a file that
/// isn't a PNG.
pub fn read_png_text(path: &Path, keywords: &[&str]) -> io::Result<Vec<(String, String)>> {
    let mut reader = BufReader::new(File::open(path)?);
    let mut sig = [0u8; 8];
    if reader.read_exact(&mut sig).is_err() || sig != PNG_SIGNATURE {
        return Ok(Vec::new());
    }
    let mut found: Vec<(String, String)> = Vec::new();
    loop {
        let mut head = [0u8; 8];
        if reader.read_exact(&mut head).is_err() {
            break; // truncated file: use what we have
        }
        let len = u32::from_be_bytes([head[0], head[1], head[2], head[3]]);
        let kind = &head[4..8];
        let is_text = kind == b"tEXt" || kind == b"iTXt";
        if is_text && len <= MAX_TEXT_CHUNK {
            let mut data = vec![0u8; len as usize];
            if reader.read_exact(&mut data).is_err() {
                break;
            }
            let parsed = if kind == b"tEXt" {
                parse_text(&data)
            } else {
                parse_itxt(&data)
            };
            if let Some((key, value)) = parsed
                && keywords.contains(&key.as_str())
                && !found.iter().any(|(k, _)| *k == key)
            {
                found.push((key, value));
                if found.len() == keywords.len() {
                    break;
                }
            }
            reader.seek_relative(4)?; // CRC
        } else {
            reader.seek(SeekFrom::Current(i64::from(len) + 4))?;
        }
        if kind == b"IEND" {
            break;
        }
    }
    Ok(found)
}

/// `keyword\0text` (Latin-1).
fn parse_text(data: &[u8]) -> Option<(String, String)> {
    let nul = data.iter().position(|b| *b == 0)?;
    let key = String::from_utf8_lossy(&data[..nul]).into_owned();
    let value: String = data[nul + 1..].iter().map(|b| *b as char).collect();
    Some((key, value))
}

/// `keyword\0 compressed(1) method(1) lang\0 translated\0 text` (UTF-8);
/// `None` for compressed chunks.
fn parse_itxt(data: &[u8]) -> Option<(String, String)> {
    let nul = data.iter().position(|b| *b == 0)?;
    let key = String::from_utf8_lossy(&data[..nul]).into_owned();
    let rest = data.get(nul + 1..)?;
    if *rest.first()? != 0 {
        return None;
    }
    let rest = rest.get(2..)?; // flag + method
    let lang_end = rest.iter().position(|b| *b == 0)?;
    let rest = &rest[lang_end + 1..];
    let translated_end = rest.iter().position(|b| *b == 0)?;
    let text = &rest[translated_end + 1..];
    Some((key, String::from_utf8_lossy(text).into_owned()))
}

/// The XMP packet of a PNG, if it has one.
pub fn read_png_xmp(path: &Path) -> io::Result<Option<String>> {
    Ok(read_png_text(path, &[XMP_KEYWORD])?
        .into_iter()
        .next()
        .map(|(_, v)| v))
}

static DERIVED_FROM: LazyLock<Regex> = LazyLock::new(|| {
    // Element form (`<xmpMM:DerivedFrom ...>...</xmpMM:DerivedFrom>`) or the
    // attribute form exiftool may write (`<xmpMM:DerivedFrom stRef:documentID=".."/>`).
    Regex::new(r"(?s)<xmpMM:DerivedFrom\b([^>]*?)(?:/>|>(.*?)</xmpMM:DerivedFrom>)").unwrap()
});
static INGREDIENTS: LazyLock<Regex> = LazyLock::new(|| {
    Regex::new(r"(?s)<xmpMM:Ingredients\b[^>]*>(.*?)</xmpMM:Ingredients>").unwrap()
});
static DOCUMENT_ID: LazyLock<Regex> = LazyLock::new(|| {
    Regex::new(r#"documentID(?:\s*=\s*["']|\s*>\s*)blake3:([0-9a-fA-F]{64})"#).unwrap()
});
static VERSION_KIND: LazyLock<Regex> = LazyLock::new(|| {
    Regex::new(r#"imageapi:versionKind(?:\s*=\s*["']|\s*>\s*)([A-Za-z]+)"#).unwrap()
});

/// Parse the link fields out of an XMP packet. `None` when it has no
/// `DerivedFrom` blake3 id (an ordinary photo's XMP, or someone else's).
pub fn parse_version_xmp(xmp: &str) -> Option<VersionXmp> {
    let derived = DERIVED_FROM.captures(xmp)?;
    let region = format!(
        "{} {}",
        derived.get(1).map_or("", |m| m.as_str()),
        derived.get(2).map_or("", |m| m.as_str())
    );
    let original_hash = DOCUMENT_ID.captures(&region)?[1].to_ascii_lowercase();

    let mut references: Vec<String> = Vec::new();
    if let Some(ing) = INGREDIENTS.captures(xmp) {
        for cap in DOCUMENT_ID.captures_iter(&ing[1]) {
            let h = cap[1].to_ascii_lowercase();
            if !references.contains(&h) {
                references.push(h);
            }
        }
    }
    let kind = VERSION_KIND
        .captures(xmp)
        .map(|c| c[1].to_ascii_lowercase());
    Some(VersionXmp {
        original_hash,
        references,
        kind,
    })
}

/// Parse the version link out of a PNG file; `None` if it has none.
pub fn read_version_xmp(path: &Path) -> io::Result<Option<VersionXmp>> {
    Ok(read_png_xmp(path)?.and_then(|x| parse_version_xmp(&x)))
}

#[cfg(test)]
pub(crate) mod test_png {
    use std::io::Write;

    fn chunk(out: &mut Vec<u8>, kind: &[u8; 4], data: &[u8]) {
        out.extend_from_slice(&(data.len() as u32).to_be_bytes());
        out.extend_from_slice(kind);
        out.extend_from_slice(data);
        out.extend_from_slice(&[0, 0, 0, 0]); // CRC: the reader never checks it
    }

    /// A PNG-shaped byte stream: signature, IHDR, optional tEXt `prompt`,
    /// optional XMP iTXt, an IDAT, IEND. Not decodable; enough for the
    /// chunk walker (and unique per `salt`, so each file hashes differently).
    pub fn png_bytes(xmp: Option<&str>, prompt: Option<&str>, salt: u8) -> Vec<u8> {
        let mut out = super::PNG_SIGNATURE.to_vec();
        chunk(&mut out, b"IHDR", &[0, 0, 0, 1, 0, 0, 0, 1, 8, 2, 0, 0, 0]);
        if let Some(p) = prompt {
            let mut d = b"prompt\0".to_vec();
            d.extend_from_slice(p.as_bytes());
            chunk(&mut out, b"tEXt", &d);
        }
        if let Some(x) = xmp {
            let mut d = b"XML:com.adobe.xmp\0\0\0\0\0".to_vec();
            d.extend_from_slice(x.as_bytes());
            chunk(&mut out, b"iTXt", &d);
        }
        chunk(&mut out, b"IDAT", &[salt; 32]);
        chunk(&mut out, b"IEND", &[]);
        out
    }

    pub fn write_png(path: &std::path::Path, xmp: Option<&str>, prompt: Option<&str>, salt: u8) {
        if let Some(dir) = path.parent() {
            std::fs::create_dir_all(dir).unwrap();
        }
        std::fs::File::create(path)
            .unwrap()
            .write_all(&png_bytes(xmp, prompt, salt))
            .unwrap();
    }

    /// The packet the producer writes, built the way `versioning.build_xmp` does.
    pub fn producer_xmp(original: &str, kind: &str, references: &[&str]) -> String {
        let ingredients = if references.is_empty() {
            String::new()
        } else {
            let items: String = references
                .iter()
                .map(|h| {
                    format!(
                        "<rdf:li rdf:parseType=\"Resource\"><stRef:documentID>blake3:{h}</stRef:documentID>\
                         <stRef:filePath>ref.jpg</stRef:filePath></rdf:li>"
                    )
                })
                .collect();
            format!("<xmpMM:Ingredients><rdf:Bag>{items}</rdf:Bag></xmpMM:Ingredients>")
        };
        format!(
            "<?xpacket begin=\"\u{feff}\" id=\"W5M0MpCehiHzreSzNTczkc9d\"?><x:xmpmeta xmlns:x=\"adobe:ns:meta/\">\
             <rdf:RDF xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\"><rdf:Description rdf:about=\"\" \
             xmlns:xmpMM=\"http://ns.adobe.com/xap/1.0/mm/\" xmlns:stRef=\"http://ns.adobe.com/xap/1.0/sType/ResourceRef#\" \
             xmlns:imageapi=\"https://git.crkc.xyz/Apps/ImageApi/ns/versions/1.0/\">\
             <xmpMM:DerivedFrom rdf:parseType=\"Resource\"><stRef:documentID>blake3:{original}</stRef:documentID>\
             <stRef:filePath>orig.jpg</stRef:filePath></xmpMM:DerivedFrom>{ingredients}\
             <imageapi:versionKind>{kind}</imageapi:versionKind></rdf:Description></rdf:RDF></x:xmpmeta><?xpacket end=\"w\"?>"
        )
    }
}

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

    const ORIG: &str = "d5fa14088629f680f4f8e5815df473cc4684a97fc937da588baf72d2693dea3a";
    const REF: &str = "0000000000000000000000000000000000000000000000000000000000000001";

    /// The packet embedded in a real archive restoration (copied verbatim).
    const REAL_XMP: &str = r#"<?xpacket begin="" id="W5M0MpCehiHzreSzNTczkc9d"?><x:xmpmeta xmlns:x="adobe:ns:meta/"><rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"><rdf:Description rdf:about="" xmlns:xmpMM="http://ns.adobe.com/xap/1.0/mm/" xmlns:stRef="http://ns.adobe.com/xap/1.0/sType/ResourceRef#" xmlns:imageapi="https://git.crkc.xyz/Apps/ImageApi/ns/versions/1.0/"><xmpMM:DerivedFrom rdf:parseType="Resource"><stRef:documentID>blake3:d5fa14088629f680f4f8e5815df473cc4684a97fc937da588baf72d2693dea3a</stRef:documentID><stRef:filePath>019_19.JPG</stRef:filePath></xmpMM:DerivedFrom><imageapi:versionKind>restored</imageapi:versionKind></rdf:Description></rdf:RDF></x:xmpmeta><?xpacket end="w"?>"#;

    #[test]
    fn parses_the_real_producer_packet() {
        let x = parse_version_xmp(REAL_XMP).unwrap();
        assert_eq!(x.original_hash, ORIG);
        assert_eq!(x.kind.as_deref(), Some("restored"));
        assert!(x.references.is_empty());
    }

    #[test]
    fn parses_references_and_keeps_them_apart_from_the_original() {
        let x = parse_version_xmp(&producer_xmp(ORIG, "edit", &[REF])).unwrap();
        assert_eq!(x.original_hash, ORIG);
        assert_eq!(x.references, vec![REF]);
        assert_eq!(x.kind.as_deref(), Some("edit"));
    }

    #[test]
    fn accepts_the_attribute_form_other_writers_use() {
        let xmp = format!(
            r#"<x:xmpmeta><rdf:RDF><rdf:Description><xmpMM:DerivedFrom stRef:documentID="blake3:{ORIG}" stRef:filePath="a.jpg"/></rdf:Description></rdf:RDF></x:xmpmeta>"#
        );
        assert_eq!(parse_version_xmp(&xmp).unwrap().original_hash, ORIG);
    }

    #[test]
    fn ignores_xmp_without_a_blake3_derived_from() {
        assert!(parse_version_xmp("<x:xmpmeta><dc:title>hi</dc:title></x:xmpmeta>").is_none());
        let other = "<xmpMM:DerivedFrom><stRef:documentID>xmp.did:1234</stRef:documentID></xmpMM:DerivedFrom>";
        assert!(parse_version_xmp(other).is_none());
        // An ingredient alone is not a link.
        let only_ing = format!(
            "<xmpMM:Ingredients><rdf:Bag><rdf:li><stRef:documentID>blake3:{REF}</stRef:documentID></rdf:li></rdf:Bag></xmpMM:Ingredients>"
        );
        assert!(parse_version_xmp(&only_ing).is_none());
    }

    #[test]
    fn reads_xmp_and_prompt_chunks_from_a_png_file() {
        let dir = tempfile::tempdir().unwrap();
        let p = dir.path().join("v.png");
        write_png(
            &p,
            Some(&producer_xmp(ORIG, "restored", &[])),
            Some(r#"{"1": {"class_type": "LoadImage"}}"#),
            1,
        );
        let got = read_version_xmp(&p).unwrap().unwrap();
        assert_eq!(got.original_hash, ORIG);
        let text = read_png_text(&p, &["prompt"]).unwrap();
        assert_eq!(text[0].1, r#"{"1": {"class_type": "LoadImage"}}"#);
    }

    #[test]
    fn plain_png_and_non_png_have_no_link() {
        let dir = tempfile::tempdir().unwrap();
        let p = dir.path().join("plain.png");
        write_png(&p, None, None, 2);
        assert!(read_version_xmp(&p).unwrap().is_none());
        let j = dir.path().join("x.jpg");
        std::fs::write(&j, b"\xff\xd8\xff\xe0 not a png").unwrap();
        assert!(read_version_xmp(&j).unwrap().is_none());
    }
}