roxify 1.13.2 → 1.13.4

package/native/streaming_decode.rs

@@ -4,17 +4,39 @@ use cipher::{KeyIvInit, StreamCipher};
 
 const PIXEL_MAGIC: &[u8] = b"PXL1";
 const MARKER_BYTES: usize = 12;
+const PACK_MAGIC: [u8; 4] = 0x524f5850u32.to_be_bytes();
 
 type Aes256Ctr = ctr::Ctr64BE<aes::Aes256>;
 
+pub type DecodeProgressCallback = Box<dyn Fn(u64, u64, &str) + Send>;
+
 pub fn streaming_decode_to_dir(png_path: &Path, out_dir: &Path) -> Result<Vec<String>, String> {
-    streaming_decode_to_dir_encrypted(png_path, out_dir, None)
+    streaming_decode_selected_to_dir_encrypted_with_progress(png_path, out_dir, None, None, None)
 }
 
 pub fn streaming_decode_to_dir_encrypted(
     png_path: &Path,
     out_dir: &Path,
     passphrase: Option<&str>,
+) -> Result<Vec<String>, String> {
+    streaming_decode_selected_to_dir_encrypted_with_progress(png_path, out_dir, None, passphrase, None)
+}
+
+pub fn streaming_decode_to_dir_encrypted_with_progress(
+    png_path: &Path,
+    out_dir: &Path,
+    passphrase: Option<&str>,
+    progress: Option<DecodeProgressCallback>,
+) -> Result<Vec<String>, String> {
+    streaming_decode_selected_to_dir_encrypted_with_progress(png_path, out_dir, None, passphrase, progress)
+}
+
+pub fn streaming_decode_selected_to_dir_encrypted_with_progress(
+    png_path: &Path,
+    out_dir: &Path,
+    files_opt: Option<&[String]>,
+    passphrase: Option<&str>,
+    progress: Option<DecodeProgressCallback>,
 ) -> Result<Vec<String>, String> {
     let file = std::fs::File::open(png_path).map_err(|e| format!("open: {}", e))?;
     let mmap = unsafe { memmap2::Mmap::map(&file).map_err(|e| format!("mmap: {}", e))? };
@@ -24,7 +46,16 @@ pub fn streaming_decode_to_dir_encrypted(
         return Err("Not a PNG file".into());
     }
 
+    if let Some(ref cb) = progress {
+        cb(2, 100, "parsing_png");
+    }
+
     let (width, height, idat_data_start, idat_data_end) = parse_png_header(data)?;
+    let total_expected = parse_rxfl_total_bytes(data).unwrap_or(0);
+
+    if let Some(ref cb) = progress {
+        cb(5, 100, "reading_header");
+    }
 
     let mut reader = DeflatePixelReader::new(data, width, height, idat_data_start, idat_data_end);
 
@@ -51,6 +82,10 @@ pub fn streaming_decode_to_dir_encrypted(
     reader.read_exact(&mut plen_buf).map_err(|e| format!("read payload_len: {}", e))?;
     let payload_len = u32::from_be_bytes(plen_buf) as u64;
 
+    if let Some(ref cb) = progress {
+        cb(8, 100, "decrypting");
+    }
+
     let payload_reader = reader.take(payload_len);
 
     let first_byte_reader = FirstByteReader::new(payload_reader);
@@ -58,10 +93,13 @@ pub fn streaming_decode_to_dir_encrypted(
 
     match enc_byte {
         0x00 => {
+            if let Some(ref cb) = progress {
+                cb(10, 100, "decompressing");
+            }
             let mut decoder = zstd::stream::Decoder::new(remaining_reader)
                 .map_err(|e| format!("zstd decoder: {}", e))?;
             decoder.window_log_max(31).map_err(|e| format!("zstd window_log_max: {}", e))?;
-            read_rox1_and_untar(decoder, out_dir)
+            read_rox1_and_unpack_with_progress(decoder, out_dir, files_opt, progress, total_expected)
         }
         0x03 => {
             let pass = passphrase.ok_or("Passphrase required for AES-CTR decryption")?;
@@ -79,23 +117,42 @@ pub fn streaming_decode_to_dir_encrypted(
             let encrypted_data_len = payload_len - 1 - 16 - 16 - hmac_size;
             let ctr_reader = CtrDecryptReader::new(r.take(encrypted_data_len), cipher);
 
+            if let Some(ref cb) = progress {
+                cb(10, 100, "decompressing");
+            }
             let mut decoder = zstd::stream::Decoder::new(ctr_reader)
                 .map_err(|e| format!("zstd decoder: {}", e))?;
             decoder.window_log_max(31).map_err(|e| format!("zstd window_log_max: {}", e))?;
-            read_rox1_and_untar(decoder, out_dir)
+            read_rox1_and_unpack_with_progress(decoder, out_dir, files_opt, progress, total_expected)
         }
         _ => Err(format!("Unsupported encryption (enc=0x{:02x}) in streaming decode", enc_byte)),
     }
 }
 
-fn read_rox1_and_untar<R: Read>(mut decoder: R, out_dir: &Path) -> Result<Vec<String>, String> {
+fn read_rox1_and_unpack_with_progress<R: Read>(
+    mut decoder: R,
+    out_dir: &Path,
+    files_opt: Option<&[String]>,
+    progress: Option<DecodeProgressCallback>,
+    total_expected: u64,
+) -> Result<Vec<String>, String> {
     let mut magic = [0u8; 4];
     decoder.read_exact(&mut magic).map_err(|e| format!("read ROX1: {}", e))?;
     if &magic != b"ROX1" {
         return Err(format!("Expected ROX1, got {:?}", magic));
     }
     std::fs::create_dir_all(out_dir).map_err(|e| format!("mkdir: {}", e))?;
-    tar_unpack_from_reader(decoder, out_dir)
+
+    let mut prefix = [0u8; 4];
+    decoder.read_exact(&mut prefix).map_err(|e| format!("read payload magic: {}", e))?;
+    let mut chained = std::io::Cursor::new(prefix).chain(decoder);
+
+    if prefix == PACK_MAGIC {
+        crate::packer::unpack_stream_to_dir(&mut chained, out_dir, files_opt, progress.as_deref(), total_expected)
+            .map_err(|e| format!("pack unpack: {}", e))
+    } else {
+        tar_unpack_from_reader_with_progress(chained, out_dir, files_opt, progress, total_expected)
+    }
 }
 
 fn parse_png_header(data: &[u8]) -> Result<(usize, usize, usize, usize), String> {
@@ -147,6 +204,30 @@ fn parse_png_header(data: &[u8]) -> Result<(usize, usize, usize, usize), String>
     Ok((width, height, idat_start, idat_end))
 }
 
+fn parse_rxfl_total_bytes(data: &[u8]) -> Option<u64> {
+    let mut pos = 8;
+    while pos + 12 <= data.len() {
+        let chunk_len = u32::from_be_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]) as usize;
+        let chunk_type = &data[pos + 4..pos + 8];
+        let chunk_data_start = pos + 8;
+
+        if chunk_type == b"rXFL" && chunk_data_start + chunk_len <= data.len() {
+            let json_bytes = &data[chunk_data_start..chunk_data_start + chunk_len];
+            if let Ok(entries) = serde_json::from_slice::<Vec<serde_json::Value>>(json_bytes) {
+                let total: u64 = entries.iter()
+                    .filter_map(|e| e.get("size").and_then(|s| s.as_u64()))
+                    .sum();
+                return Some(total);
+            }
+        } else if chunk_type == b"IEND" {
+            break;
+        }
+
+        pos = chunk_data_start + chunk_len + 4;
+    }
+    None
+}
+
 struct DeflatePixelReader<'a> {
     data: &'a [u8],
     height: usize,
@@ -298,16 +379,29 @@ impl<R: Read> Read for CtrDecryptReader<R> {
     }
 }
 
-fn tar_unpack_from_reader<R: Read>(reader: R, output_dir: &Path) -> Result<Vec<String>, String> {
+fn tar_unpack_from_reader_with_progress<R: Read>(
+    reader: R,
+    output_dir: &Path,
+    files_opt: Option<&[String]>,
+    progress: Option<DecodeProgressCallback>,
+    total_expected: u64,
+) -> Result<Vec<String>, String> {
     let buf_reader = std::io::BufReader::with_capacity(8 * 1024 * 1024, reader);
     let mut archive = tar::Archive::new(buf_reader);
     let mut written = Vec::new();
     let mut created_dirs = std::collections::HashSet::new();
+    let mut bytes_extracted: u64 = 0;
+    let mut last_pct: u64 = 10;
+    let files_filter: Option<std::collections::HashSet<&str>> = files_opt.map(|files| files.iter().map(|file| file.as_str()).collect());
+    let mut remaining = files_filter.as_ref().map(|files| files.len()).unwrap_or(usize::MAX);
 
     let entries = archive.entries().map_err(|e| format!("tar entries: {}", e))?;
     for entry in entries {
         let mut entry = entry.map_err(|e| format!("tar entry: {}", e))?;
+        let entry_size = entry.size();
         let path = entry.path().map_err(|e| format!("tar path: {}", e))?.to_path_buf();
+        let logical_path = path.to_string_lossy().replace('\\', "/");
+        let should_write = files_filter.as_ref().map(|files| files.contains(logical_path.as_str())).unwrap_or(true);
 
         let mut safe = std::path::PathBuf::new();
         for comp in path.components() {
@@ -319,6 +413,23 @@ fn tar_unpack_from_reader<R: Read>(reader: R, output_dir: &Path) -> Result<Vec<S
             continue;
         }
 
+        if !should_write {
+            std::io::copy(&mut entry, &mut std::io::sink()).map_err(|e| format!("skip {:?}: {}", safe, e))?;
+            bytes_extracted += entry_size;
+            if let Some(ref cb) = progress {
+                let pct = if total_expected > 0 {
+                    10 + (bytes_extracted * 89 / total_expected).min(89)
+                } else {
+                    (10 + (bytes_extracted / (1024 * 1024))).min(99)
+                };
+                if pct > last_pct {
+                    last_pct = pct;
+                    cb(pct, 100, "extracting");
+                }
+            }
+            continue;
+        }
+
         let dest = output_dir.join(&safe);
         if let Some(parent) = dest.parent() {
             if created_dirs.insert(parent.to_path_buf()) {
@@ -327,11 +438,34 @@ fn tar_unpack_from_reader<R: Read>(reader: R, output_dir: &Path) -> Result<Vec<S
         }
 
         let mut f = std::io::BufWriter::with_capacity(
-            (entry.size() as usize).min(4 * 1024 * 1024).max(8192),
+            (entry_size as usize).min(4 * 1024 * 1024).max(8192),
             std::fs::File::create(&dest).map_err(|e| format!("create {:?}: {}", dest, e))?,
         );
         std::io::copy(&mut entry, &mut f).map_err(|e| format!("write {:?}: {}", dest, e))?;
         written.push(safe.to_string_lossy().to_string());
+        if files_filter.is_some() {
+            remaining = remaining.saturating_sub(1);
+        }
+
+        bytes_extracted += entry_size;
+        if let Some(ref cb) = progress {
+            let pct = if total_expected > 0 {
+                10 + (bytes_extracted * 89 / total_expected).min(89)
+            } else {
+                (10 + (bytes_extracted / (1024 * 1024))).min(99)
+            };
+            if pct > last_pct {
+                last_pct = pct;
+                cb(pct, 100, "extracting");
+            }
+        }
+        if remaining == 0 {
+            break;
+        }
+    }
+
+    if let Some(ref cb) = progress {
+        cb(99, 100, "finishing");
     }
 
     Ok(written)