From d201beadc9b8e5e5ea028992e8e12968455a64b9 Mon Sep 17 00:00:00 2001 From: Alex Orlenko Date: Mon, 27 Jul 2020 23:22:52 +0100 Subject: [PATCH] Add ChunkMode enum to mark chunks as text or binary --- src/ffi/lua.rs | 3 +++ src/lib.rs | 2 +- src/lua.rs | 67 ++++++++++++++++++++++++++++++++--------------- tests/function.rs | 16 +++-------- tests/tests.rs | 56 +++++++++++++++++++++++++++++++++++++-- 5 files changed, 107 insertions(+), 37 deletions(-) diff --git a/src/ffi/lua.rs b/src/ffi/lua.rs index 16d163d..e3f441f 100644 --- a/src/ffi/lua.rs +++ b/src/ffi/lua.rs @@ -38,7 +38,10 @@ pub use super::glue::LUA_REGISTRYINDEX; #[cfg(any(feature = "lua51", feature = "luajit"))] pub use super::glue::{LUA_ENVIRONINDEX, LUA_GLOBALSINDEX}; +#[cfg(not(feature = "luajit"))] pub const LUA_SIGNATURE: &[u8] = b"\x1bLua"; +#[cfg(feature = "luajit")] +pub const LUA_SIGNATURE: &[u8] = b"\x1bLJ"; // option for multiple returns in 'lua_pcall' and 'lua_call' pub const LUA_MULTRET: c_int = -1; diff --git a/src/lib.rs b/src/lib.rs index ce60d6c..85425a2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -81,7 +81,7 @@ pub use crate::ffi::lua_State; pub use crate::error::{Error, ExternalError, ExternalResult, Result}; pub use crate::function::Function; pub use crate::hook::{Debug, DebugNames, DebugSource, DebugStack, HookTriggers}; -pub use crate::lua::{Chunk, GCMode, Lua}; +pub use crate::lua::{Chunk, ChunkMode, GCMode, Lua}; pub use crate::multi::Variadic; pub use crate::scope::Scope; pub use crate::stdlib::StdLib; diff --git a/src/lua.rs b/src/lua.rs index 665b602..286299b 100644 --- a/src/lua.rs +++ b/src/lua.rs @@ -671,6 +671,7 @@ impl Lua { source: source.as_ref(), name: None, env: None, + mode: None, } } @@ -679,32 +680,35 @@ impl Lua { source: &[u8], name: Option<&CString>, env: Option>, + mode: Option, ) -> Result> { unsafe { let _sg = StackGuard::new(self.state); assert_stack(self.state, 1); - let mode_string = match self.safe { - true => cstr!("t"), - false => cstr!("bt"), + + let mode_str = match mode { + Some(ChunkMode::Binary) if self.safe => { + return Err(Error::SafetyError( + "binary chunks are disabled in safe mode".to_string(), + )) + } + Some(ChunkMode::Binary) => cstr!("b"), + Some(ChunkMode::Text) => cstr!("t"), + None if source.starts_with(ffi::LUA_SIGNATURE) && self.safe => { + return Err(Error::SafetyError( + "binary chunks are disabled in safe mode".to_string(), + )) + } + None => cstr!("bt"), }; - match if let Some(name) = name { - ffi::luaL_loadbufferx( - self.state, - source.as_ptr() as *const c_char, - source.len(), - name.as_ptr() as *const c_char, - mode_string, - ) - } else { - ffi::luaL_loadbufferx( - self.state, - source.as_ptr() as *const c_char, - source.len(), - ptr::null(), - mode_string, - ) - } { + match ffi::luaL_loadbufferx( + self.state, + source.as_ptr() as *const c_char, + source.len(), + name.map(|n| n.as_ptr()).unwrap_or_else(|| ptr::null()), + mode_str, + ) { ffi::LUA_OK => { if let Some(env) = env { self.push_value(env)?; @@ -1809,6 +1813,14 @@ pub struct Chunk<'lua, 'a> { source: &'a [u8], name: Option, env: Option>, + mode: Option, +} + +/// Represents chunk mode (text or binary). +#[derive(Clone, Copy)] +pub enum ChunkMode { + Text, + Binary, } impl<'lua, 'a> Chunk<'lua, 'a> { @@ -1840,6 +1852,17 @@ impl<'lua, 'a> Chunk<'lua, 'a> { Ok(self) } + /// Sets whether the chunk is text or binary (autodetected by default). + /// + /// Lua does not check the consistency of binary chunks, therefore this mode is allowed only + /// for instances created with [`Lua::unsafe_new`]. + /// + /// [`Lua::unsafe_new`]: struct.Lua.html#method.unsafe_new + pub fn set_mode(mut self, mode: ChunkMode) -> Chunk<'lua, 'a> { + self.mode = Some(mode); + self + } + /// Execute this chunk of code. /// /// This is equivalent to calling the chunk function with no arguments and no return values. @@ -1879,6 +1902,7 @@ impl<'lua, 'a> Chunk<'lua, 'a> { &self.expression_source(), self.name.as_ref(), self.env.clone(), + self.mode, ) { function.call(()) } else { @@ -1905,6 +1929,7 @@ impl<'lua, 'a> Chunk<'lua, 'a> { &self.expression_source(), self.name.as_ref(), self.env.clone(), + self.mode, ) { function.call_async(()) } else { @@ -1944,7 +1969,7 @@ impl<'lua, 'a> Chunk<'lua, 'a> { /// This simply compiles the chunk without actually executing it. pub fn into_function(self) -> Result> { self.lua - .load_chunk(self.source, self.name.as_ref(), self.env) + .load_chunk(self.source, self.name.as_ref(), self.env, self.mode) } fn expression_source(&self) -> Vec { diff --git a/tests/function.rs b/tests/function.rs index 7ed8d33..f7ce7c0 100644 --- a/tests/function.rs +++ b/tests/function.rs @@ -10,7 +10,7 @@ )] extern "system" {} -use mlua::{Error, Function, Lua, Result, String}; +use mlua::{Function, Lua, Result, String}; #[test] fn test_function() -> Result<()> { @@ -92,22 +92,12 @@ fn test_rust_function() -> Result<()> { fn test_dump() -> Result<()> { let lua = unsafe { Lua::unsafe_new() }; - let concat_tmp = lua + let concat_lua = lua .load(r#"function(arg1, arg2) return arg1 .. arg2 end"#) .eval::()?; - let concat_bytecode = concat_tmp.dump(false)?; - let concat = lua.load(&concat_bytecode).into_function()?; + let concat = lua.load(&concat_lua.dump(false)?).into_function()?; assert_eq!(concat.call::<_, String>(("foo", "bar"))?, "foobar"); - let lua = Lua::new(); - match lua.load(&concat_bytecode).exec() { - Ok(_) => panic!("expected SyntaxError, got no error"), - Err(Error::SyntaxError { message: msg, .. }) => { - assert!(msg.contains("attempt to load a binary chunk")) - } - Err(e) => panic!("expected SyntaxError, got {:?}", e), - } - Ok(()) } diff --git a/tests/tests.rs b/tests/tests.rs index 64b8d2c..61eefa3 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -16,8 +16,8 @@ use std::sync::Arc; use std::{error, f32, f64, fmt}; use mlua::{ - Error, ExternalError, Function, Lua, Nil, Result, StdLib, String, Table, UserData, Value, - Variadic, + ChunkMode, Error, ExternalError, Function, Lua, Nil, Result, StdLib, String, Table, UserData, + Value, Variadic, }; #[test] @@ -56,6 +56,23 @@ fn test_safety() -> Result<()> { Ok(_) => panic!("expected RuntimeError, got no error"), } + match lua.load("1 + 1").set_mode(ChunkMode::Binary).exec() { + Err(Error::SafetyError(msg)) => { + assert!(msg.contains("binary chunks are disabled in safe mode")) + } + Err(e) => panic!("expected SafetyError, got {:?}", e), + Ok(_) => panic!("expected SafetyError, got no error"), + } + + let bytecode = lua.load("return 1 + 1").into_function()?.dump(true)?; + match lua.load(&bytecode).exec() { + Err(Error::SafetyError(msg)) => { + assert!(msg.contains("binary chunks are disabled in safe mode")) + } + Err(e) => panic!("expected SafetyError, got {:?}", e), + Ok(_) => panic!("expected SafetyError, got no error"), + } + Ok(()) } @@ -127,6 +144,41 @@ fn test_eval() -> Result<()> { Ok(()) } +#[test] +fn test_load_mode() -> Result<()> { + let lua = unsafe { Lua::unsafe_new() }; + + assert_eq!( + lua.load("1 + 1").set_mode(ChunkMode::Text).eval::()?, + 2 + ); + match lua.load("1 + 1").set_mode(ChunkMode::Binary).exec() { + Ok(_) => panic!("expected SyntaxError, got no error"), + Err(Error::SyntaxError { message: msg, .. }) => { + assert!(msg.contains("attempt to load a text chunk")) + } + Err(e) => panic!("expected SyntaxError, got {:?}", e), + }; + + let bytecode = lua.load("return 1 + 1").into_function()?.dump(true)?; + assert_eq!(lua.load(&bytecode).eval::()?, 2); + assert_eq!( + lua.load(&bytecode) + .set_mode(ChunkMode::Binary) + .eval::()?, + 2 + ); + match lua.load(&bytecode).set_mode(ChunkMode::Text).exec() { + Ok(_) => panic!("expected SyntaxError, got no error"), + Err(Error::SyntaxError { message: msg, .. }) => { + assert!(msg.contains("attempt to load a binary chunk")) + } + Err(e) => panic!("expected SyntaxError, got {:?}", e), + }; + + Ok(()) +} + #[test] fn test_lua_multi() -> Result<()> { let lua = Lua::new();