msg_tool\utils/
simple_pack.rs

1use crate::ext::io::*;
2use crate::types::*;
3use crate::utils::encoding::*;
4use anyhow::Result;
5use std::io::{Read, Seek};
6
7pub struct SimplePack<T: Read> {
8    inner: T,
9    total: u64,
10    current: u64,
11}
12
13impl<T: Read> SimplePack<T> {
14    pub fn next<'a>(&'a mut self) -> Result<Option<SimplePackEntry<'a, T>>> {
15        if self.current >= self.total {
16            return Ok(None);
17        }
18        let name = self.read_cstring()?;
19        let name = decode_to_string(Encoding::Utf8, name.as_bytes(), true)?;
20        let entry_size = self.read_u64()?;
21        Ok(Some(SimplePackEntry {
22            pack: self,
23            total: entry_size,
24            current: 0,
25            name,
26        }))
27    }
28}
29
30impl<T: Read> Read for SimplePack<T> {
31    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
32        let remaining = self.total - self.current;
33        if remaining == 0 {
34            return Ok(0);
35        }
36        let to_read = std::cmp::min(remaining, buf.len() as u64) as usize;
37        let bytes_read = self.inner.read(&mut buf[..to_read])?;
38        self.current += bytes_read as u64;
39        Ok(bytes_read)
40    }
41}
42
43pub struct SimplePackEntry<'a, T: Read> {
44    pub pack: &'a mut SimplePack<T>,
45    total: u64,
46    current: u64,
47    pub name: String,
48}
49
50impl<'a, T: Read> SimplePackEntry<'a, T> {
51    pub fn total_size(&self) -> u64 {
52        self.total
53    }
54
55    pub fn is_eof(&self) -> bool {
56        self.current >= self.total
57    }
58}
59
60impl<'a, T: Read> Drop for SimplePackEntry<'a, T> {
61    fn drop(&mut self) {
62        let to_skip = self.total - self.current;
63        if to_skip > 0 {
64            if let Err(e) = self.pack.skip(to_skip) {
65                eprintln!("Failed to skip remaining bytes in SimplePackEntry: {}", e);
66                crate::COUNTER.inc_error();
67            }
68        }
69    }
70}
71
72impl<'a, T: Read> Read for SimplePackEntry<'a, T> {
73    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
74        let remaining = self.total - self.current;
75        if remaining == 0 {
76            return Ok(0);
77        }
78        let to_read = std::cmp::min(remaining, buf.len() as u64) as usize;
79        let bytes_read = self.pack.read(&mut buf[..to_read])?;
80        self.current += bytes_read as u64;
81        Ok(bytes_read)
82    }
83}
84
85pub fn read_simple_pack<'a, T: Read + Seek + 'a>(
86    mut reader: T,
87) -> Result<SimplePack<Box<dyn Read + 'a>>> {
88    reader.read_and_equal(b"SPCK")?;
89    let flags = reader.read_u8()?;
90    // not compressed
91    if flags == 0 {
92        let pos = reader.stream_position()?;
93        let total = reader.stream_length()? - pos;
94        Ok(SimplePack {
95            inner: Box::new(reader),
96            total,
97            current: 0,
98        })
99    } else {
100        let compressed_size = reader.read_u64()?;
101        let uncompressed_size = reader.read_u64()?;
102        let compressed = reader.take(compressed_size);
103        let decompressed = zstd::stream::read::Decoder::new(compressed)?;
104        Ok(SimplePack {
105            inner: Box::new(decompressed),
106            total: uncompressed_size,
107            current: 0,
108        })
109    }
110}