123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299 |
- use std::ffi::c_void;
- use std::io::{self, Cursor};
- use std::io::{Read, Seek, Write};
- use std::os::raw::c_uint;
- use lua::lua_State;
- use lua_shared as lua;
- use crate::check_slice;
- #[derive(Debug)]
- enum KOption {
- I8,
- U8,
- I16,
- U16,
- I32,
- U32,
- Usize,
- Float,
- Double,
- Char,
- String,
- NOP,
- }
- enum Endianness {
- Little,
- Big,
- Native,
- }
- struct ReaderState<'a> {
- state: *mut c_void,
- endianness: Endianness,
- fmt: &'a [u8],
- }
- fn is_digit(byte: u8) -> bool {
- byte ^ b'0' < 10
- }
- unsafe fn read_number(state: &mut ReaderState) -> Option<usize> {
- if state.fmt.len() == 0 || !is_digit(state.fmt[0]) {
- return None;
- }
- let mut result = 0;
- while state.fmt.len() > 0 && is_digit(state.fmt[0]) {
- let (opt, rest) = state.fmt.split_at(1);
- state.fmt = rest;
- result = result * 10 + (opt[0] - b'0') as usize;
- }
- Some(result)
- }
- #[derive(Debug)]
- pub enum StructError {
- Error(*const u8),
- ArgError(i32, *const u8),
- InvalidFormatOption(*const u8, c_uint),
- IOError(io::Error),
- }
- impl From<std::io::Error> for StructError {
- fn from(e: std::io::Error) -> Self {
- StructError::IOError(e)
- }
- }
- unsafe fn get_option(state: &mut ReaderState) -> Result<Option<(KOption, usize)>, StructError> {
- if state.fmt.len() == 0 {
- return Ok(None);
- }
- let (opt, rest) = state.fmt.split_at(1);
- state.fmt = rest;
- match opt[0] {
- b'b' => Ok(Some((KOption::I8, std::mem::size_of::<i8>()))),
- b'B' => Ok(Some((KOption::U8, std::mem::size_of::<u8>()))),
- b'h' => Ok(Some((KOption::I16, std::mem::size_of::<i16>()))),
- b'H' => Ok(Some((KOption::U16, std::mem::size_of::<u16>()))),
- b'l' => Ok(Some((KOption::I32, std::mem::size_of::<i32>()))),
- b'L' => Ok(Some((KOption::U32, std::mem::size_of::<u32>()))),
- b'T' => Ok(Some((KOption::Usize, std::mem::size_of::<usize>()))),
- b'f' => Ok(Some((KOption::Float, std::mem::size_of::<f32>()))),
- b'd' => Ok(Some((KOption::Double, std::mem::size_of::<f64>()))),
- b'n' => Ok(Some((KOption::Double, std::mem::size_of::<f64>()))),
- b's' => Ok(Some((KOption::String, std::mem::size_of::<u16>()))),
- b'c' => match read_number(state) {
- Some(len) => Ok(Some((KOption::Char, len))),
- None => Err(StructError::Error(lua::cstr!(
- "missing size for format option 'c'"
- ))),
- },
- b' ' => Ok(Some((KOption::NOP, 0))),
- b'<' => {
- state.endianness = Endianness::Little;
- Ok(Some((KOption::NOP, 0)))
- }
- b'>' => {
- state.endianness = Endianness::Big;
- Ok(Some((KOption::NOP, 0)))
- }
- b'=' => {
- state.endianness = Endianness::Native;
- Ok(Some((KOption::NOP, 0)))
- }
- token @ _ => Err(StructError::InvalidFormatOption(
- lua::cstr!("invalid format option '%c'"),
- token as c_uint,
- )),
- }
- }
- // I don't fucking care. luaL_Buffer is allocated on the stack.
- // At the same time, this buffer is two times bigger than lua's string buffer.
- static mut STRING_BUFFER: [u8; 65536] = [0; 65536];
- macro_rules! pack_number {
- ($state:ident, $buffer:ident, $value:tt) => {
- match $state.endianness {
- Endianness::Little => $buffer.write(&$value.to_le_bytes())?,
- Endianness::Big => $buffer.write(&$value.to_be_bytes())?,
- Endianness::Native => $buffer.write(&$value.to_ne_bytes())?,
- }
- };
- }
- macro_rules! unpack_number {
- ($state:ident, $buffer:ident, $typ:ty) => {{
- let mut data = [0; std::mem::size_of::<$typ>()];
- $buffer.read(&mut data)?;
- match $state.endianness {
- Endianness::Little => <$typ>::from_le_bytes(data),
- Endianness::Big => <$typ>::from_be_bytes(data),
- Endianness::Native => <$typ>::from_ne_bytes(data),
- }
- }};
- }
- pub fn pack(state: lua_State, fmt: &[u8], start: i32) -> Result<&'static [u8], StructError> {
- unsafe {
- let mut arg = start;
- let mut reader_state = ReaderState {
- state: state,
- endianness: Endianness::Native,
- fmt: fmt,
- };
- let mut buffer = Cursor::new(&mut STRING_BUFFER[..]);
- while let Some((option, size)) = get_option(&mut reader_state)? {
- if let KOption::NOP = option {
- continue;
- }
- if STRING_BUFFER.len() - (buffer.seek(std::io::SeekFrom::Current(0))? as usize) < size {
- return Err(StructError::Error(lua::cstr!("buffer overflow")));
- }
- match option {
- KOption::I8 => {
- let value = lua::Lcheckinteger(reader_state.state, arg) as i8;
- pack_number!(reader_state, buffer, value);
- }
- KOption::U8 => {
- let value = lua::Lcheckinteger(reader_state.state, arg) as u8;
- pack_number!(reader_state, buffer, value);
- }
- KOption::I16 => {
- let value = lua::Lcheckinteger(reader_state.state, arg) as i16;
- pack_number!(reader_state, buffer, value);
- }
- KOption::U16 => {
- let value = lua::Lcheckinteger(reader_state.state, arg) as u16;
- pack_number!(reader_state, buffer, value);
- }
- KOption::I32 => {
- let value = lua::Lcheckinteger(reader_state.state, arg) as i32;
- pack_number!(reader_state, buffer, value);
- }
- KOption::U32 => {
- let value = lua::Lcheckinteger(reader_state.state, arg) as u32;
- pack_number!(reader_state, buffer, value);
- }
- KOption::Usize => {
- let value = lua::Lcheckinteger(reader_state.state, arg) as usize;
- pack_number!(reader_state, buffer, value);
- }
- KOption::Float => {
- let value = lua::Lchecknumber(reader_state.state, arg) as f32;
- pack_number!(reader_state, buffer, value);
- }
- KOption::Double => {
- let value = lua::Lchecknumber(reader_state.state, arg);
- pack_number!(reader_state, buffer, value);
- }
- KOption::Char => {
- let str = check_slice!(state, arg);
- if str.len() >= size {
- buffer.write(&str[..size])?;
- } else {
- buffer.write(str)?;
- for _ in 0..size - str.len() {
- buffer.write(&[0])?;
- }
- }
- }
- KOption::String => {
- let str = check_slice!(state, arg);
- if (str.len() > 1 << 16 - 2)
- || ((buffer.seek(std::io::SeekFrom::Current(0))? as usize) + size
- > STRING_BUFFER.len())
- {
- return Err(StructError::ArgError(
- arg,
- lua::cstr!("string won't fit in the buffer"),
- ));
- }
- pack_number!(reader_state, buffer, (str.len() as u16));
- buffer.write(str)?;
- }
- KOption::NOP => {}
- }
- arg += 1;
- }
- Ok(&STRING_BUFFER[..(buffer.seek(std::io::SeekFrom::Current(0))? as usize)])
- }
- }
- pub fn unpack(state: lua_State, fmt: &[u8], data: &[u8]) -> Result<i32, StructError> {
- unsafe {
- let mut nrets = 0;
- let mut reader_state = ReaderState {
- state: state,
- endianness: Endianness::Native,
- fmt: fmt,
- };
- let mut buffer = Cursor::new(data);
- while let Some((option, size)) = get_option(&mut reader_state)? {
- if let KOption::NOP = option {
- continue;
- }
- if data.len() - (buffer.seek(std::io::SeekFrom::Current(0))? as usize) < size {
- return Err(StructError::Error(lua::cstr!("data string too short")));
- }
- nrets += 1;
- match option {
- KOption::I8 => {
- let value = unpack_number!(reader_state, buffer, i8);
- lua::pushinteger(state, value as _);
- }
- KOption::U8 => {
- let value = unpack_number!(reader_state, buffer, u8);
- lua::pushinteger(state, value as _);
- }
- KOption::I16 => {
- let value = unpack_number!(reader_state, buffer, i16);
- lua::pushinteger(state, value as _);
- }
- KOption::U16 => {
- let value = unpack_number!(reader_state, buffer, u16);
- lua::pushinteger(state, value as _);
- }
- KOption::I32 => {
- let value = unpack_number!(reader_state, buffer, i32);
- lua::pushinteger(state, value as _);
- }
- KOption::U32 => {
- let value = unpack_number!(reader_state, buffer, u32);
- lua::pushinteger(state, value as _);
- }
- KOption::Usize => {
- let value = unpack_number!(reader_state, buffer, usize);
- lua::pushinteger(state, value as _);
- }
- KOption::Float => {
- let value = unpack_number!(reader_state, buffer, f32);
- lua::pushnumber(state, value as _);
- }
- KOption::Double => {
- let value = unpack_number!(reader_state, buffer, f64);
- lua::pushnumber(state, value);
- }
- KOption::Char => {
- let offset = buffer.seek(std::io::SeekFrom::Current(0))? as usize;
- lua::pushlstring(state, data.as_ptr().add(offset), size);
- }
- KOption::String => {
- let value = unpack_number!(reader_state, buffer, u16);
- let offset = buffer.seek(std::io::SeekFrom::Current(0))? as usize;
- if offset + value as usize > data.len() {
- return Err(StructError::Error(lua::cstr!("data string too short")));
- }
- buffer.seek(std::io::SeekFrom::Current(value as _))?;
- lua::pushlstring(state, data.as_ptr().add(offset), value as _);
- }
- KOption::NOP => {}
- }
- }
- Ok(nrets)
- }
- }
|