From 4af78281cfb30f94d1fe48103973c96da099adc5 Mon Sep 17 00:00:00 2001
From: sukibaby <163092272+sukibaby@users.noreply.github.com>
Date: Mon, 16 Dec 2024 07:28:15 -0800
Subject: [PATCH] Properly destroy closed WebSocketHandles

---
 src/NetworkManager.cpp | 22 ++++++++++++++++++++--
 src/NetworkManager.h   |  5 ++++-
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/NetworkManager.cpp b/src/NetworkManager.cpp
index 4c167d3c5a..1e9eda2607 100644
--- a/src/NetworkManager.cpp
+++ b/src/NetworkManager.cpp
@@ -100,6 +100,12 @@ NetworkManager::~NetworkManager()
 	// Unregister with Lua.
 	LUA->UnsetGlobal("NETWORK");
 
+	// Close all WebSocket connections
+	for (auto& handle : webSocketHandles) {
+		handle->webSocket.stop();
+	}
+
+	// Set the status to uninitialized.
 	ix::uninitNetSystem();
 }
 
@@ -262,6 +268,8 @@ WebSocketHandlePtr NetworkManager::WebSocket(const WebSocketArgs& args)
 
 	handle->webSocket.start();
 
+	webSocketHandles.push_back(handle);
+
 	return handle;
 }
 
@@ -316,6 +324,10 @@ int HttpRequestFuture::Cancel(lua_State *L)
 	return 0;
 }
 
+WebSocketHandle::~WebSocketHandle() {
+	webSocket.stop();
+}
+
 int WebSocketHandle::Collect(lua_State *L)
 {
 	void *udata = luaL_checkudata(L, 1, "WebSocketHandle");
@@ -328,11 +340,17 @@ int WebSocketHandle::Close(lua_State *L)
 {
 	void *udata = luaL_checkudata(L, 1, "WebSocketHandle");
 	auto handle = *static_cast<WebSocketHandlePtr*>(udata);
+
 	LUA->YieldLua();
 	handle->webSocket.stop();
-	handle->onClose();
 	LUA->UnyieldLua();
-	return 0;
+
+	if (handle->onClose)
+	{
+		handle->onClose();
+	}
+
+	return 1;
 }
 
 int WebSocketHandle::Send(lua_State *L)
diff --git a/src/NetworkManager.h b/src/NetworkManager.h
index 5c1cca3a84..695b7fbef3 100644
--- a/src/NetworkManager.h
+++ b/src/NetworkManager.h
@@ -110,7 +110,8 @@ class WebSocketHandle
 {
 public:
 	WebSocketHandle() {};
-
+	~WebSocketHandle();
+	
 	static int Collect(lua_State *L);
 	static int Close(lua_State *L);
 	static int Send(lua_State *L);
@@ -146,6 +147,8 @@ private:
 
 	static Preference<bool> httpEnabled;
 	static Preference<RString> httpAllowHosts;
+
+	std::vector<std::shared_ptr<WebSocketHandle>> webSocketHandles;
 };
 
 extern NetworkManager*	NETWORK;