From a6ca65aa7431b3ce10038600995ca68d98829252 Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Sun, 30 Oct 2022 11:41:09 +0000
Subject: [PATCH] Better checks and tests when trying to modify a Luau readonly
 table

---
 src/table.rs  | 41 ++++++++++++++++++++++++++++-------------
 tests/luau.rs | 33 +++++++++++++++++++++++++--------
 2 files changed, 53 insertions(+), 21 deletions(-)

diff --git a/src/table.rs b/src/table.rs
index 523ec6b..424aaf0 100644
--- a/src/table.rs
+++ b/src/table.rs
@@ -234,6 +234,9 @@ impl<'lua> Table<'lua> {
 
     /// Sets a key-value pair without invoking metamethods.
     pub fn raw_set<K: ToLua<'lua>, V: ToLua<'lua>>(&self, key: K, value: V) -> Result<()> {
+        #[cfg(feature = "luau")]
+        self.check_readonly_write()?;
+
         let lua = self.0.lua;
         let key = key.to_lua(lua)?;
         let value = value.to_lua(lua)?;
@@ -246,13 +249,7 @@ impl<'lua> Table<'lua> {
             lua.push_value(key)?;
             lua.push_value(value)?;
 
-            #[cfg(not(feature = "luau"))]
-            let protect = !lua.unlikely_memory_error();
-            // If Luau table is readonly it will throw an exception
-            #[cfg(feature = "luau")]
-            let protect = !lua.unlikely_memory_error() || self.is_readonly();
-
-            if !protect {
+            if lua.unlikely_memory_error() {
                 ffi::lua_rawset(lua.state, -3);
                 ffi::lua_pop(lua.state, 1);
                 Ok(())
@@ -309,8 +306,12 @@ impl<'lua> Table<'lua> {
 
     /// Appends a value to the back of the table without invoking metamethods.
     pub fn raw_push<V: ToLua<'lua>>(&self, value: V) -> Result<()> {
+        #[cfg(feature = "luau")]
+        self.check_readonly_write()?;
+
         let lua = self.0.lua;
         let value = value.to_lua(lua)?;
+
         unsafe {
             let _sg = StackGuard::new(lua.state);
             check_stack(lua.state, 4)?;
@@ -323,12 +324,7 @@ impl<'lua> Table<'lua> {
                 ffi::lua_rawseti(state, -2, len + 1);
             }
 
-            #[cfg(not(feature = "luau"))]
-            let protect = !lua.unlikely_memory_error();
-            // If Luau table is readonly it will throw an exception
-            #[cfg(feature = "luau")]
-            let protect = !lua.unlikely_memory_error() || self.is_readonly();
-            if !protect {
+            if lua.unlikely_memory_error() {
                 callback(lua.state);
             } else {
                 protect_lua!(lua.state, 2, 0, fn(state) callback(state))?;
@@ -339,6 +335,9 @@ impl<'lua> Table<'lua> {
 
     /// Removes the last element from the table and returns it, without invoking metamethods.
     pub fn raw_pop<V: FromLua<'lua>>(&self) -> Result<V> {
+        #[cfg(feature = "luau")]
+        self.check_readonly_write()?;
+
         let lua = self.0.lua;
         let value = unsafe {
             let _sg = StackGuard::new(lua.state);
@@ -440,6 +439,12 @@ impl<'lua> Table<'lua> {
     /// If `metatable` is `None`, the metatable is removed (if no metatable is set, this does
     /// nothing).
     pub fn set_metatable(&self, metatable: Option<Table<'lua>>) {
+        // Workaround to throw readonly error without returning Result
+        #[cfg(feature = "luau")]
+        if self.is_readonly() {
+            panic!("attempt to modify a readonly table");
+        }
+
         let lua = self.0.lua;
         unsafe {
             let _sg = StackGuard::new(lua.state);
@@ -644,6 +649,16 @@ impl<'lua> Table<'lua> {
             ffi::lua_rawequal(lua.state, -1, -2) != 0
         }
     }
+
+    #[cfg(feature = "luau")]
+    #[inline(always)]
+    pub(crate) fn check_readonly_write(&self) -> Result<()> {
+        if self.is_readonly() {
+            let err = "attempt to modify a readonly table".to_string();
+            return Err(Error::RuntimeError(err));
+        }
+        Ok(())
+    }
 }
 
 impl<'lua> PartialEq for Table<'lua> {
diff --git a/tests/luau.rs b/tests/luau.rs
index 331402f..bb8bff0 100644
--- a/tests/luau.rs
+++ b/tests/luau.rs
@@ -1,7 +1,9 @@
 #![cfg(feature = "luau")]
 
 use std::env;
+use std::fmt::Debug;
 use std::fs;
+use std::panic::{catch_unwind, AssertUnwindSafe};
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 
@@ -75,18 +77,33 @@ fn test_vectors() -> Result<()> {
 fn test_readonly_table() -> Result<()> {
     let lua = Lua::new();
 
-    let t = lua.create_table()?;
+    let t = lua.create_sequence_from([1])?;
     assert!(!t.is_readonly());
     t.set_readonly(true);
     assert!(t.is_readonly());
 
-    match t.set("key", "value") {
-        Err(Error::RuntimeError(err)) if err.contains("attempt to modify a readonly table") => {}
-        r => panic!(
-            "expected RuntimeError(...) with a specific message, got {:?}",
-            r
-        ),
-    };
+    #[track_caller]
+    fn check_readonly_error<T: Debug>(res: Result<T>) {
+        match res {
+            Err(Error::RuntimeError(e)) if e.contains("attempt to modify a readonly table") => {}
+            r => panic!("expected RuntimeError(...) with a specific message, got {r:?}"),
+        }
+    }
+
+    check_readonly_error(t.set("key", "value"));
+    check_readonly_error(t.raw_set("key", "value"));
+    check_readonly_error(t.raw_insert(1, "value"));
+    check_readonly_error(t.raw_remove(1));
+    check_readonly_error(t.push("value"));
+    check_readonly_error(t.pop::<Value>());
+    check_readonly_error(t.raw_push("value"));
+    check_readonly_error(t.raw_pop::<Value>());
+
+    // Special case
+    match catch_unwind(AssertUnwindSafe(|| t.set_metatable(None))) {
+        Ok(_) => panic!("expected panic, got nothing"),
+        Err(_) => {}
+    }
 
     Ok(())
 }