abstract_tcp_server2: fix race on shutdown

2 files changed, 9 insertions, 1 deletions

diff --git a/contrib/epee/include/net/abstract_tcp_server2.h b/contrib/epee/include/net/abstract_tcp_server2.h
index 7ca6ac872..b2c05ebb0 100644
--- a/contrib/epee/include/net/abstract_tcp_server2.h
+++ b/contrib/epee/include/net/abstract_tcp_server2.h
@@ -158,6 +158,7 @@ namespace net_utils
     std::list<boost::shared_ptr<connection<t_protocol_handler> > > m_self_refs; // add_ref/release support
     critical_section m_self_refs_lock;
     critical_section m_chunking_lock; // held while we add small chunks of the big do_send() to small do_send_chunk()
+    critical_section m_shutdown_lock; // held while shutting down
     
     t_connection_type m_connection_type;
     
diff --git a/contrib/epee/include/net/abstract_tcp_server2.inl b/contrib/epee/include/net/abstract_tcp_server2.inl
index 5b3550005..59a126163 100644
--- a/contrib/epee/include/net/abstract_tcp_server2.inl
+++ b/contrib/epee/include/net/abstract_tcp_server2.inl
@@ -649,6 +649,10 @@ PRAGMA_WARNING_DISABLE_VS(4355)
   template<class t_protocol_handler>
   bool connection<t_protocol_handler>::shutdown()
   {
+    CRITICAL_REGION_BEGIN(m_shutdown_lock);
+    if (m_was_shutdown)
+      return true;
+    m_was_shutdown = true;
     // Initiate graceful connection closure.
     m_timer.cancel();
     boost::system::error_code ignored_ec;
@@ -658,7 +662,7 @@ PRAGMA_WARNING_DISABLE_VS(4355)
       try { host_count(m_host, -1); } catch (...) { /* ignore */ }
       m_host = "";
     }
-    m_was_shutdown = true;
+    CRITICAL_REGION_END();
     m_protocol_handler.release_protocol();
     return true;
   }
@@ -667,6 +671,9 @@ PRAGMA_WARNING_DISABLE_VS(4355)
   bool connection<t_protocol_handler>::close()
   {
     TRY_ENTRY();
+    auto self = safe_shared_from_this();
+    if(!self)
+      return false;
     //_info("[sock " << socket_.native_handle() << "] Que Shutdown called.");
     m_timer.cancel();
     size_t send_que_size = 0;