Add ChunkMode enum to mark chunks as text or binary

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;

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;

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<Value<'lua>>,
+        mode: Option<ChunkMode>,
     ) -> Result<Function<'lua>> {
         unsafe {
             let _sg = StackGuard::new(self.state);
             assert_stack(self.state, 1);
-            let mode_string = match self.safe {
+
-                true => cstr!("t"),
+            let mode_str = match mode {
-                false => cstr!("bt"),
+                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 {
+            match ffi::luaL_loadbufferx(
-                ffi::luaL_loadbufferx(
                 self.state,
                 source.as_ptr() as *const c_char,
                 source.len(),
-                    name.as_ptr() as *const c_char,
+                name.map(|n| n.as_ptr()).unwrap_or_else(|| ptr::null()),
-                    mode_string,
+                mode_str,
-                )
+            ) {
-            } else {
-                ffi::luaL_loadbufferx(
-                    self.state,
-                    source.as_ptr() as *const c_char,
-                    source.len(),
-                    ptr::null(),
-                    mode_string,
-                )
-            } {
                 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<CString>,
     env: Option<Value<'lua>>,
+    mode: Option<ChunkMode>,
+}
+
+/// 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<Function<'lua>> {
         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<u8> {

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::<Function>()?;
-    let concat_bytecode = concat_tmp.dump(false)?;
+    let concat = lua.load(&concat_lua.dump(false)?).into_function()?;
-    let concat = lua.load(&concat_bytecode).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(())
 }

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,
+    ChunkMode, Error, ExternalError, Function, Lua, Nil, Result, StdLib, String, Table, UserData,
-    Variadic,
+    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::<i32>()?,
+        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::<i32>()?, 2);
+    assert_eq!(
+        lua.load(&bytecode)
+            .set_mode(ChunkMode::Binary)
+            .eval::<i32>()?,
+        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();