use crate::{
models::{Gamestate, HeaderBorrowed, HeaderOwned},
tokens::TokenLookup,
Ck3Error, Ck3ErrorKind, FailedResolveStrategy,
};
use jomini::{BinaryDeserializer, TextDeserializer};
use serde::de::{Deserialize, DeserializeOwned};
use std::io::{Read, Seek};
pub(crate) const HEADER_LEN_UPPER_BOUND: usize = 0x10000;
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Encoding {
TextZip,
BinaryZip,
Binary,
}
#[derive(Debug, Clone, Copy)]
pub enum Extraction {
InMemory,
#[cfg(feature = "mmap")]
MmapTemporaries,
}
#[derive(Debug, Clone)]
pub struct Ck3ExtractorBuilder {
extraction: Extraction,
on_failed_resolve: FailedResolveStrategy,
}
impl Default for Ck3ExtractorBuilder {
fn default() -> Self {
Ck3ExtractorBuilder::new()
}
}
impl Ck3ExtractorBuilder {
pub fn new() -> Self {
Ck3ExtractorBuilder {
extraction: Extraction::InMemory,
on_failed_resolve: FailedResolveStrategy::Ignore,
}
}
pub fn with_extraction(mut self, extraction: Extraction) -> Self {
self.extraction = extraction;
self
}
pub fn with_on_failed_resolve(mut self, strategy: FailedResolveStrategy) -> Self {
self.on_failed_resolve = strategy;
self
}
pub fn extract_header_owned(&self, data: &[u8]) -> Result<(HeaderOwned, Encoding), Ck3Error> {
self.extract_header_as(data)
}
pub fn extract_header_borrowed<'a>(
&self,
data: &'a [u8],
) -> Result<(HeaderBorrowed<'a>, Encoding), Ck3Error> {
self.extract_header_as(data)
}
pub fn extract_header_as<'de, T>(&self, data: &'de [u8]) -> Result<(T, Encoding), Ck3Error>
where
T: Deserialize<'de>,
{
let data = skip_save_prefix(&data);
let data = &data[..std::cmp::min(data.len(), HEADER_LEN_UPPER_BOUND)];
let (header, rest) = split_on_zip(data);
if sniff_is_binary(header) {
let res = BinaryDeserializer::ck3_builder()
.on_failed_resolve(self.on_failed_resolve)
.from_slice(header, &TokenLookup)?;
if rest.is_empty() {
Ok((res, Encoding::Binary))
} else {
Ok((res, Encoding::BinaryZip))
}
} else {
let res = TextDeserializer::from_utf8_slice(header)?;
Ok((res, Encoding::TextZip))
}
}
pub fn extract_save<R>(&self, reader: R) -> Result<(Gamestate, Encoding), Ck3Error>
where
R: Read + Seek,
{
self.extract_save_as(reader)
}
pub fn extract_save_as<T, R>(&self, mut reader: R) -> Result<(T, Encoding), Ck3Error>
where
R: Read + Seek,
T: DeserializeOwned,
{
let mut buffer = vec![0; HEADER_LEN_UPPER_BOUND];
read_upto(&mut reader, &mut buffer)?;
if zip_index(&buffer).is_some() {
let mut zip =
zip::ZipArchive::new(&mut reader).map_err(Ck3ErrorKind::ZipCentralDirectory)?;
match self.extraction {
Extraction::InMemory => {
melt_in_memory(&mut buffer, "gamestate", &mut zip, self.on_failed_resolve)
}
#[cfg(feature = "mmap")]
Extraction::MmapTemporaries => {
melt_with_temporary("gamestate", &mut zip, self.on_failed_resolve)
}
}
} else {
reader.read_to_end(&mut buffer)?;
let data = skip_save_prefix(&buffer);
let res = BinaryDeserializer::ck3_builder()
.on_failed_resolve(self.on_failed_resolve)
.from_slice(data, &TokenLookup)?;
Ok((res, Encoding::Binary))
}
}
}
#[derive(Debug, Clone)]
pub struct Ck3Extractor {}
impl Ck3Extractor {
pub fn builder() -> Ck3ExtractorBuilder {
Ck3ExtractorBuilder::new()
}
pub fn extract_header(data: &[u8]) -> Result<(HeaderOwned, Encoding), Ck3Error> {
Self::builder().extract_header_owned(data)
}
pub fn extract_save<R>(reader: R) -> Result<(Gamestate, Encoding), Ck3Error>
where
R: Read + Seek,
{
Self::builder().extract_save(reader)
}
}
fn melt_in_memory<T, R>(
mut buffer: &mut Vec<u8>,
name: &'static str,
zip: &mut zip::ZipArchive<R>,
on_failed_resolve: FailedResolveStrategy,
) -> Result<(T, Encoding), Ck3Error>
where
R: Read + Seek,
T: DeserializeOwned,
{
buffer.clear();
let mut zip_file = zip
.by_name(name)
.map_err(|e| Ck3ErrorKind::ZipMissingEntry(name, e))?;
if zip_file.size() > 1024 * 1024 * 200 {
return Err(Ck3ErrorKind::ZipSize(name).into());
}
buffer.reserve(zip_file.size() as usize);
zip_file
.read_to_end(&mut buffer)
.map_err(|e| Ck3ErrorKind::ZipExtraction(name, e))?;
if sniff_is_binary(&buffer) {
let res = BinaryDeserializer::ck3_builder()
.on_failed_resolve(on_failed_resolve)
.from_slice(&buffer, &TokenLookup)
.map_err(|e| Ck3ErrorKind::Deserialize {
part: Some(name.to_string()),
err: e,
})?;
Ok((res, Encoding::BinaryZip))
} else {
let res = TextDeserializer::from_utf8_slice(&buffer)?;
Ok((res, Encoding::TextZip))
}
}
#[cfg(feature = "mmap")]
fn melt_with_temporary<T, R>(
name: &'static str,
zip: &mut zip::ZipArchive<R>,
on_failed_resolve: FailedResolveStrategy,
) -> Result<(T, Encoding), Ck3Error>
where
R: Read + Seek,
T: DeserializeOwned,
{
let mut zip_file = zip
.by_name(name)
.map_err(|e| Ck3ErrorKind::ZipMissingEntry(name, e))?;
if zip_file.size() > 1024 * 1024 * 200 {
return Err(Ck3ErrorKind::ZipSize(name).into());
}
let mut mmap = memmap::MmapMut::map_anon(zip_file.size() as usize)?;
std::io::copy(&mut zip_file, &mut mmap.as_mut())
.map_err(|e| Ck3ErrorKind::ZipExtraction(name, e))?;
let buffer = &mmap[..];
if sniff_is_binary(buffer) {
let res = BinaryDeserializer::ck3_builder()
.on_failed_resolve(on_failed_resolve)
.from_slice(&buffer, &TokenLookup)
.map_err(|e| Ck3ErrorKind::Deserialize {
part: Some(name.to_string()),
err: e,
})?;
Ok((res, Encoding::BinaryZip))
} else {
let res = TextDeserializer::from_utf8_slice(&buffer)?;
Ok((res, Encoding::TextZip))
}
}
fn skip_save_prefix(data: &[u8]) -> &[u8] {
let id_line_idx = data
.iter()
.position(|&x| x == b'\n')
.map(|x| x + 1)
.unwrap_or(0);
&data[id_line_idx..]
}
fn sniff_is_binary(data: &[u8]) -> bool {
data.get(2..4).map_or(false, |x| x == [0x01, 0x00])
}
pub(crate) fn zip_index(data: &[u8]) -> Option<usize> {
twoway::find_bytes(data, &[0x50, 0x4b, 0x03, 0x04])
}
fn split_on_zip(data: &[u8]) -> (&[u8], &[u8]) {
if let Some(idx) = zip_index(data) {
data.split_at(idx)
} else {
data.split_at(data.len())
}
}
fn read_upto<R>(reader: &mut R, mut buf: &mut [u8]) -> Result<(), std::io::Error>
where
R: std::io::Read,
{
while !buf.is_empty() {
match reader.read(buf) {
Ok(0) => break,
Ok(n) => {
let tmp = buf;
buf = &mut tmp[n..];
}
Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => {}
Err(e) => return Err(e),
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_skip_save_prefix() {
let data = b"abc\n123";
let result = skip_save_prefix(&data[..]);
assert_eq!(result, b"123");
}
}