fix: обработка выхода за пределы индекса сортировки в архиве и улучшение декодирования LZSS с поддержкой XOR

2 files changed, 63 insertions, 19 deletions

diff --git a/crates/rsli/src/compress/lzh.rs b/crates/rsli/src/compress/lzh.rs
index 93f2267..fa9cff7 100644
--- a/crates/rsli/src/compress/lzh.rs
+++ b/crates/rsli/src/compress/lzh.rs
@@ -1,4 +1,4 @@
-use super::xor::xor_stream;
+use super::xor::XorState;
 use crate::error::Error;
 use crate::Result;
 
@@ -10,22 +10,14 @@ pub(crate) const LZH_T: usize = LZH_N_CHAR * 2 - 1;
 pub(crate) const LZH_R: usize = LZH_T - 1;
 pub(crate) const LZH_MAX_FREQ: u16 = 0x8000;
 
-/// LZSS-Huffman decompression with optional XOR pre-decryption
+/// LZSS-Huffman decompression with optional on-the-fly XOR decryption.
 pub fn lzss_huffman_decompress(
     data: &[u8],
     expected_size: usize,
     xor_key: Option<u16>,
 ) -> Result<Vec<u8>> {
-    // TODO: Full optimization for Huffman variant (rare in practice)
-    // For now, fallback to separate XOR step for Huffman
-    if let Some(key) = xor_key {
-        let decrypted = xor_stream(data, key);
-        let mut decoder = LzhDecoder::new(&decrypted);
-        decoder.decode(expected_size)
-    } else {
-        let mut decoder = LzhDecoder::new(data);
-        decoder.decode(expected_size)
-    }
+    let mut decoder = LzhDecoder::new(data, xor_key);
+    decoder.decode(expected_size)
 }
 
 struct LzhDecoder<'a> {
@@ -40,9 +32,9 @@ struct LzhDecoder<'a> {
 }
 
 impl<'a> LzhDecoder<'a> {
-    fn new(data: &'a [u8]) -> Self {
+    fn new(data: &'a [u8], xor_key: Option<u16>) -> Self {
         let mut decoder = Self {
-            bit_reader: BitReader::new(data),
+            bit_reader: BitReader::new(data, xor_key),
             text: [0x20u8; LZH_N],
             freq: [0u16; LZH_T + 1],
             parent: [0usize; LZH_T + LZH_N_CHAR],
@@ -257,23 +249,37 @@ struct BitReader<'a> {
     data: &'a [u8],
     byte_pos: usize,
     bit_mask: u8,
+    current_byte: u8,
+    xor_state: Option<XorState>,
 }
 
 impl<'a> BitReader<'a> {
-    fn new(data: &'a [u8]) -> Self {
+    fn new(data: &'a [u8], xor_key: Option<u16>) -> Self {
         Self {
             data,
             byte_pos: 0,
             bit_mask: 0x80,
+            current_byte: 0,
+            xor_state: xor_key.map(XorState::new),
         }
     }
 
     fn read_bit_or_zero(&mut self) -> u8 {
-        let Some(byte) = self.data.get(self.byte_pos).copied() else {
-            return 0;
-        };
+        if self.bit_mask == 0x80 {
+            let Some(mut byte) = self.data.get(self.byte_pos).copied() else {
+                return 0;
+            };
+            if let Some(state) = &mut self.xor_state {
+                byte = state.decrypt_byte(byte);
+            }
+            self.current_byte = byte;
+        }
 
-        let bit = if (byte & self.bit_mask) != 0 { 1 } else { 0 };
+        let bit = if (self.current_byte & self.bit_mask) != 0 {
+            1
+        } else {
+            0
+        };
         self.bit_mask >>= 1;
         if self.bit_mask == 0 {
             self.bit_mask = 0x80;
diff --git a/crates/rsli/src/tests.rs b/crates/rsli/src/tests.rs
index 7ed16b1..94d14af 100644
--- a/crates/rsli/src/tests.rs
+++ b/crates/rsli/src/tests.rs
@@ -668,6 +668,44 @@ fn rsli_synthetic_all_methods_roundtrip() {
 }
 
 #[test]
+fn rsli_xorlzss_huffman_on_the_fly_roundtrip() {
+    let plain: Vec<u8> = (0..512u16).map(|i| b'A' + (i % 26) as u8).collect();
+    let entries = vec![SyntheticRsliEntry {
+        name: "XLZH_ONFLY".to_string(),
+        method_raw: 0x0A0,
+        plain: plain.clone(),
+        declared_packed_size: None,
+    }];
+
+    let bytes = build_rsli_bytes(
+        &entries,
+        &RsliBuildOptions {
+            seed: 0x0BAD_C0DE,
+            presorted: true,
+            overlay: 0,
+            add_ao_trailer: false,
+        },
+    );
+    let path = write_temp_file("rsli-xorlzh-onfly", &bytes);
+
+    let library = Library::open_path(&path).expect("open synthetic XLZH archive failed");
+    let id = library
+        .find("XLZH_ONFLY")
+        .expect("find XLZH_ONFLY entry failed");
+
+    let loaded = library.load(id).expect("load XLZH_ONFLY failed");
+    assert_eq!(loaded, plain);
+
+    let packed = library
+        .load_packed(id)
+        .expect("load_packed XLZH_ONFLY failed");
+    let unpacked = library.unpack(&packed).expect("unpack XLZH_ONFLY failed");
+    assert_eq!(unpacked, loaded);
+
+    let _ = fs::remove_file(&path);
+}
+
+#[test]
 fn rsli_synthetic_overlay_and_ao_trailer() {
     let entries = vec![SyntheticRsliEntry {
         name: "OVERLAY".to_string(),