diff --git a/src/Runner.Common/BrokerServer.cs b/src/Runner.Common/BrokerServer.cs index 04d48aa0c44..1f6c01685bd 100644 --- a/src/Runner.Common/BrokerServer.cs +++ b/src/Runner.Common/BrokerServer.cs @@ -37,6 +37,7 @@ public sealed class BrokerServer : RunnerService, IBrokerServer public async Task ConnectAsync(Uri serverUri, VssCredentials credentials) { + Trace.Entering(); _brokerUri = serverUri; _connection = VssUtil.CreateRawConnection(serverUri, credentials); diff --git a/src/Runner.Listener/MessageListener.cs b/src/Runner.Listener/MessageListener.cs index d8e2d3c2a71..4f387536361 100644 --- a/src/Runner.Listener/MessageListener.cs +++ b/src/Runner.Listener/MessageListener.cs @@ -245,7 +245,6 @@ public async Task GetNextMessageAsync(CancellationToken token) // Decrypt the message body if the session is using encryption message = DecryptMessage(message); - if (message != null && message.MessageType == BrokerMigrationMessage.MessageType) { var migrationMessage = JsonUtility.FromString(message.Body); @@ -306,6 +305,10 @@ public async Task GetNextMessageAsync(CancellationToken token) Trace.Error("Catch exception during get next message."); Trace.Error(ex); + // clear out potential message for broker migration, + // in case the exception is thrown from get message from broker-listener. + message = null; + // don't retry if SkipSessionRecover = true, DT service will delete agent session to stop agent from taking more jobs. if (ex is TaskAgentSessionExpiredException && !_settings.SkipSessionRecover && (await CreateSessionAsync(token) == CreateSessionResult.Success)) { diff --git a/src/Runner.Listener/Runner.cs b/src/Runner.Listener/Runner.cs index b012b5df138..a7ef60d1cca 100644 --- a/src/Runner.Listener/Runner.cs +++ b/src/Runner.Listener/Runner.cs @@ -31,6 +31,7 @@ public sealed class Runner : RunnerService, IRunner private ITerminal _term; private bool _inConfigStage; private ManualResetEvent _completedCommand = new(false); + private IRunnerServer _runnerServer; // // Helps avoid excessive calls to Run Service when encountering non-retriable errors from /acquirejob. @@ -51,6 +52,7 @@ public override void Initialize(IHostContext hostContext) base.Initialize(hostContext); _term = HostContext.GetService(); _acquireJobThrottler = HostContext.CreateService(); + _runnerServer = HostContext.GetService(); } public async Task ExecuteCommand(CommandSettings command) diff --git a/src/Runner.Listener/RunnerConfigUpdater.cs b/src/Runner.Listener/RunnerConfigUpdater.cs index 5d4b76dfd08..c188ad7318a 100644 --- a/src/Runner.Listener/RunnerConfigUpdater.cs +++ b/src/Runner.Listener/RunnerConfigUpdater.cs @@ -211,7 +211,17 @@ private async Task UpdateRunnerCredentialsAsync(string serviceType, string confi // save the refreshed runner credentials as a separate file _store.SaveMigratedCredential(refreshedCredConfig); - await ReportTelemetryAsync("Runner credentials updated successfully."); + + if (refreshedCredConfig.Data.ContainsKey("authorizationUrlV2")) + { + HostContext.EnableAuthMigration("Credential file updated"); + await ReportTelemetryAsync("Runner credentials updated successfully. Auth migration is enabled."); + } + else + { + HostContext.DeferAuthMigration(TimeSpan.FromDays(365), "Credential file does not contain authorizationUrlV2"); + await ReportTelemetryAsync("Runner credentials updated successfully. Auth migration is disabled."); + } } private async Task VerifyRunnerQualifiedId(string runnerQualifiedId) diff --git a/src/Test/L0/Listener/BrokerMessageListenerL0.cs b/src/Test/L0/Listener/BrokerMessageListenerL0.cs index 6821ac93871..3328d1852c6 100644 --- a/src/Test/L0/Listener/BrokerMessageListenerL0.cs +++ b/src/Test/L0/Listener/BrokerMessageListenerL0.cs @@ -18,8 +18,6 @@ public sealed class BrokerMessageListenerL0 private readonly Mock _brokerServer; private readonly Mock _runnerServer; private readonly Mock _credMgr; - private Mock _store; - public BrokerMessageListenerL0() { @@ -27,7 +25,6 @@ public BrokerMessageListenerL0() _config = new Mock(); _config.Setup(x => x.LoadSettings()).Returns(_settings); _credMgr = new Mock(); - _store = new Mock(); _brokerServer = new Mock(); _runnerServer = new Mock(); } @@ -35,7 +32,7 @@ public BrokerMessageListenerL0() [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void CreatesSession() + public async Task CreatesSession() { using (TestHostContext tc = CreateTestContext()) using (var tokenSource = new CancellationTokenSource()) @@ -51,8 +48,6 @@ public async void CreatesSession() .Returns(Task.FromResult(expectedSession)); _credMgr.Setup(x => x.LoadCredentials(It.IsAny())).Returns(new VssCredentials()); - _store.Setup(x => x.GetCredentials()).Returns(new CredentialData() { Scheme = Constants.Configuration.OAuthAccessToken }); - _store.Setup(x => x.GetMigratedCredentials()).Returns(default(CredentialData)); // Act. BrokerMessageListener listener = new(); @@ -75,7 +70,6 @@ private TestHostContext CreateTestContext([CallerMemberName] String testName = " TestHostContext tc = new(this, testName); tc.SetSingleton(_config.Object); tc.SetSingleton(_credMgr.Object); - tc.SetSingleton(_store.Object); tc.SetSingleton(_brokerServer.Object); tc.SetSingleton(_runnerServer.Object); return tc; diff --git a/src/Test/L0/Listener/MessageListenerL0.cs b/src/Test/L0/Listener/MessageListenerL0.cs index 5b156fe98f8..2da28d710b5 100644 --- a/src/Test/L0/Listener/MessageListenerL0.cs +++ b/src/Test/L0/Listener/MessageListenerL0.cs @@ -51,7 +51,7 @@ private TestHostContext CreateTestContext([CallerMemberName] String testName = " [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void CreatesSession() + public async Task CreatesSession() { using (TestHostContext tc = CreateTestContext()) using (var tokenSource = new CancellationTokenSource()) @@ -95,7 +95,7 @@ public async void CreatesSession() [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void DeleteSession() + public async Task DeleteSession() { using (TestHostContext tc = CreateTestContext()) using (var tokenSource = new CancellationTokenSource()) @@ -142,7 +142,7 @@ public async void DeleteSession() [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void GetNextMessage() + public async Task GetNextMessage() { using (TestHostContext tc = CreateTestContext()) using (var tokenSource = new CancellationTokenSource()) @@ -223,7 +223,7 @@ public async void GetNextMessage() [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void GetNextMessageWithBrokerMigration() + public async Task GetNextMessageWithBrokerMigration() { using (TestHostContext tc = CreateTestContext()) using (var tokenSource = new CancellationTokenSource()) @@ -329,7 +329,7 @@ public async void GetNextMessageWithBrokerMigration() [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void CreateSessionWithOriginalCredential() + public async Task CreateSessionWithOriginalCredential() { using (TestHostContext tc = CreateTestContext()) using (var tokenSource = new CancellationTokenSource()) @@ -374,7 +374,7 @@ public async void CreateSessionWithOriginalCredential() [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void SkipDeleteSession_WhenGetNextMessageGetTaskAgentAccessTokenExpiredException() + public async Task SkipDeleteSession_WhenGetNextMessageGetTaskAgentAccessTokenExpiredException() { using (TestHostContext tc = CreateTestContext()) using (var tokenSource = new CancellationTokenSource()) diff --git a/src/Test/L0/Listener/RunnerConfigUpdaterTests.cs b/src/Test/L0/Listener/RunnerConfigUpdaterTests.cs index 48e798d3ec6..63deafe5b85 100644 --- a/src/Test/L0/Listener/RunnerConfigUpdaterTests.cs +++ b/src/Test/L0/Listener/RunnerConfigUpdaterTests.cs @@ -210,9 +210,9 @@ public async Task UpdateRunnerConfigAsync_UpdateRunnerCredentials_ShouldSucceed( var encodedConfig = Convert.ToBase64String(Encoding.UTF8.GetBytes(StringUtil.ConvertToJson(credData))); _runnerServer.Setup(x => x.RefreshRunnerConfigAsync(It.IsAny(), It.Is(s => s == "credentials"), It.IsAny(), It.IsAny())).ReturnsAsync(encodedConfig); - var _runnerConfigUpdater = new RunnerConfigUpdater(); _runnerConfigUpdater.Initialize(hc); + hc.EnableAuthMigration("L0Test"); var validRunnerQualifiedId = "valid/runner/qualifiedid/1"; var configType = "credentials"; @@ -226,6 +226,7 @@ public async Task UpdateRunnerConfigAsync_UpdateRunnerCredentials_ShouldSucceed( _runnerServer.Verify(x => x.RefreshRunnerConfigAsync(1, "credentials", It.IsAny(), It.IsAny()), Times.Once); _runnerServer.Verify(x => x.UpdateAgentUpdateStateAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.Is(s => s.Contains("Runner credentials updated successfully")), It.IsAny()), Times.Once); _configurationStore.Verify(x => x.SaveMigratedCredential(It.IsAny()), Times.Once); + Assert.False(hc.AllowAuthMigration); } } @@ -306,7 +307,7 @@ public async Task UpdateRunnerConfigAsync_RefreshRunnerSettingsFailure_ShouldRep [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async Task UpdateRunnerConfigAsync_RefreshRunnerCredetialsFailure_ShouldReportTelemetry() + public async Task UpdateRunnerConfigAsync_RefreshRunnerCredentialsFailure_ShouldReportTelemetry() { using (var hc = new TestHostContext(this)) { @@ -625,5 +626,53 @@ public async Task UpdateRunnerConfigAsync_RunnerAdminService_ShouldThrowNotSuppo _configurationStore.Verify(x => x.SaveMigratedSettings(It.IsAny()), Times.Never); } } + + [Fact] + [Trait("Level", "L0")] + [Trait("Category", "Runner")] + public async Task UpdateRunnerConfigAsync_UpdateRunnerCredentials_EnableDisableAuthMigration() + { + using (var hc = new TestHostContext(this)) + { + hc.SetSingleton(_configurationStore.Object); + hc.SetSingleton(_runnerServer.Object); + + // Arrange + var setting = new RunnerSettings { AgentId = 1, AgentName = "agent1" }; + _configurationStore.Setup(x => x.GetSettings()).Returns(setting); + var credData = new CredentialData + { + Scheme = "OAuth" + }; + credData.Data.Add("ClientId", "12345"); + credData.Data.Add("AuthorizationUrl", "https://example.com"); + credData.Data.Add("AuthorizationUrlV2", "https://example2.com"); + _configurationStore.Setup(x => x.GetCredentials()).Returns(credData); + + IOUtil.SaveObject(setting, hc.GetConfigFile(WellKnownConfigFile.Runner)); + IOUtil.SaveObject(credData, hc.GetConfigFile(WellKnownConfigFile.Credentials)); + + var encodedConfig = Convert.ToBase64String(Encoding.UTF8.GetBytes(StringUtil.ConvertToJson(credData))); + _runnerServer.Setup(x => x.RefreshRunnerConfigAsync(It.IsAny(), It.Is(s => s == "credentials"), It.IsAny(), It.IsAny())).ReturnsAsync(encodedConfig); + + var _runnerConfigUpdater = new RunnerConfigUpdater(); + _runnerConfigUpdater.Initialize(hc); + Assert.False(hc.AllowAuthMigration); + + var validRunnerQualifiedId = "valid/runner/qualifiedid/1"; + var configType = "credentials"; + var serviceType = "pipelines"; + var configRefreshUrl = "http://example.com"; + + // Act + await _runnerConfigUpdater.UpdateRunnerConfigAsync(validRunnerQualifiedId, configType, serviceType, configRefreshUrl); + + // Assert + _runnerServer.Verify(x => x.RefreshRunnerConfigAsync(1, "credentials", It.IsAny(), It.IsAny()), Times.Once); + _runnerServer.Verify(x => x.UpdateAgentUpdateStateAsync(It.IsAny(), It.IsAny(), It.IsAny(), It.Is(s => s.Contains("Runner credentials updated successfully")), It.IsAny()), Times.Once); + _configurationStore.Verify(x => x.SaveMigratedCredential(It.IsAny()), Times.Once); + Assert.True(hc.AllowAuthMigration); + } + } } } diff --git a/src/Test/L0/Listener/RunnerL0.cs b/src/Test/L0/Listener/RunnerL0.cs index b29f8835c2b..1dd24fe4d74 100644 --- a/src/Test/L0/Listener/RunnerL0.cs +++ b/src/Test/L0/Listener/RunnerL0.cs @@ -1,13 +1,13 @@ -using GitHub.DistributedTask.WebApi; -using GitHub.Runner.Listener; -using GitHub.Runner.Listener.Configuration; -using Moq; -using System; +using System; using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; -using Xunit; +using GitHub.DistributedTask.WebApi; +using GitHub.Runner.Listener; +using GitHub.Runner.Listener.Configuration; using GitHub.Services.WebApi; +using Moq; +using Xunit; using Pipelines = GitHub.DistributedTask.Pipelines; namespace GitHub.Runner.Common.Tests.Listener @@ -57,7 +57,7 @@ private JobCancelMessage CreateJobCancelMessage() [Trait("Level", "L0")] [Trait("Category", "Runner")] //process 2 new job messages, and one cancel message - public async void TestRunAsync() + public async Task TestRunAsync() { using (var hc = new TestHostContext(this)) { @@ -169,7 +169,7 @@ public async void TestRunAsync() [MemberData(nameof(RunAsServiceTestData))] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void TestExecuteCommandForRunAsService(string[] args, bool configureAsService, Times expectedTimes) + public async Task TestExecuteCommandForRunAsService(string[] args, bool configureAsService, Times expectedTimes) { using (var hc = new TestHostContext(this)) { @@ -177,6 +177,7 @@ public async void TestExecuteCommandForRunAsService(string[] args, bool configur hc.SetSingleton(_promptManager.Object); hc.SetSingleton(_messageListener.Object); hc.SetSingleton(_configStore.Object); + hc.SetSingleton(_runnerServer.Object); hc.EnqueueInstance(_acquireJobThrottler.Object); var command = new CommandSettings(hc, args); @@ -201,7 +202,7 @@ public async void TestExecuteCommandForRunAsService(string[] args, bool configur [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void TestMachineProvisionerCLI() + public async Task TestMachineProvisionerCLI() { using (var hc = new TestHostContext(this)) { @@ -209,6 +210,7 @@ public async void TestMachineProvisionerCLI() hc.SetSingleton(_promptManager.Object); hc.SetSingleton(_messageListener.Object); hc.SetSingleton(_configStore.Object); + hc.SetSingleton(_runnerServer.Object); hc.EnqueueInstance(_acquireJobThrottler.Object); var command = new CommandSettings(hc, new[] { "run" }); @@ -235,7 +237,7 @@ public async void TestMachineProvisionerCLI() [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void TestRunOnce() + public async Task TestRunOnce() { using (var hc = new TestHostContext(this)) { @@ -332,7 +334,7 @@ public async void TestRunOnce() [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void TestRunOnceOnlyTakeOneJobMessage() + public async Task TestRunOnceOnlyTakeOneJobMessage() { using (var hc = new TestHostContext(this)) { @@ -433,7 +435,7 @@ public async void TestRunOnceOnlyTakeOneJobMessage() [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void TestRunOnceHandleUpdateMessage() + public async Task TestRunOnceHandleUpdateMessage() { using (var hc = new TestHostContext(this)) { @@ -523,13 +525,14 @@ public async void TestRunOnceHandleUpdateMessage() [Fact] [Trait("Level", "L0")] [Trait("Category", "Runner")] - public async void TestRemoveLocalRunnerConfig() + public async Task TestRemoveLocalRunnerConfig() { using (var hc = new TestHostContext(this)) { hc.SetSingleton(_configurationManager.Object); hc.SetSingleton(_configStore.Object); hc.SetSingleton(_promptManager.Object); + hc.SetSingleton(_runnerServer.Object); hc.EnqueueInstance(_acquireJobThrottler.Object); var command = new CommandSettings(hc, new[] { "remove", "--local" }); diff --git a/src/Test/L0/TestHostContext.cs b/src/Test/L0/TestHostContext.cs index 86a2e80ce7d..c1cf692204f 100644 --- a/src/Test/L0/TestHostContext.cs +++ b/src/Test/L0/TestHostContext.cs @@ -103,8 +103,8 @@ public async Task Delay(TimeSpan delay, CancellationToken token) handler(this, new DelayEventArgs(delay, token)); } - // Delay zero - await Task.Delay(TimeSpan.Zero); + // Delay 10ms + await Task.Delay(TimeSpan.FromMilliseconds(10)); } public T CreateService() where T : class, IRunnerService