SafeThread: Avoid use-after-free risk in thread shutdown (#1355)

Use shared_ptr to keep data alive until the thread has terminated.

cscore/src/main/native/cpp/MjpegServerImpl.cpp

@@ -865,10 +865,7 @@ void MjpegServerImpl::ServerThreadMain() {
     }
 
     // Start it if not already started
-    {
-      auto thr = it->GetThread();
-      if (!thr) it->Start(new ConnThread{GetName()});
-    }
+    it->Start(GetName());
 
     auto nstreams =
         std::count_if(m_connThreads.begin(), m_connThreads.end(),

cscore/src/main/native/cpp/Notifier.cpp

@@ -94,10 +94,7 @@ Notifier::Notifier() { s_destroyed = false; }
 
 Notifier::~Notifier() { s_destroyed = true; }
 
-void Notifier::Start() {
-  auto thr = m_owner.GetThread();
-  if (!thr) m_owner.Start(new Thread(m_on_start, m_on_exit));
-}
+void Notifier::Start() { m_owner.Start(m_on_start, m_on_exit); }
 
 void Notifier::Stop() { m_owner.Stop(); }
 

cscore/src/main/native/cpp/Telemetry.cpp

@@ -45,10 +45,7 @@ Telemetry::Telemetry() {}
 
 Telemetry::~Telemetry() {}
 
-void Telemetry::Start() {
-  auto thr = m_owner.GetThread();
-  if (!thr) m_owner.Start(new Thread);
-}
+void Telemetry::Start() { m_owner.Start(); }
 
 void Telemetry::Stop() { m_owner.Stop(); }
 

ntcore/src/main/native/cpp/CallbackManager.h

@@ -294,8 +294,7 @@ class CallbackManager {
  protected:
   template <typename... Args>
   void DoStart(Args&&... args) {
-    auto thr = m_owner.GetThread();
-    if (!thr) m_owner.Start(new Thread(std::forward<Args>(args)...));
+    m_owner.Start(std::forward<Args>(args)...);
   }
 
   template <typename... Args>

ntcore/src/main/native/cpp/DsClient.cpp

@@ -36,7 +36,7 @@ DsClient::DsClient(Dispatcher& dispatcher, wpi::Logger& logger)
 void DsClient::Start(unsigned int port) {
   auto thr = m_owner.GetThread();
   if (!thr)
-    m_owner.Start(new Thread(m_dispatcher, m_logger, port));
+    m_owner.Start(m_dispatcher, m_logger, port);
   else
     thr->m_port = port;
 }

wpiutil/src/main/native/cpp/EventLoopRunner.cpp

@@ -39,7 +39,7 @@ class EventLoopRunner::Thread : public SafeThread {
   std::weak_ptr<UvExecFunc> m_doExec;
 };
 
-EventLoopRunner::EventLoopRunner() { m_owner.Start(new Thread); }
+EventLoopRunner::EventLoopRunner() { m_owner.Start(); }
 
 EventLoopRunner::~EventLoopRunner() {
   ExecAsync([](uv::Loop& loop) {

wpiutil/src/main/native/cpp/SafeThread.cpp

@@ -0,0 +1,65 @@
+/*----------------------------------------------------------------------------*/
+/* Copyright (c) 2015-2018 FIRST. All Rights Reserved.                        */
+/* Open Source Software - may be modified and shared by FRC teams. The code   */
+/* must be accompanied by the FIRST BSD license file in the root directory of */
+/* the project.                                                               */
+/*----------------------------------------------------------------------------*/
+
+#include "wpi/SafeThread.h"
+
+using namespace wpi;
+
+detail::SafeThreadProxyBase::SafeThreadProxyBase(
+    std::shared_ptr<SafeThread> thr)
+    : m_thread(std::move(thr)) {
+  if (!m_thread) return;
+  m_lock = std::unique_lock<wpi::mutex>(m_thread->m_mutex);
+  if (!m_thread->m_active) {
+    m_lock.unlock();
+    m_thread = nullptr;
+    return;
+  }
+}
+
+void detail::SafeThreadOwnerBase::Start(std::shared_ptr<SafeThread> thr) {
+  std::lock_guard<wpi::mutex> lock(m_mutex);
+  if (auto thr = m_thread.lock()) return;
+  std::thread stdThread([=] { thr->Main(); });
+  m_thread = thr;
+  m_nativeHandle = stdThread.native_handle();
+  stdThread.detach();
+}
+
+void detail::SafeThreadOwnerBase::Stop() {
+  std::lock_guard<wpi::mutex> lock(m_mutex);
+  if (auto thr = m_thread.lock()) {
+    thr->m_active = false;
+    thr->m_cond.notify_one();
+  }
+}
+
+void detail::swap(SafeThreadOwnerBase& lhs, SafeThreadOwnerBase& rhs) noexcept {
+  using std::swap;
+  if (&lhs == &rhs) return;
+  std::lock(lhs.m_mutex, rhs.m_mutex);
+  std::lock_guard<wpi::mutex> lock_lhs(lhs.m_mutex, std::adopt_lock);
+  std::lock_guard<wpi::mutex> lock_rhs(rhs.m_mutex, std::adopt_lock);
+  std::swap(lhs.m_thread, rhs.m_thread);
+  std::swap(lhs.m_nativeHandle, rhs.m_nativeHandle);
+}
+
+detail::SafeThreadOwnerBase::operator bool() const {
+  std::lock_guard<wpi::mutex> lock(m_mutex);
+  return !m_thread.expired();
+}
+
+std::thread::native_handle_type
+detail::SafeThreadOwnerBase::GetNativeThreadHandle() const {
+  std::lock_guard<wpi::mutex> lock(m_mutex);
+  return m_nativeHandle;
+}
+
+std::shared_ptr<SafeThread> detail::SafeThreadOwnerBase::GetThread() const {
+  std::lock_guard<wpi::mutex> lock(m_mutex);
+  return m_thread.lock();
+}

wpiutil/src/main/native/include/wpi/SafeThread.h

@@ -9,7 +9,9 @@
 #define WPIUTIL_WPI_SAFETHREAD_H_
 
 #include <atomic>
+#include <memory>
 #include <thread>
+#include <utility>
 
 #include "wpi/condition_variable.h"
 #include "wpi/mutex.h"
@@ -23,7 +25,7 @@ class SafeThread {
   virtual ~SafeThread() = default;
   virtual void Main() = 0;
 
-  wpi::mutex m_mutex;
+  mutable wpi::mutex m_mutex;
   std::atomic_bool m_active;
   wpi::condition_variable m_cond;
 };
@@ -33,20 +35,12 @@ namespace detail {
 // Non-template proxy base class for common proxy code.
 class SafeThreadProxyBase {
  public:
-  explicit SafeThreadProxyBase(SafeThread* thr) : m_thread(thr) {
-    if (!m_thread) return;
-    m_lock = std::unique_lock<wpi::mutex>(m_thread->m_mutex);
-    if (!m_thread->m_active) {
-      m_lock.unlock();
-      m_thread = nullptr;
-      return;
-    }
-  }
+  explicit SafeThreadProxyBase(std::shared_ptr<SafeThread> thr);
   explicit operator bool() const { return m_thread != nullptr; }
   std::unique_lock<wpi::mutex>& GetLock() { return m_lock; }
 
  protected:
-  SafeThread* m_thread;
+  std::shared_ptr<SafeThread> m_thread;
   std::unique_lock<wpi::mutex> m_lock;
 };
 
@@ -55,9 +49,10 @@ class SafeThreadProxyBase {
 template <typename T>
 class SafeThreadProxy : public SafeThreadProxyBase {
  public:
-  explicit SafeThreadProxy(SafeThread* thr) : SafeThreadProxyBase(thr) {}
-  T& operator*() const { return *static_cast<T*>(m_thread); }
-  T* operator->() const { return static_cast<T*>(m_thread); }
+  explicit SafeThreadProxy(std::shared_ptr<SafeThread> thr)
+      : SafeThreadProxyBase(std::move(thr)) {}
+  T& operator*() const { return *static_cast<T*>(m_thread.get()); }
+  T* operator->() const { return static_cast<T*>(m_thread.get()); }
 };
 
 // Non-template owner base class for common owner code.
@@ -65,62 +60,47 @@ class SafeThreadOwnerBase {
  public:
   void Stop();
 
-  SafeThreadOwnerBase() { m_thread = nullptr; }
+  SafeThreadOwnerBase() noexcept = default;
   SafeThreadOwnerBase(const SafeThreadOwnerBase&) = delete;
   SafeThreadOwnerBase& operator=(const SafeThreadOwnerBase&) = delete;
-  SafeThreadOwnerBase(SafeThreadOwnerBase&& other)
-      : m_thread(other.m_thread.exchange(nullptr)) {}
-  SafeThreadOwnerBase& operator=(SafeThreadOwnerBase other) {
-    SafeThread* otherthr = other.m_thread.exchange(nullptr);
-    SafeThread* curthr = m_thread.exchange(otherthr);
-    other.m_thread.exchange(curthr);  // other destructor will clean up
+  SafeThreadOwnerBase(SafeThreadOwnerBase&& other) noexcept
+      : SafeThreadOwnerBase() {
+    swap(*this, other);
+  }
+  SafeThreadOwnerBase& operator=(SafeThreadOwnerBase other) noexcept {
+    swap(*this, other);
     return *this;
   }
   ~SafeThreadOwnerBase() { Stop(); }
 
-  explicit operator bool() const { return m_thread.load(); }
+  friend void swap(SafeThreadOwnerBase& lhs, SafeThreadOwnerBase& rhs) noexcept;
+
+  explicit operator bool() const;
+
+  std::thread::native_handle_type GetNativeThreadHandle() const;
 
  protected:
-  void Start(SafeThread* thr);
-  SafeThread* GetThread() const { return m_thread.load(); }
-  std::thread::native_handle_type GetNativeThreadHandle() const {
-    return m_nativeHandle;
-  }
+  void Start(std::shared_ptr<SafeThread> thr);
+  std::shared_ptr<SafeThread> GetThread() const;
 
  private:
-  std::atomic<SafeThread*> m_thread;
-  std::atomic<std::thread::native_handle_type> m_nativeHandle;
+  mutable wpi::mutex m_mutex;
+  std::weak_ptr<SafeThread> m_thread;
+  std::thread::native_handle_type m_nativeHandle;
 };
 
-inline void SafeThreadOwnerBase::Start(SafeThread* thr) {
-  SafeThread* curthr = nullptr;
-  SafeThread* newthr = thr;
-  if (!m_thread.compare_exchange_strong(curthr, newthr)) {
-    delete newthr;
-    return;
-  }
-  std::thread stdThread([=]() {
-    newthr->Main();
-    delete newthr;
-  });
-  m_nativeHandle = stdThread.native_handle();
-  stdThread.detach();
-}
-
-inline void SafeThreadOwnerBase::Stop() {
-  SafeThread* thr = m_thread.exchange(nullptr);
-  if (!thr) return;
-  thr->m_active = false;
-  thr->m_cond.notify_one();
-}
+void swap(SafeThreadOwnerBase& lhs, SafeThreadOwnerBase& rhs) noexcept;
 
 }  // namespace detail
 
 template <typename T>
 class SafeThreadOwner : public detail::SafeThreadOwnerBase {
  public:
-  void Start() { Start(new T); }
-  void Start(T* thr) { detail::SafeThreadOwnerBase::Start(thr); }
+  template <typename... Args>
+  void Start(Args&&... args) {
+    detail::SafeThreadOwnerBase::Start(
+        std::make_shared<T>(std::forward<Args>(args)...));
+  }
 
   using Proxy = typename detail::SafeThreadProxy<T>;
   Proxy GetThread() const {