use anyhow::{Error, Result};
use log::{debug, trace, warn};
use crate::ffi;
use super::common::{copy_whole_match, get_crc};
use super::header::Header;
use super::quantum_header::QuantumHeader;
use super::DecoderType;
pub struct Decoder {
pub input_read: u32,
pub output_written: u32,
scratch: [u8; 0x6C000],
header: Header,
quantum_header: QuantumHeader,
}
impl Decoder {
pub fn parse_header(&mut self, input: &[u8]) -> Result<usize> {
self.header.parse(input)
}
pub fn parse_quantum_header(&mut self, input: &[u8]) -> Result<usize> {
self.quantum_header.parse(input, self.header.use_checksums)
}
pub fn parse_lzna_quantum_header(&mut self, input: &[u8], output_len: usize) -> Result<usize> {
self.quantum_header
.parse_lzna(input, self.header.use_checksums, output_len)
}
pub unsafe fn decode_step(
&mut self,
output: &mut [u8],
output_offset: usize,
input: &[u8],
) -> Result<()> {
let mut input_offset: usize = 0;
if output_offset < 0x40000 {
trace!("Parsing header");
input_offset += self.parse_header(input)?;
}
let is_kraken_decoder = self.header.decoder_type == DecoderType::Kraken
|| self.header.decoder_type == DecoderType::Mermaid
|| self.header.decoder_type == DecoderType::Leviathan;
let output_len = std::cmp::min(
output.len() - output_offset,
if is_kraken_decoder { 0x40000 } else { 0x4000 },
);
trace!("Checking if data is compressed");
if self.header.uncompressed {
if input.len() - input_offset < output_len {
self.input_read = 0;
self.output_written = 0;
return Ok(());
}
output[output_offset..output_offset + output_len]
.copy_from_slice(&input[input_offset..input_offset + output_len]);
self.input_read = (input_offset + output_len) as u32;
self.output_written = output_len as u32;
return Ok(());
}
if is_kraken_decoder {
trace!("Parsing quantum header");
input_offset += self.parse_quantum_header(&input[input_offset..])?;
}
else {
trace!("Parsing LZNA quantum header");
input_offset += self.parse_lzna_quantum_header(&input[input_offset..], output_len)?;
}
trace!("Checking if input buffer has enough data to decompress");
if input.len() - input_offset < self.quantum_header.compressed_size as usize {
self.input_read = 0;
self.output_written = 0;
return Ok(());
}
trace!("Checking if output buffer is large enough for decompressed data");
if output_len < self.quantum_header.compressed_size as usize {
return Err(Error::msg("Output buffer is too small"));
}
trace!("Checking if compressed size is 0");
if self.quantum_header.compressed_size == 0 {
if self.quantum_header.whole_match_distance != 0 {
if self.quantum_header.whole_match_distance > output_offset as u32 {
return Err(Error::msg("Whole match distance is too large"));
}
copy_whole_match(
output,
output_offset,
self.quantum_header.whole_match_distance as usize,
);
} else {
output[output_offset..output_offset + output_len]
.fill(self.quantum_header.checksum as u8);
}
self.input_read = input_offset as u32;
self.output_written = output_len as u32;
return Ok(());
}
trace!("Getting CRC checksum of compressed data");
let crc = get_crc(
&input[input_offset..input_offset + self.quantum_header.compressed_size as usize],
);
trace!("Checking if checksum is correct");
if self.header.use_checksums && (crc & 0xFFFFFF) != self.quantum_header.checksum {
return Err(Error::msg("Checksum mismatch"));
}
trace!("Checking if compressed size is the same as output length");
if self.quantum_header.compressed_size == output_len as u32 {
output[output_offset..output_offset + output_len]
.copy_from_slice(&input[input_offset..input_offset + output_len]);
self.input_read = (input_offset + output_len) as u32;
self.output_written = output_len as u32;
return Ok(());
}
trace!("Decompressing data using correct decoder");
debug!(
"Decompressing {} bytes to {} bytes using {:?}",
self.quantum_header.compressed_size, output_len, self.header.decoder_type
);
debug!(
"Input offset: {} Output offset: {}",
input_offset, output_offset
);
let read_bytes: i32 = match self.header.decoder_type {
DecoderType::Kraken => ffi::Kraken_DecodeQuantum(
output.as_mut_ptr().add(output_offset),
output.as_mut_ptr().add(output_offset + output_len),
output.as_mut_ptr(),
input.as_ptr().add(input_offset),
input
.as_ptr()
.add(input_offset + self.quantum_header.compressed_size as usize),
self.scratch.as_mut_ptr(),
self.scratch.as_mut_ptr().add(self.scratch.len()),
),
DecoderType::Mermaid => ffi::Mermaid_DecodeQuantum(
output.as_mut_ptr().add(output_offset),
output.as_mut_ptr().add(output_offset + output_len),
output.as_mut_ptr(),
input.as_ptr().add(input_offset),
input
.as_ptr()
.add(input_offset + self.quantum_header.compressed_size as usize),
self.scratch.as_mut_ptr(),
self.scratch.as_mut_ptr().add(self.scratch.len()),
),
DecoderType::Leviathan => ffi::Leviathan_DecodeQuantum(
output.as_mut_ptr().add(output_offset),
output.as_mut_ptr().add(output_offset + output_len),
output.as_mut_ptr(),
input.as_ptr().add(input_offset),
input
.as_ptr()
.add(input_offset + self.quantum_header.compressed_size as usize),
self.scratch.as_mut_ptr(),
self.scratch.as_mut_ptr().add(self.scratch.len()),
),
DecoderType::LZNA => {
if self.header.restart_decoder {
self.header.restart_decoder = false;
ffi::LZNA_InitLookup(self.scratch.as_mut_ptr() as *mut ffi::LznaState);
}
ffi::LZNA_DecodeQuantum(
output.as_mut_ptr().add(output_offset),
output.as_mut_ptr().add(output_offset + output_len),
output.as_mut_ptr(),
input.as_ptr().add(input_offset),
input
.as_ptr()
.add(input_offset + self.quantum_header.compressed_size as usize),
self.scratch.as_mut_ptr() as *mut ffi::LznaState,
)
}
DecoderType::Bitknit => {
if self.header.restart_decoder {
self.header.restart_decoder = false;
ffi::BitknitState_Init(self.scratch.as_mut_ptr() as *mut ffi::BitknitState);
}
ffi::Bitknit_Decode(
input.as_ptr().add(input_offset),
input
.as_ptr()
.add(input_offset + self.quantum_header.compressed_size as usize),
output.as_mut_ptr().add(output_offset),
output.as_mut_ptr().add(output_offset + output_len),
output.as_mut_ptr(),
self.scratch.as_mut_ptr() as *mut ffi::BitknitState,
) as i32
}
_ => return Err(Error::msg("Unknown decoder type")),
};
trace!("Checking if decompressed size is correct");
if read_bytes != self.quantum_header.compressed_size as i32 {
let message = format!(
"Decompressed size mismatch: {} != {}",
read_bytes, self.quantum_header.compressed_size
);
warn!("{}", message);
return Err(Error::msg(message));
}
self.input_read = input_offset as u32 + read_bytes as u32;
self.output_written = output_len as u32;
Ok(())
}
}
impl Default for Decoder {
fn default() -> Self {
Self {
input_read: 0,
output_written: 0,
scratch: [0; 0x6C000],
header: Header::default(),
quantum_header: QuantumHeader::default(),
}
}
}