libosu/db/
binary.rs

1//! Convenience traits for dealing with OSU binary data.
2
3use std::io::{Read, Write};
4
5use byteorder::{ReadBytesExt, WriteBytesExt};
6
7/// Result for Error
8pub type Result<T, E = Error> = std::result::Result<T, E>;
9
10/// Errors that could arise from reading binary beatmap data
11#[derive(Debug, Error)]
12pub enum Error {
13    /// IO Error
14    #[error("io error: {0}")]
15    Io(#[from] std::io::Error),
16
17    /// UTF8 String conversion error
18    #[error("string conversion error: {0}")]
19    Utf8(#[from] std::string::FromUtf8Error),
20
21    /// ULEB overflows 128 bits
22    #[error("uleb error")]
23    UlebOverflow,
24
25    /// Character in front of the string is not 0x0 or 0xb
26    #[error("invalid string status char: {0}")]
27    InvalidStringStatusChar(u8),
28}
29
30/// Extends Read with more reading functions specific to OSU
31pub trait ReadBytesOsu: Read {
32    /// Read data from the reader in [ULEB128][1] format
33    ///
34    /// [1]: https://en.wikipedia.org/wiki/LEB128#Unsigned_LEB128
35    fn read_uleb128(&mut self) -> Result<u128> {
36        let mut buf = [0];
37        let mut byte_index = 0;
38        self.read_exact(&mut buf)?;
39
40        let mut total = (buf[0] & 0b01111111) as u128;
41        while (buf[0] & 0b10000000) == 0b10000000 {
42            byte_index += 1;
43            if byte_index > 9 {
44                return Err(Error::UlebOverflow);
45            }
46
47            self.read_exact(&mut buf)?;
48            total += ((buf[0] & 0b01111111) as u128) << (7 * byte_index)
49        }
50
51        Ok(total)
52    }
53
54    /// Read a string from the reader
55    ///
56    /// The first byte indicates whether a string is there (0x0B) or not (0x00). Then the length is
57    /// encoded as a ULEB128 number, and then finally the string itself is encoded as UTF-8
58    fn read_uleb128_string(&mut self) -> Result<String> {
59        match self.read_u8()? {
60            // string isn't there
61            0x0 => Ok(String::new()),
62
63            // read string normally
64            0xb => {
65                let len = self.read_uleb128()?;
66                if len == 0 {
67                    return Ok(String::new());
68                }
69
70                let mut buf = vec![0; len as usize];
71                self.read_exact(&mut buf)?;
72                let string = String::from_utf8(buf)?;
73                Ok(string)
74            }
75
76            // error
77            v => Err(Error::InvalidStringStatusChar(v)),
78        }
79    }
80}
81
82impl<R: Read + ?Sized> ReadBytesOsu for R {}
83
84/// Extends Write with more writing functions specific to OSU
85pub trait WriteBytesOsu: Write {
86    /// Writes a ULEB128 value into the writer
87    fn write_uleb128(&mut self, mut n: u128) -> Result<()> {
88        while n > 0 {
89            let mut byte = (n & 0x7fu128) as u8;
90            n >>= 7;
91            if n > 0 {
92                // more bytes
93                byte |= 1 << 7;
94            }
95            self.write_u8(byte)?;
96        }
97        Ok(())
98    }
99
100    /// Writes a string value into the writer
101    fn write_uleb128_string(&mut self, string: impl AsRef<str>) -> Result<()> {
102        let string = string.as_ref();
103        if string.is_empty() {
104            self.write_u8(0x0)?;
105            return Ok(());
106        }
107
108        self.write_u8(0xb)?;
109        self.write_uleb128(string.len() as u128)?;
110        self.write_all(string.as_bytes())?;
111        Ok(())
112    }
113}
114
115impl<W: Write> WriteBytesOsu for W {}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120    use anyhow::Result;
121    use std::io::Cursor;
122
123    #[test]
124    fn test_read_uleb128() -> Result<()> {
125        let mut num = Cursor::new([0xE5, 0x8E, 0x26]);
126        assert_eq!(num.read_uleb128()?, 624485);
127        Ok(())
128    }
129
130    #[test]
131    fn test_write_uleb128() -> Result<()> {
132        let mut buf = Vec::new();
133        let mut curs = Cursor::new(&mut buf);
134        curs.write_uleb128(624485)?;
135        assert_eq!(buf, [0xe5, 0x8e, 0x26]);
136        Ok(())
137    }
138
139    #[test]
140    fn test_read_uleb128_string() -> Result<()> {
141        let text = "Hello World";
142        let mut replay_string = vec![0x0Bu8];
143        replay_string.push(text.len() as u8);
144        replay_string.extend(text.bytes());
145
146        let mut reader = Cursor::new(replay_string);
147        assert_eq!(reader.read_uleb128_string()?, text.to_string());
148
149        let mut reader_empty = Cursor::new(vec![0x00]);
150        assert_eq!(reader_empty.read_uleb128_string()?, String::new());
151
152        let mut reader_0_len = Cursor::new(vec![0x0B, 0x00]);
153        assert_eq!(reader_0_len.read_uleb128_string()?, String::new());
154        Ok(())
155    }
156
157    #[test]
158    fn test_write_uleb128_string() -> Result<()> {
159        let mut buf = Vec::new();
160        let mut curs = Cursor::new(&mut buf);
161        curs.write_uleb128_string("Hello, world!")?;
162        assert_eq!(buf, b"\x0b\x0dHello, world!");
163        Ok(())
164    }
165}