From 5af17cee87552c4115beda96d5cc3cd9e89a703f Mon Sep 17 00:00:00 2001 From: lifegpc Date: Thu, 25 Dec 2025 15:40:17 +0800 Subject: [PATCH] Add support for decode lz4 encoded psb --- Cargo.toml | 1 + check_features.py | 2 +- src/ext/psb.rs | 29 +++++++++++++++++++++++++++++ src/scripts/emote/dref.rs | 12 ++---------- src/scripts/emote/pimg.rs | 9 ++------- src/scripts/emote/psb.rs | 15 ++++++++------- src/scripts/kirikiri/scn.rs | 6 +----- 7 files changed, 44 insertions(+), 30 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index f576fd8..e37206c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -109,6 +109,7 @@ lossless-audio = ["utils-pcm"] audio-flac = ["libflac-sys", "utils-pcm"] unstable = ["msg_tool_macro/unstable"] jieba = ["jieba-rs"] +emote-psb = ["dep:emote-psb", "lz4"] # utils feature utils-bit-stream = [] utils-blowfish = ["byteorder"] diff --git a/check_features.py b/check_features.py index 2792f17..f6800d2 100644 --- a/check_features.py +++ b/check_features.py @@ -41,7 +41,7 @@ def main(): features = cargo_toml.get("features", {}) feature_names = list(features.keys()) - feature_names = [name for name in feature_names if filter_name(name)] + feature_names = [name for name in feature_names if filter_name(name) and f"dep:{name}" not in features[name]] if not feature_names: print("No features defined in Cargo.toml.") diff --git a/src/ext/psb.rs b/src/ext/psb.rs index 82440f4..c297337 100644 --- a/src/ext/psb.rs +++ b/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(stream: T) -> Result; +} + +const LZ4_SIGNATURE: u32 = 0x184D2204; + +impl PsbReaderExt for PsbReader { + fn open_psb_v2(mut stream: T) -> Result { + 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))?) + } +} diff --git a/src/scripts/emote/dref.rs b/src/scripts/emote/dref.rs index e64f34a..29f12e7 100644 --- a/src/scripts/emote/dref.rs +++ b/src/scripts/emote/dref.rs @@ -69,21 +69,13 @@ impl Dpak { pub fn new>(path: P) -> Result { 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 { - 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 }) } diff --git a/src/scripts/emote/pimg.rs b/src/scripts/emote/pimg.rs index 5f75ae4..672dc9c 100644 --- a/src/scripts/emote/pimg.rs +++ b/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(reader: R, filename: &str, config: &ExtraConfig) -> Result { - 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(reader: R, _filename: &str, config: &ExtraConfig) -> Result { + let psb = PsbReader::open_psb_v2(reader)?.to_psb_fixed(); Ok(Self { psb, overlay: config.emote_pimg_overlay, diff --git a/src/scripts/emote/psb.rs b/src/scripts/emote/psb.rs index fcdf918..ae2a447 100644 --- a/src/scripts/emote/psb.rs +++ b/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 { 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 { - 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, diff --git a/src/scripts/kirikiri/scn.rs b/src/scripts/kirikiri/scn.rs index ee9b0b5..cc6ba69 100644 --- a/src/scripts/kirikiri/scn.rs +++ b/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),