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

@@ -26,6 +26,7 @@ namespace GitHub.Runner.Listener
private TaskAgentStatus runnerStatus = TaskAgentStatus.Online;
private CancellationTokenSource _getMessagesTokenSource;
private VssCredentials _creds;
private VssCredentials _credsV2;
private TaskAgentSession _session;
private IRunnerServer _runnerServer;
private IBrokerServer _brokerServer;
@@ -35,7 +36,8 @@ namespace GitHub.Runner.Listener
private readonly TimeSpan _sessionCreationRetryInterval = TimeSpan.FromSeconds(30);
private readonly TimeSpan _sessionConflictRetryLimit = TimeSpan.FromMinutes(4);
private readonly TimeSpan _clockSkewRetryLimit = TimeSpan.FromMinutes(30);
private bool _needRefreshCredsV2 = false;
private bool _handlerInitialized = false;
public override void Initialize(IHostContext hostContext)
{
@@ -88,7 +90,8 @@ namespace GitHub.Runner.Listener
try
{
Trace.Info("Connecting to the Broker Server...");
await _brokerServer.ConnectAsync(new Uri(serverUrlV2), _creds);
_credsV2 = _credMgr.LoadCredentials(allowAuthUrlV2: true);
await _brokerServer.ConnectAsync(new Uri(serverUrlV2), _credsV2);
Trace.Info("VssConnection created");
if (!string.IsNullOrEmpty(serverUrl) &&
@@ -113,6 +116,13 @@ namespace GitHub.Runner.Listener
encounteringError = false;
}
if (!_handlerInitialized)
{
// Register event handler for auth migration state change
HostContext.AuthMigrationChanged += HandleAuthMigrationChanged;
_handlerInitialized = true;
}
return CreateSessionResult.Success;
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
@@ -131,7 +141,7 @@ namespace GitHub.Runner.Listener
Trace.Error("Catch exception during create session.");
Trace.Error(ex);
if (ex is VssOAuthTokenRequestException vssOAuthEx && _creds.Federated is VssOAuthCredential vssOAuthCred)
if (ex is VssOAuthTokenRequestException vssOAuthEx && _credsV2.Federated is VssOAuthCredential vssOAuthCred)
{
// "invalid_client" means the runner registration has been deleted from the server.
if (string.Equals(vssOAuthEx.Error, "invalid_client", StringComparison.OrdinalIgnoreCase))
@@ -162,6 +172,12 @@ namespace GitHub.Runner.Listener
return CreateSessionResult.Failure;
}
if (HostContext.AllowAuthMigration)
{
Trace.Info("Disable migration mode for 60 minutes.");
HostContext.DeferAuthMigration(TimeSpan.FromMinutes(60), $"Session creation failed with exception: {ex}");
}
if (!encounteringError) //print the message only on the first error
{
_term.WriteError($"{DateTime.UtcNow:u}: Runner connect error: {ex.Message}. Retrying until reconnected.");
@@ -178,6 +194,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)))
@@ -220,6 +241,13 @@ namespace GitHub.Runner.Listener
_getMessagesTokenSource = CancellationTokenSource.CreateLinkedTokenSource(token);
try
{
if (_needRefreshCredsV2)
{
Trace.Info("Refreshing broker connection.");
await RefreshBrokerConnectionAsync();
_needRefreshCredsV2 = false;
}
message = await _brokerServer.GetRunnerMessageAsync(_session.SessionId,
runnerStatus,
BuildConstants.RunnerPackage.Version,
@@ -299,6 +327,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 RefreshBrokerConnectionAsync();
@@ -434,9 +468,15 @@ namespace GitHub.Runner.Listener
private async Task RefreshBrokerConnectionAsync()
{
Trace.Info("Reload credentials.");
_creds = _credMgr.LoadCredentials(allowAuthUrlV2: false); // TODO: change to `true` in the next PR.
await _brokerServer.ConnectAsync(new Uri(_settings.ServerUrlV2), _creds);
_credsV2 = _credMgr.LoadCredentials(allowAuthUrlV2: true);
await _brokerServer.ConnectAsync(new Uri(_settings.ServerUrlV2), _credsV2);
Trace.Info("Connection to Broker Server recreated.");
}
private void HandleAuthMigrationChanged(object sender, EventArgs e)
{
Trace.Info($"Auth migration changed. Current allow auth migration state: {HostContext.AllowAuthMigration}");
_needRefreshCredsV2 = true;
}
}
}