From 3449d5fa52765cfbfbc883f981d8767f8af141a3 Mon Sep 17 00:00:00 2001
From: Luke Tomlinson <luketomlinson@github.com>
Date: Wed, 21 Feb 2024 16:43:01 -0500
Subject: [PATCH] Broker fixes for token refreshes and AccessDeniedException
 (#3161)

---
 src/Runner.Common/BrokerServer.cs         | 24 ++++++++++++++++++++++-
 src/Runner.Listener/MessageListener.cs    |  5 +++--
 src/Sdk/WebApi/WebApi/BrokerHttpClient.cs |  7 ++++++-
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/src/Runner.Common/BrokerServer.cs b/src/Runner.Common/BrokerServer.cs
index 77bf5d88290..0627afd3fb6 100644
--- a/src/Runner.Common/BrokerServer.cs
+++ b/src/Runner.Common/BrokerServer.cs
@@ -21,6 +21,8 @@ public interface IBrokerServer : IRunnerService
         Task DeleteSessionAsync(CancellationToken cancellationToken);
 
         Task<TaskAgentMessage> GetRunnerMessageAsync(Guid? sessionId, TaskAgentStatus status, string version, string os, string architecture, bool disableUpdate, CancellationToken token);
+
+        Task UpdateConnectionIfNeeded(Uri serverUri, VssCredentials credentials);
     }
 
     public sealed class BrokerServer : RunnerService, IBrokerServer
@@ -59,7 +61,7 @@ public Task<TaskAgentMessage> GetRunnerMessageAsync(Guid? sessionId, TaskAgentSt
         {
             CheckConnection();
             var brokerSession = RetryRequest<TaskAgentMessage>(
-                async () => await _brokerHttpClient.GetRunnerMessageAsync(sessionId, version, status, os, architecture, disableUpdate, cancellationToken), cancellationToken);
+                async () => await _brokerHttpClient.GetRunnerMessageAsync(sessionId, version, status, os, architecture, disableUpdate, cancellationToken), cancellationToken, shouldRetry: ShouldRetryException);
 
             return brokerSession;
         }
@@ -69,5 +71,25 @@ public async Task DeleteSessionAsync(CancellationToken cancellationToken)
             CheckConnection();
             await _brokerHttpClient.DeleteSessionAsync(cancellationToken);
         }
+
+        public Task UpdateConnectionIfNeeded(Uri serverUri, VssCredentials credentials)
+        {
+            if (_brokerUri != serverUri || !_hasConnection)
+            {
+                return ConnectAsync(serverUri, credentials);
+            }
+
+            return Task.CompletedTask;
+        }
+
+        public bool ShouldRetryException(Exception ex)
+        {
+            if (ex is AccessDeniedException ade && ade.ErrorCode == 1)
+            {
+                return false;
+            }
+
+            return true;
+        }
     }
 }
diff --git a/src/Runner.Listener/MessageListener.cs b/src/Runner.Listener/MessageListener.cs
index 181ada984d3..321c4d947ae 100644
--- a/src/Runner.Listener/MessageListener.cs
+++ b/src/Runner.Listener/MessageListener.cs
@@ -109,7 +109,8 @@ public async Task<Boolean> CreateSessionAsync(CancellationToken token)
                     if (_session.BrokerMigrationMessage != null)
                     {
                         Trace.Info("Runner session is in migration mode: Creating Broker session with BrokerBaseUrl: {0}", _session.BrokerMigrationMessage.BrokerBaseUrl);
-                        await _brokerServer.ConnectAsync(_session.BrokerMigrationMessage.BrokerBaseUrl, _creds);
+
+                        await _brokerServer.UpdateConnectionIfNeeded(_session.BrokerMigrationMessage.BrokerBaseUrl, _creds);
                         _session = await _brokerServer.CreateSessionAsync(taskAgentSession, token);
                         _isBrokerSession = true;
                     }
@@ -256,7 +257,7 @@ public async Task<TaskAgentMessage> GetNextMessageAsync(CancellationToken token)
 
                         var migrationMessage = JsonUtility.FromString<BrokerMigrationMessage>(message.Body);
 
-                        await _brokerServer.ConnectAsync(migrationMessage.BrokerBaseUrl, _creds);
+                        await _brokerServer.UpdateConnectionIfNeeded(migrationMessage.BrokerBaseUrl, _creds);
                         message = await _brokerServer.GetRunnerMessageAsync(_session.SessionId,
                                                                         runnerStatus,
                                                                         BuildConstants.RunnerPackage.Version,
diff --git a/src/Sdk/WebApi/WebApi/BrokerHttpClient.cs b/src/Sdk/WebApi/WebApi/BrokerHttpClient.cs
index 37380e0c31d..e9ad938fb9f 100644
--- a/src/Sdk/WebApi/WebApi/BrokerHttpClient.cs
+++ b/src/Sdk/WebApi/WebApi/BrokerHttpClient.cs
@@ -110,9 +110,14 @@ public async Task<TaskAgentMessage> GetRunnerMessageAsync(
                 return result.Value;
             }
 
+            // the only time we throw a `Forbidden` exception from Listener /messages is when the runner is
+            // disable_update and is too old to poll
             if (result.StatusCode == HttpStatusCode.Forbidden)
             {
-                throw new AccessDeniedException(result.Error);
+                throw new AccessDeniedException($"{result.Error} Runner version v{runnerVersion} is deprecated and cannot receive messages.")
+                {
+                    ErrorCode = 1
+                };
             }
 
             throw new Exception($"Failed to get job message: {result.Error}");