From e6d0d57f808f5763fc8dcba30f3c4ff314851b55 Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Mon, 21 Oct 2024 14:05:31 +0100
Subject: [PATCH] Wait for initial remote probe before connecting to server.
 Closes #3445

---
 qrenderdoc/Windows/MainWindow.cpp |  8 ++++++++
 qrenderdoc/Windows/MainWindow.h   |  1 +
 renderdoc/android/android.cpp     | 27 +++++++++++++++++++++++----
 renderdoc/core/remote_server.cpp  |  3 +++
 4 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/qrenderdoc/Windows/MainWindow.cpp b/qrenderdoc/Windows/MainWindow.cpp
index 6399da7b8..d7a61dc5e 100644
--- a/qrenderdoc/Windows/MainWindow.cpp
+++ b/qrenderdoc/Windows/MainWindow.cpp
@@ -240,6 +240,9 @@ MainWindow::MainWindow(ICaptureContext &ctx) : QMainWindow(NULL), ui(new Ui::Mai
       // do a remoteProbe immediately to populate the device list on startup.
       remoteProbe();
 
+      // allow any early-init replay host switches now that we've populated the device list
+      m_RemoteInitialProbeReady.release();
+
       // do several small sleeps so we can respond quicker when we need to shut down
       for(int i = 0; i < 50; i++)
       {
@@ -1932,6 +1935,11 @@ void MainWindow::setRemoteHost(int hostIdx)
   if(!PromptCloseCapture())
     return;
 
+  // we only want to block once before the initial remoteProbe has happened. After that once we can
+  // acquire once, we re-release so the next one can happen
+  m_RemoteInitialProbeReady.acquire();
+  m_RemoteInitialProbeReady.release();
+
   rdcarray<RemoteHost> hosts = m_Ctx.Config().GetRemoteHosts();
 
   RemoteHost host;
diff --git a/qrenderdoc/Windows/MainWindow.h b/qrenderdoc/Windows/MainWindow.h
index 26951b79f..4163db631 100644
--- a/qrenderdoc/Windows/MainWindow.h
+++ b/qrenderdoc/Windows/MainWindow.h
@@ -250,6 +250,7 @@ private:
 
   QTimer m_MessageTick;
   QSemaphore m_RemoteProbeSemaphore;
+  QSemaphore m_RemoteInitialProbeReady;
   LambdaThread *m_RemoteProbe;
 
   // m_ProbeRemoteHosts is covered by a lock. On the UI thread we copy it from the config regularly,
diff --git a/renderdoc/android/android.cpp b/renderdoc/android/android.cpp
index e88f6c59c..f86ad88da 100644
--- a/renderdoc/android/android.cpp
+++ b/renderdoc/android/android.cpp
@@ -1099,7 +1099,9 @@ struct AndroidController : public IDeviceProtocolHandler
 
     {
       SCOPED_LOCK(lock);
-      ret = devices[GetDeviceID(URL)].name;
+      auto it = devices.find(GetDeviceID(URL));
+      if(it != devices.end())
+        ret = it->second.name;
     }
 
     return ret;
@@ -1125,7 +1127,16 @@ struct AndroidController : public IDeviceProtocolHandler
     Invoke([this, &result, URL]() {
       rdcstr deviceID = GetDeviceID(URL);
 
-      Device &dev = devices[deviceID];
+      auto it = devices.find(deviceID);
+      if(it == devices.end())
+      {
+        SET_ERROR_RESULT(result, ResultCode::InternalError,
+                         "Android device %s is not a valid device ID, can't launch server",
+                         Android::GetFriendlyName(deviceID).c_str());
+        return;
+      }
+
+      Device &dev = it->second;
 
       if(!dev.active)
       {
@@ -1254,7 +1265,9 @@ struct AndroidController : public IDeviceProtocolHandler
 
     {
       SCOPED_LOCK(lock);
-      portbase = devices[deviceID].portbase;
+      auto it = devices.find(deviceID);
+      if(it != devices.end())
+        portbase = it->second.portbase;
     }
 
     if(portbase == 0)
@@ -1275,7 +1288,13 @@ struct AndroidController : public IDeviceProtocolHandler
 
     {
       SCOPED_LOCK(lock);
-      portbase = devices[deviceID].portbase;
+      auto it = devices.find(deviceID);
+      if(it == devices.end())
+      {
+        SAFE_DELETE(sock);
+        return NULL;
+      }
+      portbase = it->second.portbase;
     }
 
     return new AndroidRemoteServer(sock, deviceID, portbase);
diff --git a/renderdoc/core/remote_server.cpp b/renderdoc/core/remote_server.cpp
index c1ef72b96..bc3d5b1ff 100644
--- a/renderdoc/core/remote_server.cpp
+++ b/renderdoc/core/remote_server.cpp
@@ -1257,6 +1257,9 @@ RENDERDOC_CreateRemoteServerConnection(const rdcstr &URL, IRemoteServer **rend)
   else
     *rend = new RemoteServer(sock, deviceID);
 
+  if(*rend == NULL)
+    return RDResult(ResultCode::NetworkIOFailed);
+
   return RDResult(ResultCode::Succeeded);
 }