1 files changed, 212 insertions, 8 deletions

diff --git a/src/Bindings/LuaState.cpp b/src/Bindings/LuaState.cpp
index 200878cf7..98507e96c 100644
--- a/src/Bindings/LuaState.cpp
+++ b/src/Bindings/LuaState.cpp
@@ -20,6 +20,10 @@ extern "C"
 #include "../Entities/Entity.h"
 #include "../BlockEntities/BlockEntity.h"
 
+
+
+
+
 // fwd: "SQLite/lsqlite3.c"
 extern "C"
 {
@@ -39,6 +43,10 @@ extern "C"
 
 const cLuaState::cRet cLuaState::Return = {};
 
+/** Each Lua state stores a pointer to its creating cLuaState in Lua globals, under this name.
+This way any cLuaState can reference the main cLuaState's TrackedCallbacks, mutex etc. */
+static const char * g_CanonLuaStateGlobalName = "_CuberiteInternal_CanonLuaState";
+
 
 
 
@@ -114,6 +122,101 @@ cLuaStateTracker & cLuaStateTracker::Get(void)
 
 
 ////////////////////////////////////////////////////////////////////////////////
+// cLuaState::cCallback:
+
+bool cLuaState::cCallback::RefStack(cLuaState & a_LuaState, int a_StackPos)
+{
+	// Check if the stack contains a function:
+	if (!lua_isfunction(a_LuaState, a_StackPos))
+	{
+		return false;
+	}
+
+	// Clear any previous callback:
+	Clear();
+
+	// Add self to LuaState's callback-tracking:
+	a_LuaState.TrackCallback(*this);
+
+	// Store the new callback:
+	cCSLock Lock(m_CS);
+	m_Ref.RefStack(a_LuaState, a_StackPos);
+	return true;
+}
+
+
+
+
+
+void cLuaState::cCallback::Clear(void)
+{
+	// Free the callback reference:
+	lua_State * luaState = nullptr;
+	{
+		cCSLock Lock(m_CS);
+		if (!m_Ref.IsValid())
+		{
+			return;
+		}
+		luaState = m_Ref.GetLuaState();
+		m_Ref.UnRef();
+	}
+
+	// Remove from LuaState's callback-tracking:
+	cLuaState(luaState).UntrackCallback(*this);
+}
+
+
+
+
+
+bool cLuaState::cCallback::IsValid(void)
+{
+	cCSLock lock(m_CS);
+	return m_Ref.IsValid();
+}
+
+
+
+
+
+bool cLuaState::cCallback::IsSameLuaState(cLuaState & a_LuaState)
+{
+	cCSLock lock(m_CS);
+	if (!m_Ref.IsValid())
+	{
+		return false;
+	}
+	auto canonState = a_LuaState.QueryCanonLuaState();
+	if (canonState == nullptr)
+	{
+		return false;
+	}
+	return (m_Ref.GetLuaState() == static_cast<lua_State *>(*canonState));
+}
+
+
+
+
+
+void cLuaState::cCallback::Invalidate(void)
+{
+	cCSLock Lock(m_CS);
+	if (!m_Ref.IsValid())
+	{
+		LOGD("%s: Invalidating an already invalid callback at %p, this should not happen",
+			__FUNCTION__, reinterpret_cast<void *>(this)
+		);
+		return;
+	}
+	m_Ref.UnRef();
+}
+
+
+
+
+
+////////////////////////////////////////////////////////////////////////////////
 // cLuaState:
 
 cLuaState::cLuaState(const AString & a_SubsystemName) :
@@ -170,6 +273,10 @@ void cLuaState::Create(void)
 	luaL_openlibs(m_LuaState);
 	m_IsOwned = true;
 	cLuaStateTracker::Add(*this);
+
+	// Add the CanonLuaState value into the Lua state, so that we can get it from anywhere:
+	lua_pushlightuserdata(m_LuaState, reinterpret_cast<void *>(this));
+	lua_setglobal(m_LuaState, g_CanonLuaStateGlobalName);
 }
 
 
@@ -206,6 +313,16 @@ void cLuaState::Close(void)
 		Detach();
 		return;
 	}
+
+	// Invalidate all callbacks:
+	{
+		cCSLock Lock(m_CSTrackedCallbacks);
+		for (auto & c: m_TrackedCallbacks)
+		{
+			c->Invalidate();
+		}
+	}
+
 	cLuaStateTracker::Del(*this);
 	lua_close(m_LuaState);
 	m_LuaState = nullptr;
@@ -821,6 +938,18 @@ void cLuaState::Push(std::chrono::milliseconds a_Value)
 
 
 
+void cLuaState::Pop(int a_NumValuesToPop)
+{
+	ASSERT(IsValid());
+
+	lua_pop(m_LuaState, a_NumValuesToPop);
+	m_NumCurrentFunctionArgs -= a_NumValuesToPop;
+}
+
+
+
+
+
 bool cLuaState::GetStackValue(int a_StackPos, AString & a_Value)
 {
 	size_t len = 0;
@@ -847,6 +976,24 @@ bool cLuaState::GetStackValue(int a_StackPos, bool & a_ReturnedVal)
 
 
 
+bool cLuaState::GetStackValue(int a_StackPos, cCallback & a_Callback)
+{
+	return a_Callback.RefStack(*this, a_StackPos);
+}
+
+
+
+
+
+bool cLuaState::GetStackValue(int a_StackPos, cCallbackPtr & a_Callback)
+{
+	return a_Callback->RefStack(*this, a_StackPos);
+}
+
+
+
+
+
 bool cLuaState::GetStackValue(int a_StackPos, cPluginManager::CommandResult & a_Result)
 {
 	if (lua_isnumber(m_LuaState, a_StackPos))
@@ -1531,7 +1678,7 @@ int cLuaState::CopyStackFrom(cLuaState & a_SrcLuaState, int a_SrcStart, int a_Sr
 				LOGWARNING("%s: Unsupported value: '%s' at stack position %d. Can only copy numbers, strings, bools and classes!",
 					__FUNCTION__, lua_typename(a_SrcLuaState, t), i
 				);
-				a_SrcLuaState.LogStack("Stack where copying failed:");
+				a_SrcLuaState.LogStackValues("Stack where copying failed:");
 				lua_pop(m_LuaState, i - a_SrcStart);
 				return -1;
 			}
@@ -1558,16 +1705,16 @@ void cLuaState::ToString(int a_StackPos, AString & a_String)
 
 
 
-void cLuaState::LogStack(const char * a_Header)
+void cLuaState::LogStackValues(const char * a_Header)
 {
-	LogStack(m_LuaState, a_Header);
+	LogStackValues(m_LuaState, a_Header);
 }
 
 
 
 
 
-void cLuaState::LogStack(lua_State * a_LuaState, const char * a_Header)
+void cLuaState::LogStackValues(lua_State * a_LuaState, const char * a_Header)
 {
 	// Format string consisting only of %s is used to appease the compiler
 	LOG("%s", (a_Header != nullptr) ? a_Header : "Lua C API Stack contents:");
@@ -1592,6 +1739,21 @@ void cLuaState::LogStack(lua_State * a_LuaState, const char * a_Header)
 
 
 
+cLuaState * cLuaState::QueryCanonLuaState(void)
+{
+	// Get the CanonLuaState global from Lua:
+	auto cb = WalkToNamedGlobal(g_CanonLuaStateGlobalName);
+	if (!cb.IsValid())
+	{
+		return nullptr;
+	}
+	return reinterpret_cast<cLuaState *>(lua_touserdata(m_LuaState, -1));
+}
+
+
+
+
+
 int cLuaState::ReportFnCallErrors(lua_State * a_LuaState)
 {
 	LOGWARNING("LUA: %s", lua_tostring(a_LuaState, -1));
@@ -1626,6 +1788,50 @@ int cLuaState::BreakIntoDebugger(lua_State * a_LuaState)
 
 
 
+void cLuaState::TrackCallback(cCallback & a_Callback)
+{
+	// Get the CanonLuaState global from Lua:
+	auto canonState = QueryCanonLuaState();
+	if (canonState == nullptr)
+	{
+		LOGWARNING("%s: Lua state %p has invalid CanonLuaState!", __FUNCTION__, reinterpret_cast<void *>(m_LuaState));
+		return;
+	}
+
+	// Add the callback:
+	cCSLock Lock(canonState->m_CSTrackedCallbacks);
+	canonState->m_TrackedCallbacks.push_back(&a_Callback);
+}
+
+
+
+
+
+void cLuaState::UntrackCallback(cCallback & a_Callback)
+{
+	// Get the CanonLuaState global from Lua:
+	auto canonState = QueryCanonLuaState();
+	if (canonState == nullptr)
+	{
+		LOGWARNING("%s: Lua state %p has invalid CanonLuaState!", __FUNCTION__, reinterpret_cast<void *>(m_LuaState));
+		return;
+	}
+
+	// Remove the callback:
+	cCSLock Lock(canonState->m_CSTrackedCallbacks);
+	auto & trackedCallbacks = canonState->m_TrackedCallbacks;
+	trackedCallbacks.erase(std::remove_if(trackedCallbacks.begin(), trackedCallbacks.end(),
+		[&a_Callback](cCallback * a_StoredCallback)
+		{
+			return (a_StoredCallback == &a_Callback);
+		}
+	));
+}
+
+
+
+
+
 ////////////////////////////////////////////////////////////////////////////////
 // cLuaState::cRef:
 
@@ -1681,7 +1887,7 @@ void cLuaState::cRef::RefStack(cLuaState & a_LuaState, int a_StackPos)
 	{
 		UnRef();
 	}
-	m_LuaState = &a_LuaState;
+	m_LuaState = a_LuaState;
 	lua_pushvalue(a_LuaState, a_StackPos);  // Push a copy of the value at a_StackPos onto the stack
 	m_Ref = luaL_ref(a_LuaState, LUA_REGISTRYINDEX);
 }
@@ -1692,11 +1898,9 @@ void cLuaState::cRef::RefStack(cLuaState & a_LuaState, int a_StackPos)
 
 void cLuaState::cRef::UnRef(void)
 {
-	ASSERT(m_LuaState->IsValid());  // The reference should be destroyed before destroying the LuaState
-
 	if (IsValid())
 	{
-		luaL_unref(*m_LuaState, LUA_REGISTRYINDEX, m_Ref);
+		luaL_unref(m_LuaState, LUA_REGISTRYINDEX, m_Ref);
 	}
 	m_LuaState = nullptr;
 	m_Ref = LUA_REFNIL;