From 186f2edf0fd7f560d6afe4268efd6391745d3f37 Mon Sep 17 00:00:00 2001
From: Stephen Seo <seo.disparate@gmail.com>
Date: Thu, 11 Jan 2024 19:00:27 +0900
Subject: [PATCH] Use mutex when enabling/disabling threaded-update

---
 src/UDPC_Defines.hpp  |  2 ++
 src/UDPConnection.cpp | 27 +++++++++++++++++++++++----
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/src/UDPC_Defines.hpp b/src/UDPC_Defines.hpp
index 0104822..4a6df9f 100644
--- a/src/UDPC_Defines.hpp
+++ b/src/UDPC_Defines.hpp
@@ -276,6 +276,8 @@ public:
     std::mutex atostrBufIndexMutex;
     std::uint32_t atostrBufIndex;
 
+    std::mutex setThreadedUpdateMutex;
+
 }; // struct Context
 
 Context *verifyContext(UDPC_HContext ctx);
diff --git a/src/UDPConnection.cpp b/src/UDPConnection.cpp
index d5c358b..1c049b7 100644
--- a/src/UDPConnection.cpp
+++ b/src/UDPConnection.cpp
@@ -257,7 +257,8 @@ peerPKWhitelistMutex(),
 threadedSleepTime(std::chrono::milliseconds(UDPC_UPDATE_MS_DEFAULT)),
 keysSet(),
 atostrBufIndexMutex(),
-atostrBufIndex(0)
+atostrBufIndex(0),
+setThreadedUpdateMutex()
 {
     std::memset(atostrBuf, 0, UDPC_ATOSTR_SIZE);
 
@@ -2207,7 +2208,13 @@ UDPC_HContext UDPC_init_threaded_update_ms(
 
 int UDPC_enable_threaded_update(UDPC_HContext ctx) {
     UDPC::Context *c = UDPC::verifyContext(ctx);
-    if(!c || c->isAutoUpdating.load() || c->thread.joinable()) {
+    if (!c) {
+        return 0;
+    }
+
+    std::lock_guard<std::mutex> setThreadedLock(c->setThreadedUpdateMutex);
+
+    if(c->isAutoUpdating.load() || c->thread.joinable()) {
         return 0;
     }
 
@@ -2222,7 +2229,13 @@ int UDPC_enable_threaded_update(UDPC_HContext ctx) {
 
 int UDPC_enable_threaded_update_ms(UDPC_HContext ctx, int updateMS) {
     UDPC::Context *c = UDPC::verifyContext(ctx);
-    if(!c || c->isAutoUpdating.load() || c->thread.joinable()) {
+    if (!c) {
+        return 0;
+    }
+
+    std::lock_guard<std::mutex> setThreadedLock(c->setThreadedUpdateMutex);
+
+    if(c->isAutoUpdating.load() || c->thread.joinable()) {
         return 0;
     }
 
@@ -2243,7 +2256,13 @@ int UDPC_enable_threaded_update_ms(UDPC_HContext ctx, int updateMS) {
 
 int UDPC_disable_threaded_update(UDPC_HContext ctx) {
     UDPC::Context *c = UDPC::verifyContext(ctx);
-    if(!c || !c->isAutoUpdating.load() || !c->thread.joinable()) {
+    if (!c) {
+        return 0;
+    }
+
+    std::lock_guard<std::mutex> setThreadedLock(c->setThreadedUpdateMutex);
+
+    if(!c->isAutoUpdating.load() || !c->thread.joinable()) {
         return 0;
     }