msg_tool\utils/
simple_pack.rs1use 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 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}