roxify 1.13.3 → 1.13.5

package/native/packer.rs

@@ -255,96 +255,210 @@ fn unpack_entries_sequential(buf: &[u8], start: usize, out_dir: &Path, files_opt
     Ok(written)
 }
 
-pub fn unpack_stream_to_dir<R: std::io::Read>(reader: &mut R, out_dir: &Path, files_opt: Option<&[String]>) -> Result<Vec<String>> {
+fn unpack_progress_percent(total_expected: u64, bytes_processed: u64, file_count: usize, processed_files: usize) -> u64 {
+    if total_expected > 0 {
+        return 10 + (bytes_processed.saturating_mul(89) / total_expected).min(89);
+    }
+    if file_count > 0 {
+        return 10 + ((processed_files as u64).saturating_mul(89) / file_count as u64).min(89);
+    }
+    10
+}
+
+fn report_unpack_progress(
+    progress: Option<&(dyn Fn(u64, u64, &str) + Send)>,
+    total_expected: u64,
+    bytes_processed: u64,
+    file_count: usize,
+    processed_files: usize,
+    last_pct: &mut u64,
+) {
+    if let Some(cb) = progress {
+        let pct = unpack_progress_percent(total_expected, bytes_processed, file_count, processed_files);
+        if pct > *last_pct {
+            *last_pct = pct;
+            cb(pct, 100, "extracting");
+        }
+    }
+}
+
+pub fn unpack_stream_to_dir<R: std::io::Read>(
+    reader: &mut R,
+    out_dir: &Path,
+    files_opt: Option<&[String]>,
+    progress: Option<&(dyn Fn(u64, u64, &str) + Send)>,
+    total_expected: u64,
+) -> Result<Vec<String>> {
     let mut written = Vec::new();
-    let mut buf: Vec<u8> = Vec::new();
-    let mut pos: usize = 0;
-    let mut temp = [0u8; 64 * 1024];
     let files_filter: Option<std::collections::HashSet<String>> = files_opt.map(|l| l.iter().map(|s| s.clone()).collect());
     let mut requested = files_filter.as_ref().map(|s| s.len()).unwrap_or(usize::MAX);
+    let mut file_count = 0usize;
+    let mut processed_files = 0usize;
+    let mut bytes_processed = 0u64;
+    let mut last_pct = 10u64;
+
+    let mut magic = read_pack_u32(reader)?;
+    if magic == 0x524f5831u32 {
+        magic = read_pack_u32(reader)?;
+    }
+    if magic == 0x524f5849u32 {
+        let index_len = read_pack_u32(reader)? as u64;
+        discard_pack_bytes(reader, index_len, &mut bytes_processed, file_count, processed_files, total_expected, progress, &mut last_pct)?;
+        magic = read_pack_u32(reader)?;
+    }
+    if magic != 0x524f5850u32 {
+        return Err(anyhow::anyhow!("Invalid pack magic: 0x{:08x}", magic));
+    }
 
-        let mut header_parsed = false;
-    let debug = std::env::var("ROX_DEBUG").is_ok();
-    if debug { eprintln!("[rox debug] unpack_stream_to_dir called (out_dir={:?})", out_dir); }
-
-        loop {
-                loop {
-                        if !header_parsed {
-                if pos + 8 > buf.len() { break; }
-                if debug {
-                    eprintln!("[rox debug] buf.len={} pos={} first16={:?}", buf.len(), pos, &buf[0..std::cmp::min(16, buf.len())]);
-                    eprintln!("[rox debug] after first debug");
-                }
-                if debug { eprintln!("[rox debug] before reading magic_header"); }
-                let magic_header = u32::from_be_bytes(buf[pos..pos+4].try_into().unwrap());
-                if debug { eprintln!("[rox debug] magic_header=0x{:08x}", magic_header); }
-                if magic_header == 0x524f5850u32 {
-                                        pos += 4;
-                                        let _file_count = u32::from_be_bytes(buf[pos..pos+4].try_into().unwrap()) as usize;
-                    pos += 4;
-                    header_parsed = true;
-                    if debug { eprintln!("[rox debug] header parsed, file_count={}", _file_count); }
-                } else if magic_header == 0x524f5831u32 {
-                                        if debug { eprintln!("[rox debug] found ROX1 outer magic, skipping 4 bytes"); }
-                    pos += 4;
-                    continue;                 } else {
-                                    }
+    file_count = read_pack_u32(reader)? as usize;
+
+    for _ in 0..file_count {
+        let name_len = read_pack_u16(reader)? as usize;
+        let mut name_bytes = vec![0u8; name_len];
+        read_pack_exact(reader, &mut name_bytes)?;
+        let name = String::from_utf8_lossy(&name_bytes).to_string();
+        let size = read_pack_u64(reader)?;
+
+        let should_write = match &files_filter {
+            Some(set) => set.contains(&name),
+            None => true,
+        };
+
+        if should_write {
+            let safe = sanitize_pack_path(&name);
+            let dest = out_dir.join(&safe);
+            if let Some(parent) = dest.parent() {
+                std::fs::create_dir_all(parent).map_err(|e| anyhow::anyhow!("Cannot create parent dir {:?}: {}", parent, e))?;
+            }
+            let file = std::fs::File::create(&dest).map_err(|e| anyhow::anyhow!("Cannot write {:?}: {}", dest, e))?;
+            let mut writer = std::io::BufWriter::with_capacity(file_buffer_capacity(size), file);
+            copy_pack_bytes(reader, &mut writer, size, &mut bytes_processed, file_count, processed_files, total_expected, progress, &mut last_pct)?;
+            finalize_output_file(writer, size, &dest)?;
+            written.push(safe.to_string_lossy().to_string());
+            if files_filter.is_some() {
+                requested = requested.saturating_sub(1);
             }
+        } else {
+            discard_pack_bytes(reader, size, &mut bytes_processed, file_count, processed_files, total_expected, progress, &mut last_pct)?;
+        }
+
+        processed_files = processed_files.saturating_add(1);
+        report_unpack_progress(progress, total_expected, bytes_processed, file_count, processed_files, &mut last_pct);
 
-                        if pos + 8 > buf.len() { break; }
-            let magic = u32::from_be_bytes(buf[pos..pos+4].try_into().unwrap());
-            if magic == 0x524f5849u32 {
-                                if pos + 8 > buf.len() { break; }
-                let index_len = u32::from_be_bytes(buf[pos+4..pos+8].try_into().unwrap()) as usize;
-                if pos + 8 + index_len > buf.len() { break; }
-                                pos += 8 + index_len;
+        if requested == 0 {
+            if let Some(cb) = progress {
+                cb(99, 100, "finishing");
             }
+            return Ok(written);
+        }
+    }
 
-                                    if pos + 2 > buf.len() { break; }
-            let name_len = u16::from_be_bytes(buf[pos..pos+2].try_into().unwrap()) as usize;
-            if pos + 2 + name_len + 8 > buf.len() { break; }
-            let name = String::from_utf8_lossy(&buf[pos+2..pos+2+name_len]).to_string();
-            let size = u64::from_be_bytes(buf[pos+2+name_len..pos+2+name_len+8].try_into().unwrap()) as usize;
-            if pos + 2 + name_len + 8 + size > buf.len() { break; }
+    if let Some(cb) = progress {
+        cb(99, 100, "finishing");
+    }
 
-            let content_start = pos + 2 + name_len + 8;
-            let content_end = content_start + size;
-            let content = &buf[content_start..content_end];
+    Ok(written)
+}
 
-                        let p = Path::new(&name);
-            let mut safe = std::path::PathBuf::new();
-            for comp in p.components() {
-                if let std::path::Component::Normal(osstr) = comp {
-                    safe.push(osstr);
-                }
-            }
-            let dest = out_dir.join(&safe);
+fn read_pack_exact<R: std::io::Read>(reader: &mut R, buf: &mut [u8]) -> Result<()> {
+    reader.read_exact(buf).map_err(|e| anyhow::anyhow!("Stream read error: {}", e))
+}
 
-            if files_filter.is_none() || files_filter.as_ref().map_or(false, |s| s.contains(&name)) {
-                if let Some(parent) = dest.parent() {
-                    std::fs::create_dir_all(parent).map_err(|e| anyhow::anyhow!("Cannot create parent dir {:?}: {}", parent, e))?;
-                }
-                std::fs::write(&dest, content).map_err(|e| anyhow::anyhow!("Cannot write {:?}: {}", dest, e))?;
-                written.push(safe.to_string_lossy().to_string());
-                if let Some(_set) = files_filter.as_ref() {
-                    requested = requested.saturating_sub(1);
-                    if requested == 0 { return Ok(written); }
-                }
-            }
+fn read_pack_u16<R: std::io::Read>(reader: &mut R) -> Result<u16> {
+    let mut buf = [0u8; 2];
+    read_pack_exact(reader, &mut buf)?;
+    Ok(u16::from_be_bytes(buf))
+}
 
-            pos = content_end;                         if pos > 0 {
-                buf.drain(0..pos);
-                pos = 0;
-            }
+fn read_pack_u32<R: std::io::Read>(reader: &mut R) -> Result<u32> {
+    let mut buf = [0u8; 4];
+    read_pack_exact(reader, &mut buf)?;
+    Ok(u32::from_be_bytes(buf))
+}
+
+fn read_pack_u64<R: std::io::Read>(reader: &mut R) -> Result<u64> {
+    let mut buf = [0u8; 8];
+    read_pack_exact(reader, &mut buf)?;
+    Ok(u64::from_be_bytes(buf))
+}
+
+fn sanitize_pack_path(name: &str) -> std::path::PathBuf {
+    let p = Path::new(name);
+    let mut safe = std::path::PathBuf::new();
+    for comp in p.components() {
+        if let std::path::Component::Normal(osstr) = comp {
+            safe.push(osstr);
         }
+    }
+    safe
+}
+
+fn file_buffer_capacity(size: u64) -> usize {
+    usize::try_from(size)
+        .unwrap_or(4 * 1024 * 1024)
+        .min(4 * 1024 * 1024)
+        .max(8192)
+}
+
+fn finalize_output_file(
+    mut writer: std::io::BufWriter<std::fs::File>,
+    size: u64,
+    dest: &Path,
+) -> Result<()> {
+    std::io::Write::flush(&mut writer).map_err(|e| anyhow::anyhow!("Cannot flush {:?}: {}", dest, e))?;
+    let file = writer.into_inner().map_err(|e| anyhow::anyhow!("Cannot finalize {:?}: {}", dest, e.error()))?;
+    crate::io_advice::sync_and_drop(&file, size);
+    Ok(())
+}
 
-                match reader.read(&mut temp) {
-            Ok(0) => break,             Ok(n) => buf.extend_from_slice(&temp[..n]),
-            Err(e) => return Err(anyhow::anyhow!("Stream read error: {}", e)),
+fn copy_pack_bytes<R: std::io::Read, W: std::io::Write>(
+    reader: &mut R,
+    writer: &mut W,
+    mut remaining: u64,
+    bytes_processed: &mut u64,
+    file_count: usize,
+    processed_files: usize,
+    total_expected: u64,
+    progress: Option<&(dyn Fn(u64, u64, &str) + Send)>,
+    last_pct: &mut u64,
+) -> Result<()> {
+    let mut buf = vec![0u8; 1024 * 1024];
+    while remaining > 0 {
+        let take = remaining.min(buf.len() as u64) as usize;
+        let read = reader.read(&mut buf[..take]).map_err(|e| anyhow::anyhow!("Stream read error: {}", e))?;
+        if read == 0 {
+            return Err(anyhow::anyhow!("Truncated pack content"));
         }
+        writer.write_all(&buf[..read]).map_err(|e| anyhow::anyhow!("Stream write error: {}", e))?;
+        remaining -= read as u64;
+        *bytes_processed = bytes_processed.saturating_add(read as u64);
+        report_unpack_progress(progress, total_expected, *bytes_processed, file_count, processed_files, last_pct);
     }
+    Ok(())
+}
 
-    Ok(written)
+fn discard_pack_bytes<R: std::io::Read>(
+    reader: &mut R,
+    mut remaining: u64,
+    bytes_processed: &mut u64,
+    file_count: usize,
+    processed_files: usize,
+    total_expected: u64,
+    progress: Option<&(dyn Fn(u64, u64, &str) + Send)>,
+    last_pct: &mut u64,
+) -> Result<()> {
+    let mut buf = vec![0u8; 1024 * 1024];
+    while remaining > 0 {
+        let take = remaining.min(buf.len() as u64) as usize;
+        let read = reader.read(&mut buf[..take]).map_err(|e| anyhow::anyhow!("Stream read error: {}", e))?;
+        if read == 0 {
+            return Err(anyhow::anyhow!("Truncated pack content"));
+        }
+        remaining -= read as u64;
+        *bytes_processed = bytes_processed.saturating_add(read as u64);
+        report_unpack_progress(progress, total_expected, *bytes_processed, file_count, processed_files, last_pct);
+    }
+    Ok(())
 }
 
 #[cfg(test)]
@@ -353,6 +467,18 @@ mod stream_tests {
     use std::io::{Write, Read};
     use std::time::{SystemTime, UNIX_EPOCH};
 
+    struct ChunkedReader<R> {
+        inner: R,
+        max_chunk: usize,
+    }
+
+    impl<R: Read> Read for ChunkedReader<R> {
+        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+            let limit = buf.len().min(self.max_chunk);
+            self.inner.read(&mut buf[..limit])
+        }
+    }
+
     #[test]
     fn test_unpack_stream_to_dir() -> Result<()> {
                 let mut parts: Vec<u8> = Vec::new();
@@ -390,7 +516,7 @@ mod stream_tests {
         let tmpdir = std::env::temp_dir().join(format!("rox_unpack_test_{}", ms));
         let _ = std::fs::create_dir_all(&tmpdir);
 
-        let out = unpack_stream_to_dir(&mut dec2, &tmpdir, None)?;
+        let out = unpack_stream_to_dir(&mut dec2, &tmpdir, None, None, 0)?;
 
                 assert_eq!(out.len(), 2);
         assert!(tmpdir.join("file1.txt").exists());
@@ -432,7 +558,7 @@ mod stream_tests {
         let tmpdir = std::env::temp_dir().join(format!("rox_unpack_png_test_{}", ms));
         let _ = std::fs::create_dir_all(&tmpdir);
 
-        let out = unpack_stream_to_dir(&mut dec, &tmpdir, None)?;
+        let out = unpack_stream_to_dir(&mut dec, &tmpdir, None, None, 0)?;
 
         assert_eq!(out.len(), 2);
         assert!(tmpdir.join("file1.txt").exists());
@@ -443,5 +569,36 @@ mod stream_tests {
         let _ = std::fs::remove_dir(&tmpdir);
         Ok(())
     }
+
+    #[test]
+    fn test_unpack_stream_to_dir_large_file_small_reads() -> Result<()> {
+        let large = vec![0x5a; 2 * 1024 * 1024];
+        let mut parts: Vec<u8> = Vec::new();
+        parts.extend_from_slice(&0x524f5850u32.to_be_bytes());
+        parts.extend_from_slice(&(1u32.to_be_bytes()));
+        let name = b"big.bin";
+        parts.extend_from_slice(&(name.len() as u16).to_be_bytes());
+        parts.extend_from_slice(name);
+        parts.extend_from_slice(&(large.len() as u64).to_be_bytes());
+        parts.extend_from_slice(&large);
+
+        let reader = std::io::Cursor::new(parts);
+        let mut reader = ChunkedReader { inner: reader, max_chunk: 37 };
+
+        let ms = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis();
+        let tmpdir = std::env::temp_dir().join(format!("rox_unpack_large_stream_test_{}", ms));
+        let _ = std::fs::create_dir_all(&tmpdir);
+
+        let out = unpack_stream_to_dir(&mut reader, &tmpdir, None, None, large.len() as u64)?;
+
+        assert_eq!(out, vec!["big.bin".to_string()]);
+        let restored = std::fs::read(tmpdir.join("big.bin"))?;
+        assert_eq!(restored.len(), large.len());
+        assert_eq!(restored, large);
+
+        let _ = std::fs::remove_file(tmpdir.join("big.bin"));
+        let _ = std::fs::remove_dir(&tmpdir);
+        Ok(())
+    }
 }
 

package/native/png_chunk_writer.rs

@@ -0,0 +1,146 @@
+use std::io::{self, Write};
+
+pub const MAX_PNG_CHUNK_DATA_LEN: usize = 64 * 1024 * 1024;
+
+pub fn write_png_chunk<W: Write>(writer: &mut W, chunk_type: &[u8; 4], data: &[u8]) -> anyhow::Result<()> {
+    let len = u32::try_from(data.len())
+        .map_err(|_| anyhow::anyhow!("chunk too large: {}", data.len()))?;
+    writer.write_all(&len.to_be_bytes())?;
+    writer.write_all(chunk_type)?;
+    writer.write_all(data)?;
+
+    let mut hasher = crc32fast::Hasher::new();
+    hasher.update(chunk_type);
+    hasher.update(data);
+    writer.write_all(&hasher.finalize().to_be_bytes())?;
+    Ok(())
+}
+
+pub fn write_chunked_idat_bytes<W: Write>(writer: &mut W, data: &[u8]) -> anyhow::Result<()> {
+    write_chunked_idat_bytes_with_limit(writer, data, MAX_PNG_CHUNK_DATA_LEN)
+}
+
+fn write_chunked_idat_bytes_with_limit<W: Write>(writer: &mut W, data: &[u8], max_chunk_len: usize) -> anyhow::Result<()> {
+    anyhow::ensure!(max_chunk_len > 0, "max_chunk_len must be > 0");
+    for chunk in data.chunks(max_chunk_len) {
+        write_png_chunk(writer, b"IDAT", chunk)?;
+    }
+    Ok(())
+}
+
+pub struct ChunkedIdatWriter<'a, W: Write> {
+    writer: &'a mut W,
+    buffer: Vec<u8>,
+    max_chunk_len: usize,
+}
+
+impl<'a, W: Write> ChunkedIdatWriter<'a, W> {
+    pub fn new(writer: &'a mut W) -> Self {
+        Self::with_max_chunk_len(writer, MAX_PNG_CHUNK_DATA_LEN)
+    }
+
+    fn with_max_chunk_len(writer: &'a mut W, max_chunk_len: usize) -> Self {
+        Self {
+            writer,
+            buffer: Vec::with_capacity(max_chunk_len.max(1).min(MAX_PNG_CHUNK_DATA_LEN)),
+            max_chunk_len: max_chunk_len.max(1),
+        }
+    }
+
+    fn flush_chunk(&mut self) -> anyhow::Result<()> {
+        if self.buffer.is_empty() {
+            return Ok(());
+        }
+        write_png_chunk(self.writer, b"IDAT", &self.buffer)?;
+        self.buffer.clear();
+        Ok(())
+    }
+
+    pub fn finish(mut self) -> anyhow::Result<()> {
+        self.flush_chunk()
+    }
+}
+
+impl<W: Write> Write for ChunkedIdatWriter<'_, W> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        let mut offset = 0;
+        while offset < buf.len() {
+            if self.buffer.len() == self.max_chunk_len {
+                self.flush_chunk().map_err(io_error)?;
+            }
+            let space = self.max_chunk_len - self.buffer.len();
+            let take = space.min(buf.len() - offset);
+            self.buffer.extend_from_slice(&buf[offset..offset + take]);
+            offset += take;
+        }
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.flush_chunk().map_err(io_error)?;
+        self.writer.flush()
+    }
+}
+
+fn io_error(err: anyhow::Error) -> io::Error {
+    io::Error::other(err.to_string())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn split_large_idat_stream_into_multiple_chunks() {
+        let mut png = Vec::new();
+        png.extend_from_slice(&[137, 80, 78, 71, 13, 10, 26, 10]);
+
+        let mut ihdr = [0u8; 13];
+        ihdr[0..4].copy_from_slice(&1u32.to_be_bytes());
+        ihdr[4..8].copy_from_slice(&1u32.to_be_bytes());
+        ihdr[8] = 8;
+        ihdr[9] = 2;
+
+        write_png_chunk(&mut png, b"IHDR", &ihdr).unwrap();
+        write_chunked_idat_bytes_with_limit(&mut png, &[1, 2, 3, 4, 5, 6, 7, 8, 9], 4).unwrap();
+        write_png_chunk(&mut png, b"IEND", &[]).unwrap();
+
+        let chunks = crate::png_utils::extract_png_chunks(&png).unwrap();
+        let idat_sizes: Vec<usize> = chunks.into_iter()
+            .filter(|chunk| chunk.name == "IDAT")
+            .map(|chunk| chunk.data.len())
+            .collect();
+
+        assert_eq!(idat_sizes, vec![4, 4, 1]);
+    }
+
+    #[test]
+    fn chunked_idat_writer_flushes_multiple_chunks() {
+        let mut png = Vec::new();
+        png.extend_from_slice(&[137, 80, 78, 71, 13, 10, 26, 10]);
+
+        let mut ihdr = [0u8; 13];
+        ihdr[0..4].copy_from_slice(&1u32.to_be_bytes());
+        ihdr[4..8].copy_from_slice(&1u32.to_be_bytes());
+        ihdr[8] = 8;
+        ihdr[9] = 2;
+
+        write_png_chunk(&mut png, b"IHDR", &ihdr).unwrap();
+        {
+            let mut writer = ChunkedIdatWriter::with_max_chunk_len(&mut png, 3);
+            writer.write_all(&[1, 2]).unwrap();
+            writer.write_all(&[3, 4, 5]).unwrap();
+            writer.write_all(&[6, 7]).unwrap();
+            writer.finish().unwrap();
+        }
+        write_png_chunk(&mut png, b"IEND", &[]).unwrap();
+
+        let chunks = crate::png_utils::extract_png_chunks(&png).unwrap();
+        let idat_sizes: Vec<usize> = chunks.into_iter()
+            .filter(|chunk| chunk.name == "IDAT")
+            .map(|chunk| chunk.data.len())
+            .collect();
+
+        assert_eq!(idat_sizes, vec![3, 3, 1]);
+    }
+}

package/native/png_utils.rs

@@ -7,6 +7,8 @@ use std::io::{Cursor, Read, Seek, SeekFrom};
 struct PngSignature([u8; 8]);
 
 const PNG_SIG: PngSignature = PngSignature([137, 80, 78, 71, 13, 10, 26, 10]);
+const HEADER_VERSION_V1: u8 = 1;
+const HEADER_VERSION_V2: u8 = 2;
 
 #[derive(Debug, Clone)]
 pub struct PngChunk {
@@ -367,21 +369,10 @@ fn extract_payload_from_embedded_nn(png_data: &[u8]) -> Result<Vec<u8>, String>
         }
         found.ok_or("PXL1 not found in reconstructed pixels")?
     };
-    let mut idx = pos + 4;
-    if idx + 2 > logical_rgb.len() { return Err("Truncated header in embedded NN".to_string()); }
-    let _version = logical_rgb[idx]; idx += 1;
-    let name_len = logical_rgb[idx] as usize; idx += 1;
-    if idx + name_len > logical_rgb.len() { return Err("Truncated name in embedded NN".to_string()); }
-    idx += name_len;
-    if idx + 4 > logical_rgb.len() { return Err("Truncated payload length in embedded NN".to_string()); }
-    let payload_len = ((logical_rgb[idx] as u32) << 24)
-        | ((logical_rgb[idx+1] as u32) << 16)
-        | ((logical_rgb[idx+2] as u32) << 8)
-        | (logical_rgb[idx+3] as u32);
-    idx += 4;
-    let end = idx + (payload_len as usize);
+    let header = parse_pixel_payload_header(&logical_rgb, pos)?;
+    let end = header.payload_offset + header.payload_len;
     if end > logical_rgb.len() { return Err("Truncated payload in embedded NN".to_string()); }
-    Ok(logical_rgb[idx..end].to_vec())
+    Ok(logical_rgb[header.payload_offset..end].to_vec())
 }
 
 pub fn extract_name_from_png(png_data: &[u8]) -> Option<String> {
@@ -420,6 +411,64 @@ fn extract_name_direct(png_data: &[u8]) -> Option<String> {
     String::from_utf8(raw[idx..idx + name_len].to_vec()).ok()
 }
 
+struct PixelPayloadHeader {
+    payload_offset: usize,
+    payload_len: usize,
+}
+
+fn parse_pixel_payload_header(buf: &[u8], pos: usize) -> Result<PixelPayloadHeader, String> {
+    let mut idx = pos + 4;
+    if idx + 2 > buf.len() {
+        return Err("Truncated header".to_string());
+    }
+
+    let version = buf[idx];
+    idx += 1;
+    let name_len = buf[idx] as usize;
+    idx += 1;
+    if idx + name_len > buf.len() {
+        return Err("Truncated name".to_string());
+    }
+    idx += name_len;
+
+    let payload_len = match version {
+        HEADER_VERSION_V1 => {
+            if idx + 4 > buf.len() {
+                return Err("Truncated payload length".to_string());
+            }
+            let len = u32::from_be_bytes([buf[idx], buf[idx + 1], buf[idx + 2], buf[idx + 3]]) as u64;
+            idx += 4;
+            len
+        }
+        HEADER_VERSION_V2 => {
+            if idx + 8 > buf.len() {
+                return Err("Truncated payload length64".to_string());
+            }
+            let len = u64::from_be_bytes([
+                buf[idx],
+                buf[idx + 1],
+                buf[idx + 2],
+                buf[idx + 3],
+                buf[idx + 4],
+                buf[idx + 5],
+                buf[idx + 6],
+                buf[idx + 7],
+            ]);
+            idx += 8;
+            len
+        }
+        other => return Err(format!("Unsupported header version {}", other)),
+    };
+
+    let payload_len = usize::try_from(payload_len)
+        .map_err(|_| "Payload too large for this platform".to_string())?;
+
+    Ok(PixelPayloadHeader {
+        payload_offset: idx,
+        payload_len,
+    })
+}
+
 fn extract_name_from_embedded_nn(png_data: &[u8]) -> Result<String, String> {
     let (pixels, width, height) = decode_to_rgba_grid(png_data)?;
     let logical_rgb = reconstruct_logical_pixels_from_nn(&pixels, width, height)?;
@@ -435,21 +484,10 @@ fn extract_name_from_embedded_nn(png_data: &[u8]) -> Result<String, String> {
 fn extract_payload_direct(png_data: &[u8]) -> Result<Vec<u8>, String> {
     let raw = decode_to_rgb(png_data)?;
     let pos = find_pixel_header(&raw)?;
-    let mut idx = pos + 4;
-    if idx + 2 > raw.len() { return Err("Truncated header".to_string()); }
-    let _version = raw[idx]; idx += 1;
-    let name_len = raw[idx] as usize; idx += 1;
-    if idx + name_len > raw.len() { return Err("Truncated name".to_string()); }
-    idx += name_len;
-    if idx + 4 > raw.len() { return Err("Truncated payload length".to_string()); }
-    let payload_len = ((raw[idx] as u32) << 24)
-        | ((raw[idx+1] as u32) << 16)
-        | ((raw[idx+2] as u32) << 8)
-        | (raw[idx+3] as u32);
-    idx += 4;
-    let end = idx + (payload_len as usize);
+    let header = parse_pixel_payload_header(&raw, pos)?;
+    let end = header.payload_offset + header.payload_len;
     if end > raw.len() { return Err("Truncated payload".to_string()); }
-    let payload = raw[idx..end].to_vec();
+    let payload = raw[header.payload_offset..end].to_vec();
     Ok(payload)
 }
 
@@ -482,19 +520,8 @@ fn extract_file_list_from_embedded_nn(png_data: &[u8]) -> Result<String, String>
     let (pixels, width, height) = decode_to_rgba_grid(png_data)?;
     let logical_rgb = reconstruct_logical_pixels_from_nn(&pixels, width, height)?;
     let pos = find_pixel_header(&logical_rgb)?;
-    let mut idx = pos + 4;
-    if idx + 2 > logical_rgb.len() { return Err("Truncated".to_string()); }
-    idx += 1;
-    let name_len = logical_rgb[idx] as usize; idx += 1;
-    if idx + name_len > logical_rgb.len() { return Err("Truncated".to_string()); }
-    idx += name_len;
-    if idx + 4 > logical_rgb.len() { return Err("Truncated".to_string()); }
-    let payload_len = ((logical_rgb[idx] as u32) << 24)
-        | ((logical_rgb[idx+1] as u32) << 16)
-        | ((logical_rgb[idx+2] as u32) << 8)
-        | (logical_rgb[idx+3] as u32);
-    idx += 4;
-    idx += payload_len as usize;
+    let header = parse_pixel_payload_header(&logical_rgb, pos)?;
+    let mut idx = header.payload_offset + header.payload_len;
     if idx + 8 > logical_rgb.len() { return Err("No file list in embedded NN".to_string()); }
     if &logical_rgb[idx..idx + 4] != b"rXFL" { return Err("No rXFL marker in embedded NN".to_string()); }
     idx += 4;
@@ -511,19 +538,8 @@ fn extract_file_list_from_embedded_nn(png_data: &[u8]) -> Result<String, String>
 fn extract_file_list_direct(png_data: &[u8]) -> Result<String, String> {
     let raw = decode_to_rgb(png_data)?;
     let pos = find_pixel_header(&raw)?;
-    let mut idx = pos + 4;
-    if idx + 2 > raw.len() { return Err("Truncated header".to_string()); }
-    idx += 1;
-    let name_len = raw[idx] as usize; idx += 1;
-    if idx + name_len > raw.len() { return Err("Truncated name".to_string()); }
-    idx += name_len;
-    if idx + 4 > raw.len() { return Err("Truncated payload length".to_string()); }
-    let payload_len = ((raw[idx] as u32) << 24)
-        | ((raw[idx+1] as u32) << 16)
-        | ((raw[idx+2] as u32) << 8)
-        | (raw[idx+3] as u32);
-    idx += 4;
-    idx += payload_len as usize;
+    let header = parse_pixel_payload_header(&raw, pos)?;
+    let mut idx = header.payload_offset + header.payload_len;
     if idx + 8 > raw.len() { return Err("No file list in pixel data".to_string()); }
     if &raw[idx..idx + 4] != b"rXFL" { return Err("No rXFL marker in pixel data".to_string()); }
     idx += 4;