libtw2_demo/
reader.rs

1use arrayvec::ArrayVec;
2use binrw::BinRead;
3use libtw2_common::digest::Sha256;
4use libtw2_common::num::Cast;
5use libtw2_huffman::instances::TEEWORLDS as HUFFMAN;
6use libtw2_packer::Unpacker;
7use std::io;
8use thiserror::Error;
9use warn::wrap;
10use warn::Warn;
11
12use crate::format::TickMarker;
13use crate::format::Warning;
14use crate::format::MAX_SNAPSHOT_SIZE;
15use crate::{format, DemoKind};
16
17#[derive(Error, Debug)]
18#[error(transparent)]
19pub enum ReadError {
20    Io(#[from] io::Error),
21    Binrw(#[from] binrw::Error),
22    Huffman(#[from] libtw2_huffman::DecompressionError),
23    #[error("Unexpected data end during secondary decompression of message")]
24    MessageVarIntUnexpectedEnd,
25    #[error("Too big decompressed size during secondary decompression of message")]
26    MessageVarIntTooLong,
27    #[error("Tick number did not increase")]
28    NotIncreasingTick,
29    #[error("The first snapshot is only a snapshot-delta")]
30    StartingDeltaSnapshot,
31    #[error("The tick number overflowed")]
32    TickOverflow,
33}
34
35impl ReadError {
36    pub fn io_error(self) -> Result<io::Error, ReadError> {
37        match self {
38            ReadError::Io(io) => Ok(io),
39            ReadError::Binrw(binrw::Error::Io(io)) => Ok(io),
40            err => Err(err),
41        }
42    }
43}
44
45trait SeekableRead: io::Read + io::Seek {}
46impl<T: io::Read + io::Seek> SeekableRead for T {}
47
48pub struct Reader {
49    data: Box<dyn SeekableRead>,
50    start: format::HeaderStart,
51    current_tick: Option<i32>,
52    raw: [u8; MAX_SNAPSHOT_SIZE],
53    huffman: ArrayVec<[u8; MAX_SNAPSHOT_SIZE]>,
54}
55
56impl Reader {
57    pub fn new<W, R>(mut data: R, warn: &mut W) -> Result<Reader, ReadError>
58    where
59        W: Warn<Warning>,
60        R: io::Read + io::Seek + 'static,
61    {
62        let start = format::HeaderStart::read(&mut data)?;
63        start.header.check(warn);
64        start.timeline_markers.check(warn);
65        Ok(Self {
66            data: Box::new(data),
67            start: start,
68            current_tick: None,
69            raw: [0; MAX_SNAPSHOT_SIZE],
70            huffman: ArrayVec::new(),
71        })
72    }
73    pub fn version(&self) -> format::Version {
74        self.start.version
75    }
76    pub fn net_version(&self) -> &[u8] {
77        self.start.header.net_version.raw()
78    }
79    pub fn map_name(&self) -> &[u8] {
80        self.start.header.map_name.raw()
81    }
82    pub fn map_size(&self) -> u32 {
83        self.start.header.map_size.assert_u32()
84    }
85    pub fn map_data(&self) -> &[u8] {
86        &self.start.map
87    }
88    pub fn map_crc(&self) -> u32 {
89        self.start.header.map_crc
90    }
91    pub fn kind(&self) -> DemoKind {
92        self.start.header.kind
93    }
94    pub fn length(&self) -> i32 {
95        self.start.header.length
96    }
97    pub fn timestamp(&self) -> &[u8] {
98        self.start.header.timestamp.raw()
99    }
100    pub fn timeline_markers(&self) -> &[i32] {
101        self.start.timeline_markers.markers()
102    }
103    pub fn map_sha256(&self) -> Option<Sha256> {
104        self.start
105            .map_sha256
106            .as_ref()
107            .map(|sha| Sha256(sha.sha_256))
108    }
109    pub fn read_chunk<W>(&mut self, warn: &mut W) -> Result<Option<format::RawChunk>, ReadError>
110    where
111        W: Warn<Warning>,
112    {
113        use crate::format::ChunkHeader;
114        use crate::format::DataKind;
115        use crate::format::RawChunk;
116
117        let chunk_header = match ChunkHeader::read(&mut self.data, self.start.version, warn)? {
118            Some(ch) => ch,
119            None => return Ok(None),
120        };
121        match chunk_header {
122            ChunkHeader::Tick {
123                marker: TickMarker::Absolute(t),
124                keyframe,
125            } => {
126                if let Some(previous) = self.current_tick {
127                    if previous >= t {
128                        return Err(ReadError::NotIncreasingTick);
129                    }
130                }
131                self.current_tick = Some(t);
132                Ok(Some(RawChunk::Tick {
133                    tick: t,
134                    keyframe: keyframe,
135                }))
136            }
137            ChunkHeader::Tick {
138                marker: TickMarker::Delta(d),
139                keyframe,
140            } => match self.current_tick {
141                None => Err(ReadError::StartingDeltaSnapshot),
142                Some(t) => match t.checked_add(d.i32()) {
143                    None => Err(ReadError::TickOverflow),
144                    Some(new_t) => {
145                        self.current_tick = Some(new_t);
146                        Ok(Some(RawChunk::Tick {
147                            tick: new_t,
148                            keyframe: keyframe,
149                        }))
150                    }
151                },
152            },
153            ChunkHeader::Data { kind, size } => {
154                if kind == DataKind::Unknown {
155                    return Ok(Some(RawChunk::Unknown));
156                }
157                let raw_data = &mut self.raw[..size.usize()];
158                self.data.read_exact(raw_data)?;
159                self.huffman.clear();
160                HUFFMAN.decompress(raw_data, &mut self.huffman)?;
161                Ok(Some(match kind {
162                    DataKind::Unknown => RawChunk::Unknown,
163                    DataKind::Snapshot => RawChunk::Snapshot(&self.huffman),
164                    DataKind::SnapshotDelta => RawChunk::SnapshotDelta(&self.huffman),
165                    DataKind::Message => {
166                        let mut unpacker = Unpacker::new(&self.huffman);
167                        let mut len = 0;
168                        let mut buffer = self.raw.chunks_mut(4);
169                        while !unpacker.is_empty() {
170                            let n: i32 = unpacker
171                                .read_int(wrap(warn))
172                                .map_err(|_| ReadError::MessageVarIntUnexpectedEnd)?;
173                            buffer
174                                .next()
175                                .ok_or(ReadError::MessageVarIntTooLong)?
176                                .copy_from_slice(&n.to_le_bytes());
177                            len += 4;
178                        }
179                        RawChunk::Message(&self.raw[..len])
180                    }
181                }))
182            }
183        }
184    }
185}
libtw2_demo/reader.rs

libtw2_demo/
reader.rs
