Add support for decode lz4 encoded psb

src/ext/psb.rs

@@ -1,4 +1,7 @@
 //!Extensions for emote_psb crate.
+use crate::ext::io::*;
+use anyhow::Result;
+use emote_psb::PsbReader;
 use emote_psb::VirtualPsb;
 use emote_psb::header::PsbHeader;
 use emote_psb::types::collection::*;
@@ -12,6 +15,7 @@ use serde::ser::SerializeStruct;
 use serde::{Deserialize, Serialize};
 use std::cmp::PartialEq;
 use std::collections::{BTreeMap, HashMap};
+use std::io::{Read, Seek};
 use std::ops::{Index, IndexMut};
 
 const NONE: PsbValueFixed = PsbValueFixed::None;
@@ -1192,3 +1196,28 @@ impl VirtualPsbExt for VirtualPsb {
         VirtualPsbFixed::new(header, resources, extra, root.to_psb_fixed())
     }
 }
+
+/// Trait to extend PSB reader functionality.
+pub trait PsbReaderExt {
+    /// Opens a PSB v2 file from a stream, handling other formats like LZ4 compression.
+    fn open_psb_v2<T: Read + Seek>(stream: T) -> Result<VirtualPsb>;
+}
+
+const LZ4_SIGNATURE: u32 = 0x184D2204;
+
+impl PsbReaderExt for PsbReader {
+    fn open_psb_v2<T: Read + Seek>(mut stream: T) -> Result<VirtualPsb> {
+        let signature = stream.peek_u32_at(0)?;
+        if signature == LZ4_SIGNATURE {
+            let mut decoder = lz4::Decoder::new(stream)?;
+            let mut mem_stream = MemWriter::new();
+            std::io::copy(&mut decoder, &mut mem_stream)?;
+            return Self::open_psb_v2(MemReader::new(mem_stream.into_inner()));
+        }
+        let mut file = PsbReader::open_psb(stream)
+            .map_err(|e| anyhow::anyhow!("Failed to open PSB: {:?}", e))?;
+        Ok(file
+            .load()
+            .map_err(|e| anyhow::anyhow!("Failed to load PSB: {:?}", e))?)
+    }
+}

src/scripts/emote/dref.rs

@@ -69,21 +69,13 @@ impl Dpak {
     pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
         let f = std::fs::File::open(path)?;
         let mut f = std::io::BufReader::new(f);
-        let mut psb = PsbReader::open_psb(&mut f)
-            .map_err(|e| anyhow::anyhow!("Failed to read PSB from DPAK: {:?}", e))?;
-        let psb = psb
-            .load()
-            .map_err(|e| anyhow::anyhow!("Failed to load PSB from DPAK: {:?}", e))?;
+        let psb = PsbReader::open_psb_v2(&mut f)?;
         let psb = psb.to_psb_fixed();
         Ok(Self { psb })
     }
 
     pub fn load_from_data(data: &[u8]) -> Result<Self> {
-        let mut psb = PsbReader::open_psb(MemReaderRef::new(data))
-            .map_err(|e| anyhow::anyhow!("Failed to read PSB from DPAK data: {:?}", e))?;
-        let psb = psb
-            .load()
-            .map_err(|e| anyhow::anyhow!("Failed to load PSB from DPAK data: {:?}", e))?;
+        let psb = PsbReader::open_psb_v2(MemReaderRef::new(data))?;
         let psb = psb.to_psb_fixed();
         Ok(Self { psb })
     }

src/scripts/emote/pimg.rs

@@ -109,13 +109,8 @@ impl PImg {
     /// * `reader` - The reader containing the PImg script data
     /// * `filename` - The name of the file
     /// * `config` - Extra configuration options
-    pub fn new<R: Read + Seek>(reader: R, filename: &str, config: &ExtraConfig) -> Result<Self> {
-        let mut psb = PsbReader::open_psb(reader)
-            .map_err(|e| anyhow::anyhow!("Failed to open PSB from {}: {:?}", filename, e))?;
-        let psb = psb
-            .load()
-            .map_err(|e| anyhow::anyhow!("Failed to load PSB from {}: {:?}", filename, e))?
-            .to_psb_fixed();
+    pub fn new<R: Read + Seek>(reader: R, _filename: &str, config: &ExtraConfig) -> Result<Self> {
+        let psb = PsbReader::open_psb_v2(reader)?.to_psb_fixed();
         Ok(Self {
             psb,
             overlay: config.emote_pimg_overlay,

src/scripts/emote/psb.rs

@@ -67,7 +67,7 @@ impl ScriptBuilder for PsbBuilder {
     }
 
     fn extensions(&self) -> &'static [&'static str] {
-        &[]
+        &["psb"]
     }
 
     fn script_type(&self) -> &'static ScriptType {
@@ -77,6 +77,12 @@ impl ScriptBuilder for PsbBuilder {
     fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
         if buf_len >= 4 && buf.starts_with(b"PSB\0") {
             return Some(10);
+        } else if buf_len >= 4 && buf.starts_with(&[0x04, 0x22, 0x4D, 0x18]) {
+            for i in 4..buf_len - 4 {
+                if buf[i..i + 4] == *b"PSB\0" {
+                    return Some(10);
+                }
+            }
         }
         None
     }
@@ -110,12 +116,7 @@ impl Psb {
         encoding: Encoding,
         config: &ExtraConfig,
     ) -> Result<Self> {
-        let mut psb = PsbReader::open_psb(reader)
-            .map_err(|e| anyhow::anyhow!("Failed to open psb file: {:?}", e))?;
-        let psb = psb
-            .load()
-            .map_err(|e| anyhow::anyhow!("Failed to load psb: {:?}", e))?
-            .to_psb_fixed();
+        let psb = PsbReader::open_psb_v2(reader)?.to_psb_fixed();
         Ok(Self {
             psb,
             encoding,

src/scripts/kirikiri/scn.rs

@@ -142,11 +142,7 @@ impl ScnScript {
             return Self::new(MemReader::new(decoded), filename, config);
         }
         reader.rewind()?;
-        let mut psb = PsbReader::open_psb(reader)
-            .map_err(|e| anyhow::anyhow!("Failed to open PSB from {}: {:?}", filename, e))?;
-        let psb = psb
-            .load()
-            .map_err(|e| anyhow::anyhow!("Failed to load PSB from {}: {:?}", filename, e))?;
+        let psb = PsbReader::open_psb_v2(reader)?;
         Ok(Self {
             psb: psb.to_psb_fixed(),
             language_index: config.kirikiri_language_index.unwrap_or(0),