From ee22c0f714bfd24485259583acfea1235951cf5a Mon Sep 17 00:00:00 2001
From: Benau <Benau@users.noreply.github.com>
Date: Fri, 28 Feb 2020 22:49:06 +0800
Subject: [PATCH] Stop STKHost if cannot register public address

---
 src/network/protocols/connect_to_server.cpp | 18 +++++++++++-------
 src/network/protocols/connect_to_server.hpp |  2 +-
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/network/protocols/connect_to_server.cpp b/src/network/protocols/connect_to_server.cpp
index e03c84ff4..9a120685c 100644
--- a/src/network/protocols/connect_to_server.cpp
+++ b/src/network/protocols/connect_to_server.cpp
@@ -244,12 +244,15 @@ void ConnectToServer::asynchronousUpdate()
             {
                 STKHost::get()->setPublicAddress(
                     m_server->useIPV6Connection() ? AF_INET6 : AF_INET);
-                if (!STKHost::get()->getValidPublicAddress().empty())
-                    registerWithSTKServer();
+                if (STKHost::get()->getValidPublicAddress().empty() ||
+                    registerWithSTKServer() == false)
+                {
+                    // Set to DONE will stop STKHost is not connected
+                    m_state = DONE;
+                    break;
+                }
             }
-            // Set to DONE will stop STKHost is not connected
-            m_state = STKHost::get()->getValidPublicAddress().empty() ?
-                DONE : GOT_SERVER_ADDRESS;
+            m_state = GOT_SERVER_ADDRESS;
             break;
         }
         case GOT_SERVER_ADDRESS:
@@ -467,7 +470,7 @@ bool ConnectToServer::tryConnect(int timeout, int retry, bool another_port,
 // ----------------------------------------------------------------------------
 /** Register this client with the STK server.
  */
-void ConnectToServer::registerWithSTKServer()
+bool ConnectToServer::registerWithSTKServer()
 {
     // Our public address is now known, register details with
     // STK server
@@ -498,13 +501,14 @@ void ConnectToServer::registerWithSTKServer()
     if(result->get("success", &success) && success == "yes")
     {
         Log::debug("ConnectToServer", "Address registered successfully.");
+        return true;
     }
     else
     {
         irr::core::stringc error(request->getInfo().c_str());
         Log::error("ConnectToServer", "Failed to register client address: %s",
             error.c_str());
-        m_state = DONE;
+        return false;
     }
 }   // registerWithSTKServer
 
diff --git a/src/network/protocols/connect_to_server.hpp b/src/network/protocols/connect_to_server.hpp
index 338e70b2e..fe3d3842b 100644
--- a/src/network/protocols/connect_to_server.hpp
+++ b/src/network/protocols/connect_to_server.hpp
@@ -53,7 +53,7 @@ private:
     std::atomic<ConnectState> m_state;
 
     void getClientServerInfo();
-    void registerWithSTKServer();
+    bool registerWithSTKServer();
     bool tryConnect(int timeout, int retry, bool another_port = false,
                     bool ipv6 = false);
     static ENetAddress m_server_address;