From a65ac083b429b49bff4937117c1daf8f555558a4 Mon Sep 17 00:00:00 2001 From: Tingluo Huang Date: Wed, 16 Feb 2022 16:10:18 -0500 Subject: [PATCH] Skip DeleteAgentSession when the acess token has been revoked. (#1692) --- src/Runner.Listener/MessageListener.cs | 32 ++++++--- src/Test/L0/Listener/MessageListenerL0.cs | 86 +++++++++++++++++++---- 2 files changed, 95 insertions(+), 23 deletions(-) diff --git a/src/Runner.Listener/MessageListener.cs b/src/Runner.Listener/MessageListener.cs index 71e5e43f5..81afd737d 100644 --- a/src/Runner.Listener/MessageListener.cs +++ b/src/Runner.Listener/MessageListener.cs @@ -1,18 +1,18 @@ -using GitHub.DistributedTask.WebApi; -using GitHub.Runner.Listener.Configuration; -using GitHub.Services.Common; using System; using System.Collections.Generic; +using System.Diagnostics; +using System.IO; +using System.Runtime.InteropServices; +using System.Security.Cryptography; +using System.Text; using System.Threading; using System.Threading.Tasks; -using System.Security.Cryptography; -using System.IO; -using System.Text; -using GitHub.Services.OAuth; -using System.Diagnostics; -using System.Runtime.InteropServices; +using GitHub.DistributedTask.WebApi; using GitHub.Runner.Common; +using GitHub.Runner.Listener.Configuration; using GitHub.Runner.Sdk; +using GitHub.Services.Common; +using GitHub.Services.OAuth; namespace GitHub.Runner.Listener { @@ -33,6 +33,7 @@ namespace GitHub.Runner.Listener private IRunnerServer _runnerServer; private TaskAgentSession _session; private TimeSpan _getNextMessageRetryInterval; + private bool _accessTokenRevoked = false; private readonly TimeSpan _sessionCreationRetryInterval = TimeSpan.FromSeconds(30); private readonly TimeSpan _sessionConflictRetryLimit = TimeSpan.FromMinutes(4); private readonly TimeSpan _clockSkewRetryLimit = TimeSpan.FromMinutes(30); @@ -111,6 +112,7 @@ namespace GitHub.Runner.Listener catch (TaskAgentAccessTokenExpiredException) { Trace.Info("Runner OAuth token has been revoked. Session creation failed."); + _accessTokenRevoked = true; throw; } catch (Exception ex) @@ -154,9 +156,16 @@ namespace GitHub.Runner.Listener { if (_session != null && _session.SessionId != Guid.Empty) { - using (var ts = new CancellationTokenSource(TimeSpan.FromSeconds(30))) + if (!_accessTokenRevoked) { - await _runnerServer.DeleteAgentSessionAsync(_settings.PoolId, _session.SessionId, ts.Token); + using (var ts = new CancellationTokenSource(TimeSpan.FromSeconds(30))) + { + await _runnerServer.DeleteAgentSessionAsync(_settings.PoolId, _session.SessionId, ts.Token); + } + } + else + { + Trace.Warning("Runner OAuth token has been revoked. Skip deleting session."); } } } @@ -205,6 +214,7 @@ namespace GitHub.Runner.Listener catch (TaskAgentAccessTokenExpiredException) { Trace.Info("Runner OAuth token has been revoked. Unable to pull message."); + _accessTokenRevoked = true; throw; } catch (Exception ex) diff --git a/src/Test/L0/Listener/MessageListenerL0.cs b/src/Test/L0/Listener/MessageListenerL0.cs index a830e9c99..05d5d14d1 100644 --- a/src/Test/L0/Listener/MessageListenerL0.cs +++ b/src/Test/L0/Listener/MessageListenerL0.cs @@ -1,18 +1,18 @@ -using GitHub.DistributedTask.WebApi; -using GitHub.Services.Common; -using GitHub.Services.WebApi; -using GitHub.Runner.Listener; -using GitHub.Runner.Listener.Configuration; -using Moq; -using System; -using System.Runtime.CompilerServices; -using System.Threading.Tasks; -using Xunit; -using System.Threading; -using System.Reflection; +using System; using System.Collections.Generic; using System.IO; +using System.Reflection; +using System.Runtime.CompilerServices; using System.Security.Cryptography; +using System.Threading; +using System.Threading.Tasks; +using GitHub.DistributedTask.WebApi; +using GitHub.Runner.Listener; +using GitHub.Runner.Listener.Configuration; +using GitHub.Services.Common; +using GitHub.Services.WebApi; +using Moq; +using Xunit; namespace GitHub.Runner.Common.Tests.Listener { @@ -256,5 +256,67 @@ namespace GitHub.Runner.Common.Tests.Listener tokenSource.Token), Times.Once()); } } + + [Fact] + [Trait("Level", "L0")] + [Trait("Category", "Runner")] + public async void SkipDeleteSession_WhenGetNextMessageGetTaskAgentAccessTokenExpiredException() + { + using (TestHostContext tc = CreateTestContext()) + using (var tokenSource = new CancellationTokenSource()) + { + Tracing trace = tc.GetTrace(); + + // Arrange. + var expectedSession = new TaskAgentSession(); + PropertyInfo sessionIdProperty = expectedSession.GetType().GetProperty("SessionId", BindingFlags.Instance | BindingFlags.NonPublic | BindingFlags.Public); + Assert.NotNull(sessionIdProperty); + sessionIdProperty.SetValue(expectedSession, Guid.NewGuid()); + + _runnerServer + .Setup(x => x.CreateAgentSessionAsync( + _settings.PoolId, + It.Is(y => y != null), + tokenSource.Token)) + .Returns(Task.FromResult(expectedSession)); + + _credMgr.Setup(x => x.LoadCredentials()).Returns(new VssCredentials()); + _store.Setup(x => x.GetCredentials()).Returns(new CredentialData() { Scheme = Constants.Configuration.OAuthAccessToken }); + _store.Setup(x => x.GetMigratedCredentials()).Returns(default(CredentialData)); + + // Act. + MessageListener listener = new MessageListener(); + listener.Initialize(tc); + + bool result = await listener.CreateSessionAsync(tokenSource.Token); + Assert.True(result); + + _runnerServer + .Setup(x => x.GetAgentMessageAsync( + _settings.PoolId, expectedSession.SessionId, It.IsAny(), tokenSource.Token)) + .Throws(new TaskAgentAccessTokenExpiredException("test")); + try + { + await listener.GetNextMessageAsync(tokenSource.Token); + } + catch (TaskAgentAccessTokenExpiredException) + { + //expected + } + finally + { + await listener.DeleteSessionAsync(); + } + + //Assert + _runnerServer + .Verify(x => x.GetAgentMessageAsync( + _settings.PoolId, expectedSession.SessionId, It.IsAny(), tokenSource.Token), Times.Once); + + _runnerServer + .Verify(x => x.DeleteAgentSessionAsync( + _settings.PoolId, expectedSession.SessionId, It.IsAny()), Times.Never); + } + } } }