From 653d431078c8dfbba3decfe17e9e44f93d9ceda6 Mon Sep 17 00:00:00 2001 From: Cheena Malhotra <13396919+cheenamalhotra@users.noreply.github.com> Date: Wed, 23 Sep 2026 11:06:09 -0700 Subject: [PATCH] Fix token expiry eviction in connection pool V2 (#4734) Validate access token expiry before general checkout, preserving return and transaction-affinity behavior. Cover callback cache refresh and physical reuse across both pool implementations and sync/async opens. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../SqlConnection.xml | 11 +- .../ConnectionPool/ChannelDbConnectionPool.cs | 20 +- .../AADFedAuthTokenRefreshTest.cs | 107 ++++ .../DbConnectionPoolAccessTokenTest.cs | 484 ++++++++++++++++++ 4 files changed, 616 insertions(+), 6 deletions(-) create mode 100644 src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/DbConnectionPoolAccessTokenTest.cs diff --git a/doc/snippets/Microsoft.Data.SqlClient/SqlConnection.xml b/doc/snippets/Microsoft.Data.SqlClient/SqlConnection.xml index 7a9af8b6b4..d663d13a36 100644 --- a/doc/snippets/Microsoft.Data.SqlClient/SqlConnection.xml +++ b/doc/snippets/Microsoft.Data.SqlClient/SqlConnection.xml @@ -361,9 +361,14 @@ The following example creates a and a are the same, they will be grouped into the same connection pool. - When using a token callback function, the connection manages - refreshing the tokens returned by the callback. The application is - not responsible for knowing when tokens expire. + The driver manages token refresh and discards pooled connections with + expired or nearly expired tokens before reuse. Connections in use or + reused within the same active transaction are unaffected. + + + New physical connections invoke the callback only when a cached token + is missing or needs refreshing. Reusing a pooled connection does not + invoke the callback. This property is mutually exclusive with the diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPool.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPool.cs index 9017f3f920..6e60f6cdbb 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPool.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPool.cs @@ -704,7 +704,9 @@ private void PutConnectionInIdleChannel(DbConnectionInternal connection, bool pr connection.SetReturnedTime(_timeProvider.GetUtcNow().UtcDateTime); } - if (!IsLiveConnection(connection, probeLiveness)) + // Match V1: check token expiry on general checkout, not return. An expired connection + // may remain idle, but will be discarded before it can be reused outside its transaction. + if (!IsLiveConnection(connection, probeLiveness, checkAccessTokenExpiry: false)) { RemoveConnection(connection); return; @@ -1264,9 +1266,21 @@ _connectionCreationRateLimiter is not null && /// Whether to poll the physical connection to confirm it is still alive. Pass false when /// running on a thread that must not block; the remaining checks are all cheap and local. /// + /// + /// Validate the token before general checkout, but not when returning a connection to the pool. + /// /// Returns true if the connection is live and unexpired, otherwise returns false. - private bool IsLiveConnection(DbConnectionInternal connection, bool probeLiveness = true) + private bool IsLiveConnection(DbConnectionInternal connection, bool probeLiveness = true, bool checkAccessTokenExpiry = true) { + if (checkAccessTokenExpiry && connection.IsAccessTokenExpired) + { + SqlClientEventSource.Log.TryPoolerTraceEvent( + "ChannelDbConnectionPool.IsLiveConnection | INFO | {0}, Connection {1}, will not be reused because its access token has expired or is about to expire.", + Id, + connection.ObjectID); + return false; + } + // Connection has been sitting idle longer than the configured idle timeout. // Checked before the (potentially expensive) liveness probe so an idle-expired // connection is discarded without an SNI round-trip. @@ -1548,7 +1562,7 @@ private async Task GetInternalConnection( { // Skip the liveness/idle/generation gate at the bottom of the loop: // GetFromTransactedPool has already probed liveness, and a transacted - // connection is exempt from idle-timeout, load-balance and + // connection is exempt from token-expiry, idle-timeout, load-balance and // clear-generation eviction because closing it would abort its // (possibly distributed) transaction. break; diff --git a/src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/AADFedAuthTokenRefreshTest/AADFedAuthTokenRefreshTest.cs b/src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/AADFedAuthTokenRefreshTest/AADFedAuthTokenRefreshTest.cs index e5616776b1..c79b2edcb7 100644 --- a/src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/AADFedAuthTokenRefreshTest/AADFedAuthTokenRefreshTest.cs +++ b/src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/AADFedAuthTokenRefreshTest/AADFedAuthTokenRefreshTest.cs @@ -4,7 +4,13 @@ using System; using System.Diagnostics; +using System.Reflection; +using System.Threading; +using System.Threading.Tasks; +using Azure.Core; using Microsoft.Data.SqlClient.ManualTesting.Tests.SQL.Common.SystemDataInternals; +using Microsoft.Data.SqlClient.ManualTesting.Tests.SystemDataInternals; +using Microsoft.Data.SqlClient.Tests.Common; using Xunit; using Xunit.Abstractions; @@ -87,6 +93,107 @@ public void FedAuthTokenRefreshTest() } } + /// + /// Verifies both pools replace connections with expired or nearly expired tokens and + /// invoke the callback when the cached token also needs refreshing, for Open and OpenAsync. + /// + [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.IsAADPasswordConnStrSetup))] + [InlineData(false, false, -1)] + [InlineData(false, true, -1)] + [InlineData(true, false, -1)] + [InlineData(true, true, -1)] + [InlineData(false, false, 5)] + [InlineData(false, true, 5)] + [InlineData(true, false, 5)] + [InlineData(true, true, 5)] + public async Task AccessTokenCallback_PooledConnectionIsReplacedOnExpiry(bool usePoolV2, bool async, int expiresInSeconds) + { + using var poolVersion = new ConnectionPoolVersionScope(usePoolV2); + string[] credentialKeys = { "Authentication", "User ID", "Password", "UID", "PWD" }; + var builder = new SqlConnectionStringBuilder( + DataTestUtility.RemoveKeysInConnStr(DataTestUtility.AADPasswordConnectionString, credentialKeys)) + { + Pooling = true, + MinPoolSize = 0, + MaxPoolSize = 1, + ConnectTimeout = 30, + Enlist = false + }; + var credential = DataTestUtility.GetTokenCredential(); + SqlAuthenticationToken callbackToken = null; + int callbackInvocations = 0; + using var connection = new SqlConnection(builder.ConnectionString) + { + AccessTokenCallback = async (parameters, cancellationToken) => + { + Interlocked.Increment(ref callbackInvocations); + const string suffix = "/.default"; + string scope = parameters.Resource.EndsWith(suffix) ? parameters.Resource : parameters.Resource + suffix; + AccessToken token = await credential.GetTokenAsync(new TokenRequestContext(new[] { scope }), cancellationToken); + callbackToken = new SqlAuthenticationToken(token.Token, token.ExpiresOn); + return callbackToken; + } + }; + + Task OpenConnection() + { + if (async) + { + return connection.OpenAsync(); + } + connection.Open(); + return Task.CompletedTask; + } + + // The empty pool requires a physical login, which invokes the callback and caches its token. + await OpenConnection(); + object original = connection.GetInternalConnection(); + Assert.NotNull(callbackToken); + int callbackCountAfterLogin = callbackInvocations; + Assert.True(callbackCountAfterLogin > 0); + // Close returns the physical connection to the pool; reopening reuses it without authentication. + connection.Close(); + await OpenConnection(); + Assert.Same(original, connection.GetInternalConnection()); + Assert.Equal(callbackCountAfterLogin, callbackInvocations); + + // The connection's expiry controls eviction; the cached token's expiry controls callback refresh. + // Age both without changing the real token or waiting. Five seconds is within the 30-second + // checkout buffer; minus one second covers an already expired token. + DateTimeOffset expiry = DateTimeOffset.UtcNow.AddSeconds(expiresInSeconds); + object cachedContext = FedAuthTokenHelper.GetAuthenticationContextValue(connection); + FieldInfo cacheExpiryField = cachedContext.GetType().GetField("_expirationTime", BindingFlags.Instance | BindingFlags.NonPublic); + Assert.NotNull(cacheExpiryField); + cacheExpiryField.SetValue(cachedContext, expiry.UtcDateTime); + FieldInfo tokenField = original.GetType().GetField("_fedAuthToken", BindingFlags.Instance | BindingFlags.NonPublic); + Assert.NotNull(tokenField); + object expiringToken = Activator.CreateInstance(tokenField.FieldType, + BindingFlags.Instance | BindingFlags.NonPublic, binder: null, + args: new object[] { new SqlAuthenticationToken(callbackToken.AccessToken, expiry) }, + culture: null); + tokenField.SetValue(original, expiringToken); + // Expiry is checked on checkout, not return, so Close does not invoke the callback. + connection.Close(); + Assert.Equal(callbackCountAfterLogin, callbackInvocations); + + // Checkout discards the expired physical connection rather than reauthenticating it. + // Its replacement needs a login, and the expiring cache entry forces another callback. + // The credential may still return the same valid token; callback invocation is what matters. + await OpenConnection(); + Assert.NotSame(original, connection.GetInternalConnection()); + Assert.True(callbackInvocations > callbackCountAfterLogin); + object replacement = connection.GetInternalConnection(); + int callbackCountAfterRefresh = callbackInvocations; + // The replacement now has a valid token, so another reopen reuses it without another callback. + connection.Close(); + await OpenConnection(); + Assert.Same(replacement, connection.GetInternalConnection()); + Assert.Equal(callbackCountAfterRefresh, callbackInvocations); + using SqlCommand command = connection.CreateCommand(); + command.CommandText = "SELECT 1"; + Assert.Equal(1, async ? await command.ExecuteScalarAsync() : command.ExecuteScalar()); + } + [Conditional("DEBUG")] private void LogInfo(string message) { diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/DbConnectionPoolAccessTokenTest.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/DbConnectionPoolAccessTokenTest.cs new file mode 100644 index 0000000000..c6e810cdd2 --- /dev/null +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/DbConnectionPoolAccessTokenTest.cs @@ -0,0 +1,484 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. +// See the LICENSE file in the project root for more information. + +using System; +using System.Collections.Generic; +using System.Data.Common; +using System.Reflection; +using System.Threading; +using System.Threading.Tasks; +using System.Transactions; +using Microsoft.Data.ProviderBase; +using Microsoft.Data.SqlClient.ConnectionPool; +using Microsoft.Data.SqlClient.Tests.Common; +using Microsoft.SqlServer.TDS.PreLogin; +using Microsoft.SqlServer.TDS.Servers; +using Xunit; + +using static Microsoft.Data.SqlClient.UnitTests.ConnectionPool.PoolTestHarness; + +namespace Microsoft.Data.SqlClient.UnitTests.ConnectionPool +{ + /// + /// Verifies checkout-time token eviction separately from pool token-cache refresh. + /// The collection isolates pool-version switches; simulated logins need no Azure credentials. + /// + [Collection(SimulatedServerTestCollection.Name)] + public class DbConnectionPoolAccessTokenTest + { + private static readonly TimeSpan WaitTimeout = TimeSpan.FromSeconds(30); + + /// + /// Both pools defer token validation until checkout, reusing valid connections and replacing expired ones. + /// + [Theory] + [InlineData(PoolImplementation.WaitHandle, false, false)] + [InlineData(PoolImplementation.WaitHandle, false, true)] + [InlineData(PoolImplementation.WaitHandle, true, false)] + [InlineData(PoolImplementation.WaitHandle, true, true)] + [InlineData(PoolImplementation.Channel, false, false)] + [InlineData(PoolImplementation.Channel, false, true)] + [InlineData(PoolImplementation.Channel, true, false)] + [InlineData(PoolImplementation.Channel, true, true)] + public void Checkout_ValidatesIdleAccessToken(PoolImplementation implementation, bool async, bool expired) + { + using var fixture = new TokenPool(implementation); + using var owner = new SqlConnection(); + TokenConnection original = Request(fixture.Pool, owner, async); + original.Expired = expired; + int checks = original.ExpiryChecks; + fixture.Pool.ReturnInternalConnection(original, owner); + + // Like V1, returning a connection does not evaluate its token or refresh credentials. + Assert.Equal(checks, original.ExpiryChecks); + Assert.False(original.Disposed); + Assert.Equal(1, fixture.Pool.IdleCount); + + // This hits idle checkout. With one pool slot, eviction must free capacity for its replacement. + TokenConnection served = Request(fixture.Pool, owner, async); + Assert.Equal(!expired, ReferenceEquals(original, served)); + Assert.Equal(expired, original.Disposed); + Assert.True(original.ExpiryChecks > checks); + Assert.False(served.Expired); + Assert.Equal(1, fixture.Pool.Count); + fixture.Pool.ReturnInternalConnection(served, owner); + } + + /// + /// Freshly created connections must pass the expiry gate before activation, just like idle connections. + /// + [Theory] + [InlineData(PoolImplementation.WaitHandle, false)] + [InlineData(PoolImplementation.WaitHandle, true)] + [InlineData(PoolImplementation.Channel, false)] + [InlineData(PoolImplementation.Channel, true)] + public void Checkout_RejectsNewConnectionWithExpiredToken(PoolImplementation implementation, bool async) + { + // Only the first creation is expired, so retrying can succeed without waiting for time to pass. + using var fixture = new TokenPool(implementation, expireFirstCreation: true); + using var owner = new SqlConnection(); + TokenConnection served = Request(fixture.Pool, owner, async); + + Assert.Equal(2, fixture.Factory.Created.Count); + TokenConnection expired = fixture.Factory.Created[0]; + Assert.True(expired.Disposed); + // Activation would assign the connection to the caller; expiry must be rejected before then. + Assert.Equal(0, expired.Activations); + Assert.Same(fixture.Factory.Created[1], served); + Assert.Equal(1, fixture.Pool.Count); + fixture.Pool.ReturnInternalConnection(served, owner); + } + + /// + /// A direct channel handoff must validate expiry even though it bypasses the idle fast-path check. + /// + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task WaitingCheckout_RejectsExpiredToken(bool async) + { + using var fixture = new TokenPool(PoolImplementation.Channel); + var pool = (ChannelDbConnectionPool)fixture.Pool; + using var owner = new SqlConnection(); + using var waitingOwner = new SqlConnection(); + TokenConnection original = Request(pool, owner, async); + Task pending = Task.Run(() => Request(pool, waitingOwner, async)); + try + { + // Wait until the request is reading the channel, not merely scheduled on another thread. + // Returning the expired connection then exercises the post-wait gate, not idle lookup. + Assert.True(SpinWait.SpinUntil(() => pool.Reclaimer.ParkedWaiters == 1, WaitTimeout), + "The request did not reach the idle-channel wait."); + original.Expired = true; + pool.ReturnInternalConnection(original, owner); + + Assert.Same(pending, await Task.WhenAny(pending, Task.Delay(WaitTimeout))); + TokenConnection served = await pending; + Assert.NotSame(original, served); + Assert.True(original.Disposed); + Assert.False(served.Expired); + Assert.Equal(1, pool.Count); + pool.ReturnInternalConnection(served, waitingOwner); + } + finally + { + pool.Shutdown(); + Assert.Same(pending, await Task.WhenAny(pending, Task.Delay(WaitTimeout))); + } + } + + /// + /// Transaction affinity overrides expiry eviction until completion, avoiding disruption of the active transaction. + /// + [Theory] + [InlineData(PoolImplementation.WaitHandle, false)] + [InlineData(PoolImplementation.WaitHandle, true)] + [InlineData(PoolImplementation.Channel, false)] + [InlineData(PoolImplementation.Channel, true)] + public void TransactionCheckout_PreservesExpiredConnectionUntilTransactionEnds(PoolImplementation implementation, bool async) + { + using var fixture = new TokenPool(implementation); + using var owner = new SqlConnection(); + using var transaction = new CommittableTransaction(); + TokenConnection original; + using (var scope = new TransactionScope(transaction)) + { + original = Request(fixture.Pool, owner, async); + original.Expired = true; + int checks = original.ExpiryChecks; + fixture.Pool.ReturnInternalConnection(original, owner); + + // Return parks this connection in the transacted store, which takes precedence over idle reuse. + // The same transaction must get it back without checking expiry or breaking enlistment. + TokenConnection enlisted = Request(fixture.Pool, owner, async); + Assert.Same(original, enlisted); + Assert.Equal(checks, original.ExpiryChecks); + Assert.False(original.Disposed); + fixture.Pool.ReturnInternalConnection(enlisted, owner); + scope.Complete(); + } + transaction.Commit(); + + // Completion releases the connection to general circulation, where expiry eviction applies again. + TokenConnection served = Request(fixture.Pool, owner, async); + Assert.NotSame(original, served); + Assert.True(original.Disposed); + fixture.Pool.ReturnInternalConnection(served, owner); + } + + /// + /// Physical-connection expiry triggers replacement; only an expiring cached token requires another callback. + /// + [Theory] + [InlineData(false, false, -1, false)] + [InlineData(false, true, -1, false)] + [InlineData(true, false, -1, false)] + [InlineData(true, true, -1, false)] + [InlineData(false, false, 300, false)] + [InlineData(false, true, 300, false)] + [InlineData(true, false, 300, false)] + [InlineData(true, true, 300, false)] + [InlineData(false, false, -1, true)] + [InlineData(false, true, -1, true)] + [InlineData(true, false, -1, true)] + [InlineData(true, true, -1, true)] + [InlineData(false, false, 300, true)] + [InlineData(false, true, 300, true)] + [InlineData(true, false, 300, true)] + [InlineData(true, true, 300, true)] + public async Task AccessTokenCallback_CheckoutRejectsExpiredToken(bool usePoolV2, bool async, int expiresInSeconds, bool expireCachedToken) + { + using var poolVersion = new ConnectionPoolVersionScope(usePoolV2); + using var server = new TdsServer(new TdsServerArguments + { + FedAuthRequiredPreLoginOption = TdsPreLoginFedAuthRequiredOption.FedAuthRequired + }); + server.Start(); + var builder = new SqlConnectionStringBuilder + { + DataSource = $"localhost,{server.EndPoint.Port}", + Encrypt = SqlConnectionEncryptOption.Optional, + MaxPoolSize = 1, + ConnectTimeout = 600, + Enlist = false + }; + int callbackInvocations = 0; + using var connection = new SqlConnection(builder.ConnectionString) + { + AccessTokenCallback = (_, _) => + { + Interlocked.Increment(ref callbackInvocations); + return Task.FromResult(new SqlAuthenticationToken("invalid", DateTimeOffset.UtcNow.AddHours(2))); + } + }; + + try + { + // The first physical login populates the token cache; ordinary reopen only reuses the socket. + await OpenConnection(connection, async); + var original = Assert.IsType(connection.InnerConnection); + Assert.False(original.IsAccessTokenExpired); + Assert.Equal(1, callbackInvocations); + connection.Close(); + await OpenConnection(connection, async); + Assert.Same(original, connection.InnerConnection); + Assert.Equal(1, callbackInvocations); + + // Change metadata, not wall-clock time: -1 is expired and 300 is within the 600-second buffer. + // The physical connection and the pool cache hold separate expiry values. + FieldInfo? tokenField = original.GetType().GetField("_fedAuthToken", BindingFlags.Instance | BindingFlags.NonPublic); + Assert.NotNull(tokenField); + DateTimeOffset expiry = DateTimeOffset.UtcNow.AddSeconds(expiresInSeconds); + tokenField!.SetValue(original, new SqlFedAuthToken(new SqlAuthenticationToken("invalid", expiry))); + Assert.Equal(expiresInSeconds > 0, expiry > DateTimeOffset.UtcNow); + Assert.True(original.IsAccessTokenExpired); + IDbConnectionPool pool = original.Pool; + var cachedToken = Assert.Single(pool.AuthenticationContexts); + if (expireCachedToken) + { + // Leaving the cache fresh models another login already having refreshed the pool's token. + // Aging it instead requires the replacement login to invoke the callback. + pool.AuthenticationContexts[cachedToken.Key] = new DbConnectionPoolAuthenticationContext( + cachedToken.Value.AccessToken, expiry.UtcDateTime); + } + connection.Close(); + Assert.Equal(1, callbackInvocations); + + // Checkout evicts the old physical connection. Its replacement either uses the fresh cache + // or invokes the callback and updates the cache; eviction alone must not force token acquisition. + await OpenConnection(connection, async); + Assert.NotSame(original, connection.InnerConnection); + Assert.False(connection.InnerConnection.IsAccessTokenExpired); + Assert.Equal(expireCachedToken ? 2 : 1, callbackInvocations); + Assert.True(pool.AuthenticationContexts[cachedToken.Key].ExpirationTime > expiry.UtcDateTime); + + // Once replaced, reuse must neither create another physical connection nor invoke the callback. + DbConnectionInternal replacement = connection.InnerConnection; + connection.Close(); + await OpenConnection(connection, async); + Assert.Same(replacement, connection.InnerConnection); + Assert.Equal(expireCachedToken ? 2 : 1, callbackInvocations); + } + finally + { + SqlConnection.ClearPool(connection); + } + } + + /// + /// New physical logins refresh expired or nearly expired cached tokens but reuse sufficiently valid ones. + /// + [Theory] + [InlineData(false, false, -1)] + [InlineData(false, true, -1)] + [InlineData(true, false, -1)] + [InlineData(true, true, -1)] + [InlineData(false, false, 300)] + [InlineData(false, true, 300)] + [InlineData(true, false, 300)] + [InlineData(true, true, 300)] + [InlineData(false, false, 3600)] + [InlineData(false, true, 3600)] + [InlineData(true, false, 3600)] + [InlineData(true, true, 3600)] + public async Task AccessTokenCallback_NewPhysicalConnectionHonorsCachedTokenExpiry(bool usePoolV2, bool async, int expiresInSeconds) + { + using var poolVersion = new ConnectionPoolVersionScope(usePoolV2); + using var server = new TdsServer(new TdsServerArguments + { + FedAuthRequiredPreLoginOption = TdsPreLoginFedAuthRequiredOption.FedAuthRequired + }); + server.Start(); + var builder = new SqlConnectionStringBuilder + { + DataSource = $"localhost,{server.EndPoint.Port}", + Encrypt = SqlConnectionEncryptOption.Optional, + MaxPoolSize = 2, + Enlist = false + }; + int callbackInvocations = 0; + using var first = new SqlConnection(builder.ConnectionString) + { + AccessTokenCallback = (_, _) => + { + int invocation = Interlocked.Increment(ref callbackInvocations); + return Task.FromResult(new SqlAuthenticationToken($"invalid-{invocation}", DateTimeOffset.UtcNow.AddHours(2))); + } + }; + // The callback delegate is part of the pool key; share it to exercise the same token cache. + using var second = new SqlConnection(builder.ConnectionString) + { + AccessTokenCallback = first.AccessTokenCallback + }; + + try + { + await OpenConnection(first, async); + Assert.Equal(1, callbackInvocations); + IDbConnectionPool pool = first.InnerConnection.Pool; + var entry = Assert.Single(pool.AuthenticationContexts); + var cachedToken = new DbConnectionPoolAuthenticationContext( + entry.Value.AccessToken, DateTime.UtcNow.AddSeconds(expiresInSeconds)); + pool.AuthenticationContexts[entry.Key] = cachedToken; + + // Keep the first connection checked out so the second must perform a physical login. + await OpenConnection(second, async); + Assert.NotSame(first.InnerConnection, second.InnerConnection); + Assert.Same(pool, second.InnerConnection.Pool); + Assert.False(second.InnerConnection.IsAccessTokenExpired); + // -1 and 300 seconds require refresh (the cache's 10-minute window); 3600 seconds + // is beyond even its 45-minute opportunistic refresh window, so it must reuse the token. + bool refreshed = expiresInSeconds <= 600; + Assert.Equal(refreshed ? 2 : 1, callbackInvocations); + DbConnectionPoolAuthenticationContext current = pool.AuthenticationContexts[entry.Key]; + if (refreshed) + { + Assert.NotSame(cachedToken, current); + Assert.NotEqual(cachedToken.AccessToken, current.AccessToken); + Assert.True(current.ExpirationTime > cachedToken.ExpirationTime); + } + else + { + Assert.Same(cachedToken, current); + } + + // Physical reuse skips authentication entirely, regardless of which cache branch ran above. + DbConnectionInternal reused = second.InnerConnection; + second.Close(); + await OpenConnection(second, async); + Assert.Same(reused, second.InnerConnection); + Assert.Equal(refreshed ? 2 : 1, callbackInvocations); + } + finally + { + SqlConnection.ClearPool(first); + } + } + + /// + /// Exercises the selected public open API, bounding asynchronous opens with cancellation. + /// + /// Connection to open against the simulated server. + /// Whether to use OpenAsync instead of Open. + /// A task completing when the connection is open. + private static async Task OpenConnection(SqlConnection connection, bool async) + { + if (async) + { + using var cancellation = new CancellationTokenSource(WaitTimeout); + await connection.OpenAsync(cancellation.Token); + } + else + { + connection.Open(); + } + } + + /// + /// Requests a stub connection directly from a pool, handling inline and deferred completion. + /// + /// Pool under test. + /// Owner passed to connection activation. + /// Whether to supply the completion source used by asynchronous opens. + /// The connection assigned to the owner. + private static TokenConnection Request(IDbConnectionPool pool, SqlConnection owner, bool async) + { + // Async opens carry the ambient transaction in AsyncState so worker threads preserve affinity. + TaskCompletionSource? completion = async + ? new TaskCompletionSource(Transaction.Current, TaskCreationOptions.RunContinuationsAsynchronously) + : null; + bool completed = pool.TryGetConnection(owner, completion, TimeoutTimer.StartNew(WaitTimeout), out DbConnectionInternal? connection); + if (!completed) + { + Assert.NotNull(completion); + Assert.True(completion!.Task.Wait(WaitTimeout), "The connection request did not complete."); + connection = completion.Task.GetAwaiter().GetResult(); + } + return Assert.IsType(connection); + } + + /// + /// Uses one pool slot and the harness's frozen clock to isolate checkout decisions from maintenance. + /// + private sealed class TokenPool : IDisposable + { + internal TokenFactory Factory { get; } + internal IDbConnectionPool Pool { get; } + + /// Creates an isolated pool backed by controllable token connections. + /// Pool implementation to exercise. + /// Whether the first created connection starts expired. + internal TokenPool(PoolImplementation implementation, bool expireFirstCreation = false) + { + Factory = new TokenFactory(expireFirstCreation); + Pool = ConstructPool(implementation, Factory, maxPoolSize: 1, creationTimeout: 30000); + } + + /// Stops maintenance and disposes idle connections and any left checked out by a failed test. + public void Dispose() + { + Pool.Shutdown(); + Pool.Clear(); + foreach (TokenConnection connection in Factory.Created) + { + if (!connection.Disposed) + { + connection.Dispose(); + } + } + } + } + + /// + /// Records physical creations and can expire the first one to force a checkout retry. + /// + private sealed class TokenFactory(bool expireFirstCreation) : SqlConnectionFactory + { + internal List Created { get; } = new(); + + /// + protected override DbConnectionInternal CreateConnection(SqlConnectionOptions options, ConnectionPoolKey poolKey, + DbConnectionPoolGroupProviderInfo poolGroupProviderInfo, IDbConnectionPool pool, DbConnection owningConnection, TimeoutTimer timeout) + { + var connection = new TokenConnection { Expired = expireFirstCreation && Created.Count == 0 }; + Created.Add(connection); + return connection; + } + } + + /// + /// Makes expiry deterministic and records whether the pool checks, activates, or disposes the connection. + /// + private sealed class TokenConnection : ChannelDbConnectionPoolTest.StubDbConnectionInternal + { + internal bool Expired { get; set; } + internal int ExpiryChecks { get; private set; } + internal int Activations { get; private set; } + internal bool Disposed { get; private set; } + + internal override bool IsAccessTokenExpired + { + get + { + ExpiryChecks++; + return Expired; + } + } + + /// + protected override void Activate(Transaction transaction) + { + Activations++; + base.Activate(transaction); + } + + /// + public override void Dispose() + { + Disposed = true; + base.Dispose(); + } + } + } +}