fromsoftware_shared/ext/
read.rs

1use std::io;
2
3/// An extension on top of [io::Read] that reads data with length delimiters
4/// that was written using
5/// [LengthDelimitedWriteExt](super::LengthDelimitedWriteExt). This is
6/// particularly useful for adding extra data before or after data structures
7/// written by FSW games.
8///
9/// The specific format used by this trait is a little-endian four-byte header
10/// indicating the length of the data in bytes, followed by the data itself.
11pub trait LengthDelimitedReadExt {
12    /// Reads a length-delimited amount of data from the underlying reader into
13    /// a newly-allocated [Vec].
14    ///
15    /// This is the dual of [LengthDelimitedWriteExt.write_delimited].
16    ///
17    /// Unlike [io::Read.read], this will return an error result if there's not
18    /// enough data available to read the full result.
19    fn read_delimited(&mut self) -> io::Result<Vec<u8>>;
20
21    /// Reads a length-delimited amount of data from the underlying reader into
22    /// a newly-allocated [String].
23    ///
24    /// This is the dual of [LengthDelimitedWriteExt.write_str_delimited].
25    ///
26    /// Unlike [io::Read.read], this will return an error result if there's not
27    /// enough data available to read the full result. This will also return an
28    /// error if the data isn't valid UTF-8.
29    fn read_str_delimited(&mut self) -> io::Result<String> {
30        String::from_utf8(self.read_delimited()?).map_err(io::Error::other)
31    }
32}
33
34impl<T: ?Sized + io::Read> LengthDelimitedReadExt for T {
35    fn read_delimited(&mut self) -> io::Result<Vec<u8>> {
36        let mut delimiter = [0; std::mem::size_of::<u32>()];
37        if self.read(&mut delimiter)? < delimiter.len() {
38            return Err(io::Error::other("Couldn't read delimiter"));
39        }
40
41        let delimiter: usize = u32::from_le_bytes(delimiter)
42            .try_into()
43            .map_err(io::Error::other)?;
44
45        // Don't allow a delimiter to force us to allocate massive amounts of
46        // memory. Read at most 16kb at a time and only continue if the stream
47        // actually has that amount of data available.
48        const CHUNK_SIZE: usize = 0x4000;
49        let mut result = Vec::with_capacity(std::cmp::min(delimiter, CHUNK_SIZE));
50        while result.len() < delimiter {
51            let prev_len = result.len();
52            result.resize(std::cmp::min(delimiter, prev_len + CHUNK_SIZE), 0);
53            let size = self.read(&mut result[prev_len..])?;
54
55            let expected = result.len() - prev_len;
56            if size < expected {
57                return Err(io::Error::other(format!(
58                    "Expected {} bytes but only {} were read",
59                    expected, size
60                )));
61            }
62        }
63
64        Ok(result)
65    }
66}
67
68#[cfg(test)]
69mod test {
70    use super::LengthDelimitedReadExt;
71    use crate::ext::LengthDelimitedWriteExt;
72
73    #[test]
74    fn write_and_read_small_bytes() {
75        let mut writer = Vec::new();
76        assert_eq!(writer.write_delimited(&[1, 2, 3, 4, 5]).unwrap(), 9);
77
78        let mut reader = &writer[..];
79        assert_eq!(reader.read_delimited().unwrap(), &[1, 2, 3, 4, 5]);
80    }
81
82    #[test]
83    fn write_and_read_small_string() {
84        let mut writer = Vec::new();
85        assert_eq!(writer.write_str_delimited("hello!").unwrap(), 10);
86
87        let mut reader = &writer[..];
88        assert_eq!(reader.read_str_delimited().unwrap(), "hello!");
89    }
90
91    #[test]
92    fn write_and_read_large_bytes() {
93        let mut data = [0; 0x5000];
94        for (i, byte) in data.iter_mut().enumerate() {
95            *byte = (i % 255).try_into().unwrap();
96        }
97
98        let mut writer = Vec::new();
99        assert_eq!(writer.write_delimited(&data[..]).unwrap(), 0x5004);
100
101        let mut reader = &writer[..];
102        assert_eq!(reader.read_delimited().unwrap(), &data);
103    }
104}
fromsoftware_shared/ext/read.rs

fromsoftware_shared/ext/
read.rs
