faf_replay_parser/

reader.rs

use byteorder::{LittleEndian, ReadBytesExt};
use ordered_hash_map::OrderedHashMap;

use super::lua::*;

use std::error::Error;
use std::ffi::CString;
use std::fmt;
use std::io;
use std::string::FromUtf8Error;

#[derive(Debug)]
pub enum ReplayReadError {
    Desynced(u32),
    Malformed(&'static str),
    MalformedUtf8(FromUtf8Error),
    IO(std::io::Error),
}
pub type ReplayResult<T> = Result<T, ReplayReadError>;

impl Error for ReplayReadError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::IO(e) => Some(e),
            Self::MalformedUtf8(e) => Some(e),
            _ => None,
        }
    }
}

impl fmt::Display for ReplayReadError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        use ReplayReadError::*;

        match self {
            Desynced(tick) => write!(f, "desynced at tick {}", tick),
            Malformed(msg) => write!(f, "{}", msg),
            MalformedUtf8(e) => write!(f, "{}", e),
            IO(e) => write!(f, "{}", e),
        }
    }
}

impl From<FromUtf8Error> for ReplayReadError {
    fn from(err: FromUtf8Error) -> ReplayReadError {
        ReplayReadError::MalformedUtf8(err)
    }
}

impl From<std::io::Error> for ReplayReadError {
    fn from(err: std::io::Error) -> ReplayReadError {
        ReplayReadError::IO(err)
    }
}

impl From<LuaTypeError> for ReplayReadError {
    fn from(_err: LuaTypeError) -> ReplayReadError {
        ReplayReadError::Malformed("invalud lua type")
    }
}

/// Extensions to [`std::io::Read`]
pub trait ReplayReadExt: io::Read {
    /// Convenience wrapper around [`std::io::Read::read_to_end`] for allocating a new `Vec`.
    fn vec_read_to_end(&mut self) -> io::Result<Vec<u8>> {
        let mut res = Vec::new();
        self.read_to_end(&mut res)?;
        Ok(res)
    }

    /// Convenience wrapper around [`std::io::Read::read_exact`] for allocating a new `Vec`.
    fn vec_read_exact(&mut self, n: usize) -> io::Result<Vec<u8>> {
        // Don't pre-allocate more than 1 MiB.
        let capacity = std::cmp::min(n, crate::MIB);
        // TODO: Use #![feature(read_initializer)] to skip initialization here. Since read requries
        // that the buffer be initialized, we need to write 0's to it for now.
        let mut vec = vec![0; capacity];
        self.read_exact(&mut vec)?;

        Ok(vec)
    }

    /// Convenience wrapper around [`std::io::Read::read_exact`] for reading data into a Vec,
    /// resizing it if more space is needed.
    fn read_exact_to_vec(&mut self, n: usize, buf: &mut Vec<u8>) -> io::Result<()> {
        let mut remaining = n;
        // Read in 1 MiB chunks in case the requested size is really big.
        while remaining > 0 {
            let chunk = std::cmp::min(remaining, crate::MIB);
            let index = n - remaining;

            buf.resize(chunk, 0);
            self.read_exact(&mut buf[index..index + chunk])?;
            remaining -= chunk;
        }
        Ok(())
    }

    /// Ignore the next `n` bytes.
    fn skip(&mut self, n: usize) -> io::Result<()> {
        self.vec_read_exact(n)?;
        Ok(())
    }

    /// Ignore the next `n` bytes by reading them into the buffer. This may perform multiple calls
    /// to `read` so make sure the buffer is sufficiently large.
    ///
    /// If frequent skips are needed, this can be more efficient than calling `skip` as the buffer
    /// can be reused.
    fn skip_buf(&mut self, n: usize, buf: &mut [u8]) -> io::Result<()> {
        let mut remaining = n;
        while remaining > 0 {
            let to_read = std::cmp::min(remaining, buf.len());
            unsafe {
                match self.read(buf.get_unchecked_mut(..to_read)) {
                    Ok(0) => {
                        return Err(io::Error::new(
                            io::ErrorKind::UnexpectedEof,
                            "failed to skip desired bytes",
                        ));
                    }
                    Ok(n) => remaining -= n,
                    Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                    Err(e) => return Err(e),
                }
            }
        }
        Ok(())
    }

    fn read_i32_le(&mut self) -> io::Result<i32> {
        self.read_i32::<LittleEndian>()
    }

    fn read_u32_le(&mut self) -> io::Result<u32> {
        self.read_u32::<LittleEndian>()
    }

    fn read_u16_le(&mut self) -> io::Result<u16> {
        self.read_u16::<LittleEndian>()
    }

    fn read_i16_le(&mut self) -> io::Result<i16> {
        self.read_i16::<LittleEndian>()
    }

    fn read_f32_le(&mut self) -> io::Result<f32> {
        self.read_f32::<LittleEndian>()
    }

    /// Re-export of `ReadBytesExt::read_u8`
    fn read_u8(&mut self) -> io::Result<u8> {
        ReadBytesExt::read_u8(self)
    }

    fn read_bool(&mut self) -> io::Result<bool> {
        Ok(self.read_u8()? == 1)
    }
}

/// Implement `ReplayReadExt` for all types that implement `Read`.
impl<R: io::Read + ?Sized> ReplayReadExt for R {}

/// Extensions to `std::io::BufRead`
pub trait ReplayBufReadExt: io::BufRead + ReplayReadExt {
    /// Convenience wrapper around `std::io::BufRead::read_until` for allocating a new `Vec`.
    fn vec_read_until(&mut self, byte: u8) -> io::Result<Vec<u8>> {
        let mut vec = Vec::new();
        self.read_until(byte, &mut vec)?;

        Ok(vec)
    }

    /// Read a null terminated string from the stream into a new `Vec`.
    fn vec_read_null_term(&mut self) -> io::Result<Vec<u8>> {
        let mut vec = Vec::new();
        self.read_until(0x00, &mut vec)?;
        vec.pop();

        Ok(vec)
    }

    /// Like `vec_read_null_term` but initialize the `Vec` with given capacity.
    fn vec_read_null_term_with_capacity(&mut self, capacity: usize) -> io::Result<Vec<u8>> {
        let mut vec = Vec::with_capacity(capacity);
        self.read_until(0x00, &mut vec)?;
        vec.pop();

        Ok(vec)
    }

    /// Read a null terminated string from the stream and decode it as UTF-8. If it does not
    /// decode succesfully, a MalformedUtf8 error is returned.
    fn read_string(&mut self) -> io::Result<String> {
        String::from_utf8(self.vec_read_null_term()?).map_err(|_| {
            io::Error::new(
                io::ErrorKind::InvalidData,
                "stream did not contain valid UTF-8",
            )
        })
    }

    /// Like `read_string` but initialize the buffer with given capacity.
    fn read_string_with_capacity(&mut self, capacity: usize) -> io::Result<String> {
        String::from_utf8(self.vec_read_null_term_with_capacity(capacity)?).map_err(|_| {
            io::Error::new(
                io::ErrorKind::InvalidData,
                "stream did not contain valid UTF-8",
            )
        })
    }

    /// Read a null terminated string from the stream and wrap it in a CString.
    fn read_c_string(&mut self) -> io::Result<CString> {
        unsafe { Ok(CString::from_vec_unchecked(self.vec_read_null_term()?)) }
    }

    /// Read a lua object.
    fn read_lua_object(&mut self) -> ReplayResult<LuaObject> {
        let lua_type = self.read_u8()?;
        self.read_lua_object_as(lua_type)
    }

    fn read_lua_object_as(&mut self, lua_type: u8) -> ReplayResult<LuaObject> {
        match lua_type {
            LUA_FLOAT_MARKER => Ok(LuaObject::Float(self.read_f32_le()?)),
            LUA_STRING_MARKER => Ok(LuaObject::String(self.read_c_string()?)),
            LUA_NIL_MARKER => {
                self.read_u8()?;
                Ok(LuaObject::Nil)
            }
            LUA_BOOL_MARKER => Ok(LuaObject::Bool(self.read_bool()?)),
            LUA_TABLE_MARKER => {
                let mut res = OrderedHashMap::new();
                loop {
                    match self.read_u8()? {
                        LUA_END_MARKER => break,
                        next_type => {
                            let mut key = self.read_lua_object_as(next_type)?;
                            if let LuaObject::String(s) = key {
                                key = LuaObject::from(
                                    s.into_string().map_err(|e| LuaTypeError::from(e))?,
                                )
                            }
                            let value = self.read_lua_object()?;
                            res.insert(key, value);
                        }
                    }
                }
                Ok(LuaObject::Table(res))
            }
            _ => Err(LuaTypeError {})?,
        }
    }
}

/// Implement `ReplayBufReadExt` for all types that implement `Read`.
impl<R: io::BufRead + ReplayReadExt + ?Sized> ReplayBufReadExt for R {}

#[cfg(test)]
mod tests {
    use super::*;
    use pretty_assertions::assert_eq;

    use std::io::Cursor;

    #[test]
    fn test_skip_buf() {
        let mut cur = Cursor::new(vec![0u8; 100]);
        let mut buf = [0; 10];

        cur.skip_buf(17, &mut buf).unwrap();
        assert_eq!(cur.position(), 17);

        cur.skip_buf(25, &mut buf).unwrap();
        assert_eq!(cur.position(), 17 + 25);
    }

    #[test]
    fn read_empty_string() {
        let mut cur = Cursor::new(b"\x00");
        assert_eq!(cur.read_string().unwrap(), "")
    }

    #[test]
    fn read_string() {
        let data: &[u8] = b"This is a c-style string\x00";
        let mut cur = Cursor::new(data);
        assert_eq!(cur.read_string().unwrap(), "This is a c-style string")
    }

    #[test]
    fn read_copyright_string() {
        // Taken directly from replays using the "Total Mayhem" mod which was an official GPG mod.
        // "Copyright © 2006, Gas Powered Games" but as a single byte encoding
        // Strings like these should not be decoded by the parser, as they would generate an error
        let data: &[u8] = &[
            67, 111, 112, 121, 114, 105, 103, 104, 116, 32, 169, 32, 50, 48, 48, 54, 44, 32, 71,
            97, 115, 32, 80, 111, 119, 101, 114, 101, 100, 32, 71, 97, 109, 101, 115, 0,
        ];
        let mut cur = Cursor::new(data);
        cur.read_string().unwrap_err();
    }

    #[test]
    fn read_i32_le() {
        let mut data: &[u8] = b"\xef\xbe\xad\x7e";
        assert_eq!(data.read_i32_le().unwrap(), 0x7eadbeef)
    }

    #[test]
    fn read_i16_le() {
        let mut data: &[u8] = b"\xad\xde";
        assert_eq!(data.read_u16_le().unwrap(), 0xdead)
    }

    #[test]
    fn vec_read_until() {
        let mut data: &[u8] = &[1, 2, 3, 4, 5];
        assert_eq!(data.vec_read_until(3).unwrap(), vec![1, 2, 3]);
    }

    #[test]
    fn read_lua_table() {
        let data: &[u8] = b"\x04\x05";
        let mut cur = Cursor::new(data);
        let table = cur.read_lua_object().unwrap();
        let hashmap = table.into_hashmap().unwrap();
        assert!(hashmap.is_empty())
    }

    #[test]
    fn test_replay_read_error_display() {
        use std::io::Read;

        assert_eq!(
            format!("{}", ReplayReadError::Desynced(10)),
            "desynced at tick 10"
        );
        assert_eq!(
            format!("{}", ReplayReadError::Malformed("bad something or other")),
            "bad something or other"
        );
        assert_eq!(
            format!(
                "{}",
                <ReplayReadError as From<FromUtf8Error>>::from(
                    String::from_utf8(b"\xFF".to_vec()).unwrap_err()
                )
            ),
            "invalid utf-8 sequence of 1 bytes from index 0"
        );

        let mut buf: &[u8] = b"";
        assert_eq!(
            format!(
                "{}",
                ReplayReadError::from(
                    (&mut buf)
                        .read_exact(vec![0; 5].as_mut_slice())
                        .unwrap_err()
                )
            ),
            "failed to fill whole buffer"
        );
    }
}