Skip to content

Enable auth migration based on config refresh. #3786

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Apr 3, 2025
Merged

Enable auth migration based on config refresh. #3786

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e3a985c
Enable auth migration based on config refresh.
TingluoHuang Apr 3, 2025
f6b4067
.
TingluoHuang Apr 3, 2025
f929e2e
.
TingluoHuang Apr 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/Runner.Common/BrokerServer.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ public sealed class BrokerServer : RunnerService, IBrokerServer

public async Task ConnectAsync(Uri serverUri, VssCredentials credentials)
{
Trace.Entering();
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add verbose trace.

_brokerUri = serverUri;

_connection = VssUtil.CreateRawConnection(serverUri, credentials);
Expand Down
5 changes: 4 additions & 1 deletion src/Runner.Listener/MessageListener.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -245,7 +245,6 @@ public async Task<TaskAgentMessage> 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<BrokerMigrationMessage>(message.Body);
Expand Down Expand Up @@ -306,6 +305,10 @@ public async Task<TaskAgentMessage> 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;
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in case we got migration message from pipelines, but fail to get message broker-listener.
we currently return the migration message to the caller which the caller don't know how to handle the migration message.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good find


// 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))
{
Expand Down
2 changes: 2 additions & 0 deletions src/Runner.Listener/Runner.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ public sealed class Runner : RunnerService, IRunner
private ITerminal _term;
private bool _inConfigStage;
private ManualResetEvent _completedCommand = new(false);
private IRunnerServer _runnerServer;

// <summary>
// Helps avoid excessive calls to Run Service when encountering non-retriable errors from /acquirejob.
Expand All @@ -51,6 +52,7 @@ public override void Initialize(IHostContext hostContext)
base.Initialize(hostContext);
_term = HostContext.GetService<ITerminal>();
_acquireJobThrottler = HostContext.CreateService<IErrorThrottler>();
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we should always use GetService instead of CreateService for singleton service.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was originally intended not to be a singleton. Why switch to singleton?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The error throttler helps us back off when encountering successive, non-retriable errors from /acquirejob.

Copy link
Member Author

@TingluoHuang TingluoHuang Apr 3, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i didn't find any place we use IErrorThrottler outside of the Runner.cs, so i thought we only need one.
I guess, if the original idea is to create IErrorThrottler as needed, then CreateService make more sense.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i changed it back for now.

_runnerServer = HostContext.GetService<IRunnerServer>();
}

public async Task<int> ExecuteCommand(CommandSettings command)
Expand Down
12 changes: 11 additions & 1 deletion src/Runner.Listener/RunnerConfigUpdater.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is acting as disable the migration via FF from the service.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm fuzzy on this detail. This PR seems fine for now. I will catch up on this detail on the next PR.

await ReportTelemetryAsync("Runner credentials updated successfully. Auth migration is disabled.");
}
}

private async Task<bool> VerifyRunnerQualifiedId(string runnerQualifiedId)
Expand Down
8 changes: 1 addition & 7 deletions src/Test/L0/Listener/BrokerMessageListenerL0.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,24 +18,21 @@ public sealed class BrokerMessageListenerL0
private readonly Mock<IBrokerServer> _brokerServer;
private readonly Mock<IRunnerServer> _runnerServer;
private readonly Mock<ICredentialManager> _credMgr;
private Mock<IConfigurationStore> _store;
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we don't use this any more.



public BrokerMessageListenerL0()
{
_settings = new RunnerSettings { AgentId = 1, AgentName = "myagent", PoolId = 123, PoolName = "default", ServerUrl = "http://myserver", WorkFolder = "_work", ServerUrlV2 = "http://myserverv2" };
_config = new Mock<IConfigurationManager>();
_config.Setup(x => x.LoadSettings()).Returns(_settings);
_credMgr = new Mock<ICredentialManager>();
_store = new Mock<IConfigurationStore>();
_brokerServer = new Mock<IBrokerServer>();
_runnerServer = new Mock<IRunnerServer>();
}

[Fact]
[Trait("Level", "L0")]
[Trait("Category", "Runner")]
public async void CreatesSession()
public async Task CreatesSession()
{
using (TestHostContext tc = CreateTestContext())
using (var tokenSource = new CancellationTokenSource())
Expand All @@ -51,8 +48,6 @@ public async void CreatesSession()
.Returns(Task.FromResult(expectedSession));

_credMgr.Setup(x => x.LoadCredentials(It.IsAny<bool>())).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();
Expand All @@ -75,7 +70,6 @@ private TestHostContext CreateTestContext([CallerMemberName] String testName = "
TestHostContext tc = new(this, testName);
tc.SetSingleton<IConfigurationManager>(_config.Object);
tc.SetSingleton<ICredentialManager>(_credMgr.Object);
tc.SetSingleton<IConfigurationStore>(_store.Object);
tc.SetSingleton<IBrokerServer>(_brokerServer.Object);
tc.SetSingleton<IRunnerServer>(_runnerServer.Object);
return tc;
Expand Down
12 changes: 6 additions & 6 deletions src/Test/L0/Listener/MessageListenerL0.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -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())
Expand Down
53 changes: 51 additions & 2 deletions src/Test/L0/Listener/RunnerConfigUpdaterTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int>(), It.Is<string>(s => s == "credentials"), It.IsAny<string>(), It.IsAny<CancellationToken>())).ReturnsAsync(encodedConfig);


var _runnerConfigUpdater = new RunnerConfigUpdater();
_runnerConfigUpdater.Initialize(hc);
hc.EnableAuthMigration("L0Test");

var validRunnerQualifiedId = "valid/runner/qualifiedid/1";
var configType = "credentials";
Expand All @@ -226,6 +226,7 @@ public async Task UpdateRunnerConfigAsync_UpdateRunnerCredentials_ShouldSucceed(
_runnerServer.Verify(x => x.RefreshRunnerConfigAsync(1, "credentials", It.IsAny<string>(), It.IsAny<CancellationToken>()), Times.Once);
_runnerServer.Verify(x => x.UpdateAgentUpdateStateAsync(It.IsAny<int>(), It.IsAny<ulong>(), It.IsAny<string>(), It.Is<string>(s => s.Contains("Runner credentials updated successfully")), It.IsAny<CancellationToken>()), Times.Once);
_configurationStore.Verify(x => x.SaveMigratedCredential(It.IsAny<CredentialData>()), Times.Once);
Assert.False(hc.AllowAuthMigration);
}
}

Expand Down Expand Up @@ -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))
{
Expand Down Expand Up @@ -625,5 +626,53 @@ public async Task UpdateRunnerConfigAsync_RunnerAdminService_ShouldThrowNotSuppo
_configurationStore.Verify(x => x.SaveMigratedSettings(It.IsAny<RunnerSettings>()), Times.Never);
}
}

[Fact]
[Trait("Level", "L0")]
[Trait("Category", "Runner")]
public async Task UpdateRunnerConfigAsync_UpdateRunnerCredentials_EnableDisableAuthMigration()
{
using (var hc = new TestHostContext(this))
{
hc.SetSingleton<IConfigurationStore>(_configurationStore.Object);
hc.SetSingleton<IRunnerServer>(_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<int>(), It.Is<string>(s => s == "credentials"), It.IsAny<string>(), It.IsAny<CancellationToken>())).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<string>(), It.IsAny<CancellationToken>()), Times.Once);
_runnerServer.Verify(x => x.UpdateAgentUpdateStateAsync(It.IsAny<int>(), It.IsAny<ulong>(), It.IsAny<string>(), It.Is<string>(s => s.Contains("Runner credentials updated successfully")), It.IsAny<CancellationToken>()), Times.Once);
_configurationStore.Verify(x => x.SaveMigratedCredential(It.IsAny<CredentialData>()), Times.Once);
Assert.True(hc.AllowAuthMigration);
}
}
}
}
29 changes: 16 additions & 13 deletions src/Test/L0/Listener/RunnerL0.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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))
{
Expand Down Expand Up @@ -169,14 +169,15 @@ 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))
{
hc.SetSingleton<IConfigurationManager>(_configurationManager.Object);
hc.SetSingleton<IPromptManager>(_promptManager.Object);
hc.SetSingleton<IMessageListener>(_messageListener.Object);
hc.SetSingleton<IConfigurationStore>(_configStore.Object);
hc.SetSingleton<IRunnerServer>(_runnerServer.Object);
hc.EnqueueInstance<IErrorThrottler>(_acquireJobThrottler.Object);

var command = new CommandSettings(hc, args);
Expand All @@ -201,14 +202,15 @@ 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))
{
hc.SetSingleton<IConfigurationManager>(_configurationManager.Object);
hc.SetSingleton<IPromptManager>(_promptManager.Object);
hc.SetSingleton<IMessageListener>(_messageListener.Object);
hc.SetSingleton<IConfigurationStore>(_configStore.Object);
hc.SetSingleton<IRunnerServer>(_runnerServer.Object);
hc.EnqueueInstance<IErrorThrottler>(_acquireJobThrottler.Object);

var command = new CommandSettings(hc, new[] { "run" });
Expand All @@ -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))
{
Expand Down Expand Up @@ -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))
{
Expand Down Expand Up @@ -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))
{
Expand Down Expand Up @@ -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<IConfigurationManager>(_configurationManager.Object);
hc.SetSingleton<IConfigurationStore>(_configStore.Object);
hc.SetSingleton<IPromptManager>(_promptManager.Object);
hc.SetSingleton<IRunnerServer>(_runnerServer.Object);
hc.EnqueueInstance<IErrorThrottler>(_acquireJobThrottler.Object);

var command = new CommandSettings(hc, new[] { "remove", "--local" });
Expand Down
4 changes: 2 additions & 2 deletions src/Test/L0/TestHostContext.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

zero will make the test case stuck in case we use the HostContext.Delay in some background task.

// Delay 10ms
await Task.Delay(TimeSpan.FromMilliseconds(10));
}

public T CreateService<T>() where T : class, IRunnerService
Expand Down