From 651ea42e009490cda02d8a10428e37c9733fda31 Mon Sep 17 00:00:00 2001 From: Luke Tomlinson Date: Wed, 7 Feb 2024 11:24:40 -0500 Subject: [PATCH] Handle ForceTokenRefresh message (#3133) * Handle ForceTokenRefresh message * move to constants * format --- src/Runner.Listener/BrokerMessageListener.cs | 9 +++++++-- src/Runner.Listener/MessageListener.cs | 7 +++++++ src/Runner.Listener/Runner.cs | 4 ++++ src/Sdk/DTWebApi/WebApi/TaskAgentMessageTypes.cs | 10 ++++++++++ 4 files changed, 28 insertions(+), 2 deletions(-) create mode 100644 src/Sdk/DTWebApi/WebApi/TaskAgentMessageTypes.cs diff --git a/src/Runner.Listener/BrokerMessageListener.cs b/src/Runner.Listener/BrokerMessageListener.cs index 3781855b7..6767d0beb 100644 --- a/src/Runner.Listener/BrokerMessageListener.cs +++ b/src/Runner.Listener/BrokerMessageListener.cs @@ -273,7 +273,7 @@ namespace GitHub.Runner.Listener } // re-create VssConnection before next retry - await RefreshBrokerConnection(); + await RefreshBrokerConnectionAsync(); Trace.Info("Sleeping for {0} seconds before retrying.", _getNextMessageRetryInterval.TotalSeconds); await HostContext.Delay(_getNextMessageRetryInterval, token); @@ -303,6 +303,11 @@ namespace GitHub.Runner.Listener } } + public async Task RefreshListenerTokenAsync(CancellationToken cancellationToken) + { + await RefreshBrokerConnectionAsync(); + } + public async Task DeleteMessageAsync(TaskAgentMessage message) { await Task.CompletedTask; @@ -398,7 +403,7 @@ namespace GitHub.Runner.Listener } } - private async Task RefreshBrokerConnection() + private async Task RefreshBrokerConnectionAsync() { var configManager = HostContext.GetService(); _settings = configManager.LoadSettings(); diff --git a/src/Runner.Listener/MessageListener.cs b/src/Runner.Listener/MessageListener.cs index 04e0e5727..181ada984 100644 --- a/src/Runner.Listener/MessageListener.cs +++ b/src/Runner.Listener/MessageListener.cs @@ -25,6 +25,8 @@ namespace GitHub.Runner.Listener Task DeleteSessionAsync(); Task GetNextMessageAsync(CancellationToken token); Task DeleteMessageAsync(TaskAgentMessage message); + + Task RefreshListenerTokenAsync(CancellationToken token); void OnJobStatus(object sender, JobStatusEventArgs e); } @@ -379,6 +381,11 @@ namespace GitHub.Runner.Listener } } + public async Task RefreshListenerTokenAsync(CancellationToken cancellationToken) + { + await _runnerServer.RefreshConnectionAsync(RunnerConnectionType.MessageQueue, TimeSpan.FromSeconds(60)); + } + private TaskAgentMessage DecryptMessage(TaskAgentMessage message) { if (_session.EncryptionKey == null || diff --git a/src/Runner.Listener/Runner.cs b/src/Runner.Listener/Runner.cs index 263da5fef..59f2f71e2 100644 --- a/src/Runner.Listener/Runner.cs +++ b/src/Runner.Listener/Runner.cs @@ -596,6 +596,10 @@ namespace GitHub.Runner.Listener Trace.Info($"Service requests the hosted runner to shutdown. Reason: '{HostedRunnerShutdownMessage.Reason}'."); return Constants.Runner.ReturnCode.Success; } + else if (string.Equals(message.MessageType, TaskAgentMessageTypes.ForceTokenRefresh)) + { + await _listener.RefreshListenerTokenAsync(messageQueueLoopTokenSource.Token); + } else { Trace.Error($"Received message {message.MessageId} with unsupported message type {message.MessageType}."); diff --git a/src/Sdk/DTWebApi/WebApi/TaskAgentMessageTypes.cs b/src/Sdk/DTWebApi/WebApi/TaskAgentMessageTypes.cs new file mode 100644 index 000000000..5c8f95ed7 --- /dev/null +++ b/src/Sdk/DTWebApi/WebApi/TaskAgentMessageTypes.cs @@ -0,0 +1,10 @@ +using System; +using System.Runtime.Serialization; + +namespace GitHub.DistributedTask.WebApi +{ + public sealed class TaskAgentMessageTypes + { + public static readonly string ForceTokenRefresh = "ForceTokenRefresh"; + } +}