From 7f11260bf03d5c9f2129affde6c14ca5e3bfb382 Mon Sep 17 00:00:00 2001
From: Bartosz Taudul <wolf.pld@gmail.com>
Date: Tue, 12 Feb 2019 01:41:09 +0100
Subject: [PATCH] Handle dropped connection during handshake.

---
 common/TracyProtocol.hpp |  3 ++-
 server/TracyView.cpp     | 30 ++++++++++++++++++++++++++----
 server/TracyWorker.cpp   |  6 +++++-
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/common/TracyProtocol.hpp b/common/TracyProtocol.hpp
index ed707f97..1bbebc69 100644
--- a/common/TracyProtocol.hpp
+++ b/common/TracyProtocol.hpp
@@ -37,7 +37,8 @@ enum HandshakeStatus : uint8_t
     HandshakePending,
     HandshakeWelcome,
     HandshakeProtocolMismatch,
-    HandshakeNotAvailable
+    HandshakeNotAvailable,
+    HandshakeDropped
 };
 
 enum { WelcomeMessageProgramNameSize = 64 };
diff --git a/server/TracyView.cpp b/server/TracyView.cpp
index 95c13fa1..ba508f3f 100644
--- a/server/TracyView.cpp
+++ b/server/TracyView.cpp
@@ -515,13 +515,19 @@ void View::DrawTextContrast( ImDrawList* draw, const ImVec2& pos, uint32_t color
 bool View::Draw()
 {
     HandshakeStatus status = (HandshakeStatus)s_instance->m_worker.GetHandshakeStatus();
-    if( status == HandshakeProtocolMismatch )
+    switch( status )
     {
+    case HandshakeProtocolMismatch:
         ImGui::OpenPopup( "Protocol mismatch" );
-    }
-    else if( status == HandshakeNotAvailable )
-    {
+        break;
+    case HandshakeNotAvailable:
         ImGui::OpenPopup( "Client not ready" );
+        break;
+    case HandshakeDropped:
+        ImGui::OpenPopup( "Client disconnected" );
+        break;
+    default:
+        break;
     }
 
     const auto& failure = s_instance->m_worker.GetFailureType();
@@ -562,6 +568,22 @@ bool View::Draw()
         ImGui::EndPopup();
     }
 
+    if( ImGui::BeginPopupModal( "Client disconnected", nullptr, ImGuiWindowFlags_AlwaysAutoResize ) )
+    {
+#ifdef TRACY_EXTENDED_FONT
+        TextCentered( ICON_FA_HANDSHAKE );
+#endif
+        ImGui::TextUnformatted( "The client you are trying to connect to has disconnected during the initial\nconnection handshake. Please check your network configuration." );
+        ImGui::Separator();
+        if( ImGui::Button( "Will do" ) )
+        {
+            ImGui::CloseCurrentPopup();
+            ImGui::EndPopup();
+            return false;
+        }
+        ImGui::EndPopup();
+    }
+
     if( ImGui::BeginPopupModal( "Instrumentation failure", nullptr, ImGuiWindowFlags_AlwaysAutoResize ) )
     {
         const auto& data = s_instance->m_worker.GetFailureData();
diff --git a/server/TracyWorker.cpp b/server/TracyWorker.cpp
index c9c179fc..33da2e47 100644
--- a/server/TracyWorker.cpp
+++ b/server/TracyWorker.cpp
@@ -1503,7 +1503,11 @@ void Worker::Exec()
     uint32_t protocolVersion = ProtocolVersion;
     m_sock.Send( &protocolVersion, sizeof( protocolVersion ) );
     HandshakeStatus handshake;
-    if( !m_sock.Read( &handshake, sizeof( handshake ), 10, ShouldExit ) ) goto close;
+    if( !m_sock.Read( &handshake, sizeof( handshake ), 10, ShouldExit ) )
+    {
+        m_handshake.store( HandshakeDropped, std::memory_order_relaxed );
+        goto close;
+    }
     m_handshake.store( handshake, std::memory_order_relaxed );
     switch( handshake )
     {