From 0cfff3143978aede4f0741df1cd81d5c79e76e3b Mon Sep 17 00:00:00 2001
From: Peter Johnson <johnson.peter@gmail.com>
Date: Thu, 10 Oct 2024 23:03:55 -0600
Subject: [PATCH] [ntcore] Fix use-after-free on connection termination (#7177)

The stream can close (e.g. due to an error) while in the middle of writing. The close callback would immediately destroy the connection object, resulting in the writing code having a use-after-free. Fix this by deferring the deletion to the loop main using a single-shot timer.
---
 ntcore/src/main/native/cpp/NetworkServer.cpp  | 4 +++-
 ntcore/src/main/native/cpp/net/ServerImpl.cpp | 5 ++---
 ntcore/src/main/native/cpp/net/ServerImpl.h   | 2 +-
 3 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/ntcore/src/main/native/cpp/NetworkServer.cpp b/ntcore/src/main/native/cpp/NetworkServer.cpp
index b6d8e0d79f..86d20b454f 100644
--- a/ntcore/src/main/native/cpp/NetworkServer.cpp
+++ b/ntcore/src/main/native/cpp/NetworkServer.cpp
@@ -125,7 +125,9 @@ void NetworkServer::ServerConnection::UpdateOutgoingTimer(uint32_t repeatMs) {
 void NetworkServer::ServerConnection::ConnectionClosed() {
   // don't call back into m_server if it's being destroyed
   if (!m_outgoingTimer->IsLoopClosing()) {
-    m_server.m_serverImpl.RemoveClient(m_clientId);
+    uv::Timer::SingleShot(m_outgoingTimer->GetLoopRef(), uv::Timer::Time{0},
+                          [client = m_server.m_serverImpl.RemoveClient(
+                               m_clientId)]() mutable { client.reset(); });
     m_server.RemoveConnection(this);
   }
   m_outgoingTimer->Close();
diff --git a/ntcore/src/main/native/cpp/net/ServerImpl.cpp b/ntcore/src/main/native/cpp/net/ServerImpl.cpp
index bbbc85d0a0..96eedf5ed0 100644
--- a/ntcore/src/main/native/cpp/net/ServerImpl.cpp
+++ b/ntcore/src/main/native/cpp/net/ServerImpl.cpp
@@ -1252,7 +1252,7 @@ int ServerImpl::AddClient3(std::string_view connInfo, bool local,
   return index;
 }
 
-void ServerImpl::RemoveClient(int clientId) {
+std::shared_ptr<void> ServerImpl::RemoveClient(int clientId) {
   DEBUG3("RemoveClient({})", clientId);
   auto& client = m_clients[clientId];
 
@@ -1288,8 +1288,7 @@ void ServerImpl::RemoveClient(int clientId) {
   DeleteTopic(client->m_metaPub);
   DeleteTopic(client->m_metaSub);
 
-  // delete the client
-  client.reset();
+  return std::move(client);
 }
 
 bool ServerImpl::PersistentChanged() {
diff --git a/ntcore/src/main/native/cpp/net/ServerImpl.h b/ntcore/src/main/native/cpp/net/ServerImpl.h
index d84a06ce29..937ba0e533 100644
--- a/ntcore/src/main/native/cpp/net/ServerImpl.h
+++ b/ntcore/src/main/native/cpp/net/ServerImpl.h
@@ -80,7 +80,7 @@ class ServerImpl final {
   int AddClient3(std::string_view connInfo, bool local,
                  net3::WireConnection3& wire, Connected3Func connected,
                  SetPeriodicFunc setPeriodic);
-  void RemoveClient(int clientId);
+  std::shared_ptr<void> RemoveClient(int clientId);
 
   void ConnectionsChanged(const std::vector<ConnectionInfo>& conns);