Support auth migration using authUrlV2 in Runner/MessageListener. (#3787)

This commit is contained in:
Tingluo Huang
2025-04-10 12:58:33 -04:00
committed by GitHub
parent f1b5b5bd5c
commit d5ccbd10d1
6 changed files with 1204 additions and 12 deletions

View File

@@ -55,6 +55,9 @@ namespace GitHub.Runner.Listener
private TaskAgentStatus runnerStatus = TaskAgentStatus.Online;
private CancellationTokenSource _getMessagesTokenSource;
private VssCredentials _creds;
private VssCredentials _credsV2;
private bool _needRefreshCredsV2 = false;
private bool _handlerInitialized = false;
public override void Initialize(IHostContext hostContext)
{
@@ -120,6 +123,13 @@ namespace GitHub.Runner.Listener
encounteringError = false;
}
if (!_handlerInitialized)
{
Trace.Info("Registering AuthMigrationChanged event handler.");
HostContext.AuthMigrationChanged += HandleAuthMigrationChanged;
_handlerInitialized = true;
}
return CreateSessionResult.Success;
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
@@ -185,6 +195,11 @@ namespace GitHub.Runner.Listener
{
if (_session != null && _session.SessionId != Guid.Empty)
{
if (_handlerInitialized)
{
HostContext.AuthMigrationChanged -= HandleAuthMigrationChanged;
}
if (!_accessTokenRevoked)
{
using (var ts = new CancellationTokenSource(TimeSpan.FromSeconds(30)))
@@ -249,7 +264,15 @@ namespace GitHub.Runner.Listener
{
var migrationMessage = JsonUtility.FromString<BrokerMigrationMessage>(message.Body);
await _brokerServer.UpdateConnectionIfNeeded(migrationMessage.BrokerBaseUrl, _creds);
_credsV2 = _credMgr.LoadCredentials(allowAuthUrlV2: true);
await _brokerServer.UpdateConnectionIfNeeded(migrationMessage.BrokerBaseUrl, _credsV2);
if (_needRefreshCredsV2)
{
Trace.Info("Refreshing credentials for V2.");
await _brokerServer.ForceRefreshConnection(_credsV2);
_needRefreshCredsV2 = false;
}
message = await _brokerServer.GetRunnerMessageAsync(_session.SessionId,
runnerStatus,
BuildConstants.RunnerPackage.Version,
@@ -341,6 +364,12 @@ namespace GitHub.Runner.Listener
encounteringError = true;
}
if (HostContext.AllowAuthMigration)
{
Trace.Info("Disable migration mode for 60 minutes.");
HostContext.DeferAuthMigration(TimeSpan.FromMinutes(60), $"Get next message failed with exception: {ex}");
}
// re-create VssConnection before next retry
await _runnerServer.RefreshConnectionAsync(RunnerConnectionType.MessageQueue, TimeSpan.FromSeconds(60));
@@ -401,8 +430,8 @@ namespace GitHub.Runner.Listener
public async Task RefreshListenerTokenAsync()
{
await _runnerServer.RefreshConnectionAsync(RunnerConnectionType.MessageQueue, TimeSpan.FromSeconds(60));
_creds = _credMgr.LoadCredentials(allowAuthUrlV2: false); // TODO: change to `true` in next PR
await _brokerServer.ForceRefreshConnection(_creds);
_credsV2 = _credMgr.LoadCredentials(allowAuthUrlV2: true);
await _brokerServer.ForceRefreshConnection(_credsV2);
}
private TaskAgentMessage DecryptMessage(TaskAgentMessage message)
@@ -533,5 +562,11 @@ namespace GitHub.Runner.Listener
return true;
}
}
private void HandleAuthMigrationChanged(object sender, EventArgs e)
{
Trace.Info($"Auth migration changed. Current allow auth migration state: {HostContext.AllowAuthMigration}");
_needRefreshCredsV2 = true;
}
}
}