Use protected mode for lua callbacks wherever possible

7 files changed, 61 insertions, 14 deletions

diff --git a/src/lua/include/lua_thread.hpp b/src/lua/include/lua_thread.hpp
index 939d0cda..4c5198aa 100644
--- a/src/lua/include/lua_thread.hpp
+++ b/src/lua/include/lua_thread.hpp
@@ -19,6 +19,8 @@ namespace lua {
 
 class Allocator;
 
+auto CallProtected(lua_State*, int nargs, int nresults) -> int;
+
 class LuaThread {
  public:
   static auto Start(system_fsm::ServiceLocator&, lv_obj_t* lvgl_root = nullptr)
diff --git a/src/lua/lua_database.cpp b/src/lua/lua_database.cpp
index 79916115..4a7c82a5 100644
--- a/src/lua/lua_database.cpp
+++ b/src/lua/lua_database.cpp
@@ -14,6 +14,7 @@
 #include "esp_log.h"
 #include "lauxlib.h"
 #include "lua.h"
+#include "lua_thread.hpp"
 #include "lvgl.h"
 
 #include "database.hpp"
@@ -108,7 +109,7 @@ static auto db_iterate(lua_State* state) -> int {
       } else {
         lua_pushnil(state);
       }
-      lua_call(state, 1, 0);
+      CallProtected(state, 1, 0);
       luaL_unref(state, LUA_REGISTRYINDEX, callback_ref);
     });
   });
@@ -139,7 +140,6 @@ static auto push_iterator(
   lua_pushcclosure(state, db_iterate, 1);
 }
 
-
 static auto record_text(lua_State* state) -> int {
   LuaRecord* data = reinterpret_cast<LuaRecord*>(
       luaL_checkudata(state, 1, kDbRecordMetatable));
diff --git a/src/lua/lua_thread.cpp b/src/lua/lua_thread.cpp
index eb2f5107..e1df3087 100644
--- a/src/lua/lua_thread.cpp
+++ b/src/lua/lua_thread.cpp
@@ -10,9 +10,11 @@
 
 #include "esp_heap_caps.h"
 #include "esp_log.h"
+#include "event_queue.hpp"
 #include "lua.hpp"
 #include "luavgl.h"
 #include "service_locator.hpp"
+#include "ui_events.hpp"
 
 namespace lua {
 
@@ -63,6 +65,7 @@ auto LuaThread::Start(system_fsm::ServiceLocator& services, lv_obj_t* lvgl_root)
 
   // FIXME: luavgl init should probably be a part of the bridge.
   if (lvgl_root) {
+    luavgl_set_pcall(state, CallProtected);
     luavgl_set_root(state, lvgl_root);
     luaL_requiref(state, "lvgl", luaopen_lvgl, true);
     lua_pop(state, 1);
@@ -85,14 +88,46 @@ auto LuaThread::RunScript(const std::string& path) -> bool {
   if (res != LUA_OK) {
     return false;
   }
-  res = lua_pcall(state_, 0, 0, 0);
-  if (res) {
-    const char* msg = lua_tostring(state_, -1);
-    lua_writestring(msg, strlen(msg));
-    lua_writeline();
-    lua_pop(state_, 1);
-  }
+  CallProtected(state_, 0, 0);
   return true;
 }
 
+static int msg_handler(lua_State* L) {
+  if (!lua_isstring(L, 1)) {
+    return 1;
+  }
+
+  const char* msg = lua_tostring(L, 1);
+  if (msg == NULL) {                         /* is error object not a string? */
+    if (luaL_callmeta(L, 1, "__tostring") && /* does it have a metamethod */
+        lua_type(L, -1) == LUA_TSTRING)      /* that produces a string? */
+      return 1;                              /* that is the message */
+    else
+      msg = lua_pushfstring(L, "(error object is a %s value)",
+                            luaL_typename(L, 1));
+  }
+
+  /* append a standard traceback */
+  luaL_traceback(L, L, msg, 1);
+  return 1;
+}
+
+auto CallProtected(lua_State* s, int nargs, int nresults) -> int {
+  int base = lua_gettop(s) - nargs;
+  // Place our message handler under the function to be called.
+  lua_pushcfunction(s, msg_handler);
+  lua_insert(s, base);
+
+  // Invoke the function.
+  int ret = lua_pcall(s, nargs, nresults, base);
+  if (ret != LUA_OK) {
+    events::Ui().Dispatch(ui::OnLuaError{.message = lua_tostring(s, -1)});
+  }
+
+  // Clean up our message handler
+  lua_remove(s, base);
+
+  return ret;
+}
+
 }  // namespace lua
diff --git a/src/lua/property.cpp b/src/lua/property.cpp
index f0383dd8..c63d243f 100644
--- a/src/lua/property.cpp
+++ b/src/lua/property.cpp
@@ -10,6 +10,7 @@
 #include <string>
 
 #include "lua.hpp"
+#include "lua_thread.hpp"
 #include "lvgl.h"
 #include "service_locator.hpp"
 
@@ -49,9 +50,8 @@ static auto property_bind(lua_State* state) -> int {
   // ...and another copy, since we return the original closure.
   lua_pushvalue(state, 2);
 
-  // FIXME: This should ideally be lua_pcall, for safety.
   p->PushValue(*state);
-  lua_call(state, 1, 0);  // Invoke the initial binding.
+  CallProtected(state, 1, 0); // Invoke the initial binding.
 
   lua_pushstring(state, kBindingsTable);
   lua_gettable(state, LUA_REGISTRYINDEX);  // REGISTRY[kBindingsTable]
@@ -229,7 +229,7 @@ auto Property::Update(const LuaValue& v) -> void {
     }
 
     PushValue(*b.first);      // push the argument
-    lua_call(b.first, 1, 0);  // invoke the closure
+    CallProtected(b.first, 1, 0); // invoke the closure
   }
 }
 
diff --git a/src/ui/include/ui_events.hpp b/src/ui/include/ui_events.hpp
index b8dd459c..6a6be304 100644
--- a/src/ui/include/ui_events.hpp
+++ b/src/ui/include/ui_events.hpp
@@ -24,6 +24,10 @@ struct OnStorageChange : tinyfsm::Event {
 
 struct OnSystemError : tinyfsm::Event {};
 
+  struct OnLuaError : tinyfsm::Event {
+    std::string message;
+  };
+
 namespace internal {
 
 struct RecordSelected : tinyfsm::Event {
diff --git a/src/ui/include/ui_fsm.hpp b/src/ui/include/ui_fsm.hpp
index d3ea7eac..9e81259a 100644
--- a/src/ui/include/ui_fsm.hpp
+++ b/src/ui/include/ui_fsm.hpp
@@ -64,6 +64,7 @@ class UiState : public tinyfsm::Fsm<UiState> {
   void react(const audio::QueueUpdate&);
 
   virtual void react(const system_fsm::KeyLockChanged&);
+  virtual void react(const OnLuaError&) {}
 
   virtual void react(const internal::RecordSelected&) {}
   virtual void react(const internal::IndexSelected&) {}
@@ -124,6 +125,8 @@ class Lua : public UiState {
   void entry() override;
   void exit() override;
 
+  void react(const OnLuaError&) override;
+
   void react(const internal::IndexSelected&) override;
   void react(const internal::ShowNowPlaying&) override;
   void react(const internal::ShowSettingsPage&) override;
diff --git a/src/ui/ui_fsm.cpp b/src/ui/ui_fsm.cpp
index ed0624df..5b4ea2a7 100644
--- a/src/ui/ui_fsm.cpp
+++ b/src/ui/ui_fsm.cpp
@@ -238,8 +238,7 @@ auto Lua::PushLuaScreen(lua_State* s) -> int {
 
   // Call the constructor for this screen.
   lua_settop(s, 1);  // Make sure the function is actually at top of stack
-  // FIXME: This should ideally be lua_pcall, for safety.
-  lua_call(s, 0, 1);
+  lua::CallProtected(s, 0, 1);
 
   // Store the reference for the table the constructor returned.
   new_screen->SetObjRef(s);
@@ -262,6 +261,10 @@ void Lua::exit() {
   lv_group_set_default(NULL);
 }
 
+void Lua::react(const OnLuaError& err) {
+  ESP_LOGE("lua", "%s", err.message.c_str());
+}
+
 void Lua::react(const internal::IndexSelected& ev) {
   auto db = sServices->database().lock();
   if (!db) {