Fix AccessTokenCallback to match AccessToken behavior for TNIR and connection pool keys

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/Connection/SqlConnectionInternal.cs

@@ -145,6 +145,19 @@ internal class SqlConnectionInternal : DbConnectionInternal, IDisposable
         // @TODO: Probably a good idea to introduce a delegate type
         internal readonly Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticationToken>> _accessTokenCallback;
 
+        /// <summary>
+        /// True when the caller supplied a federated authentication access token directly, either
+        /// as a literal token via <see cref="global::Microsoft.Data.SqlClient.SqlConnection.AccessToken"/> or as a token provider
+        /// via <see cref="global::Microsoft.Data.SqlClient.SqlConnection.AccessTokenCallback"/>.
+        /// </summary>
+        /// <remarks>
+        /// Both paths represent the same "caller-supplied token" authentication mode, so they must
+        /// always be treated identically. Use this property rather than testing the underlying
+        /// fields individually.
+        /// </remarks>
+        internal bool IsAccessTokenProvided =>
+            _accessTokenInBytes != null || _accessTokenCallback != null;
+
         // @TODO: Should be private and accessed via internal property
         // @TODO: Rename to match naming conventions
         internal bool _cleanSQLDNSCaching = false;
@@ -3115,7 +3128,9 @@ private void LoginNoFailover(
             #if NET
             bool isParallel = connectionOptions.MultiSubnetFailover;
             #else
-            bool disableTnir = ShouldDisableTnir(connectionOptions);
+            bool disableTnir = ShouldDisableTnir(
+                connectionOptions,
+                isAccessTokenProvided: IsAccessTokenProvided);
             bool isParallel = connectionOptions.MultiSubnetFailover ||
                               (connectionOptions.TransparentNetworkIPResolution && !disableTnir);
             #endif
@@ -3911,12 +3926,29 @@ private void ResolveExtendedServerName(ServerInfo serverInfo, bool aliasLookup,
         }
 
         #if NETFRAMEWORK
-        private bool ShouldDisableTnir(SqlConnectionOptions connectionOptions)
+        /// <summary>
+        /// Determines whether Transparent Network IP Resolution (TNIR) should be disabled for this
+        /// connection attempt.
+        /// </summary>
+        /// <param name="connectionOptions">The parsed connection options.</param>
+        /// <param name="isAccessTokenProvided">
+        /// True when the caller supplied a federated authentication access token directly, either
+        /// via <see cref="SqlConnection.AccessToken"/> or
+        /// <see cref="SqlConnection.AccessTokenCallback"/>.
+        /// </param>
+        /// <returns>
+        /// True when TNIR should be disabled. TNIR is disabled by default for Azure SQL endpoints
+        /// and for federated authentication, but an explicit
+        /// <c>TransparentNetworkIPResolution</c> keyword always takes precedence.
+        /// </returns>
+        internal static bool ShouldDisableTnir(
+            SqlConnectionOptions connectionOptions,
+            bool isAccessTokenProvided)
         {
             bool isAzureEndPoint = ADP.IsAzureSqlServerEndpoint(connectionOptions.DataSource);
 
             // @TODO: Turn into a HashSet and just check the list instead of this MESS.
-            bool isFedAuthEnabled = _accessTokenInBytes != null ||
+            bool isFedAuthEnabled = isAccessTokenProvided ||
                                     #pragma warning disable 0618 // Type or member is obsolete
                                     connectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryPassword ||
                                     #pragma warning restore 0618 // Type or member is obsolete

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlConnection.cs

@@ -763,8 +763,15 @@ public string AccessToken
                     CheckAndThrowOnInvalidCombinationOfConnectionOptionAndAccessToken(ConnectionOptions);
                 }
 
-                // Need to call ConnectionString_Set to do proper pool group check
-                ConnectionString_Set(new ConnectionPoolKey(_connectionString, credential: _credential, accessToken: value, accessTokenCallback: null, sspiContextProvider: null));
+                // Need to call ConnectionString_Set to do proper pool group check.
+                // Preserve the other authentication state so it isn't dropped from the pool key
+                // (see the ConnectionString setter, which is the reference for this pattern).
+                ConnectionString_Set(new ConnectionPoolKey(
+                    _connectionString,
+                    credential: _credential,
+                    accessToken: value,
+                    accessTokenCallback: _accessTokenCallback,
+                    sspiContextProvider: _sspiContextProvider));
                 _accessToken = value;
             }
         }
@@ -787,7 +794,12 @@ public Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticati
                     CheckAndThrowOnInvalidCombinationOfConnectionOptionAndAccessTokenCallback(ConnectionOptions);
                 }
 
-                ConnectionString_Set(new ConnectionPoolKey(_connectionString, credential: _credential, accessToken: null, accessTokenCallback: value, sspiContextProvider: null));
+                ConnectionString_Set(new ConnectionPoolKey(
+                    _connectionString,
+                    credential: _credential,
+                    accessToken: _accessToken,
+                    accessTokenCallback: value,
+                    sspiContextProvider: _sspiContextProvider));
                 _accessTokenCallback = value;
             }
         }
@@ -804,7 +816,12 @@ public SspiContextProvider SspiContextProvider
                     throw ADP.OpenConnectionPropertySet(nameof(SspiContextProvider), InnerConnection.State);
                 }
 
-                ConnectionString_Set(new ConnectionPoolKey(_connectionString, credential: _credential, accessToken: null, accessTokenCallback: null, sspiContextProvider: value));
+                ConnectionString_Set(new ConnectionPoolKey(
+                    _connectionString,
+                    credential: _credential,
+                    accessToken: _accessToken,
+                    accessTokenCallback: _accessTokenCallback,
+                    sspiContextProvider: value));
                 _sspiContextProvider = value;
             }
         }

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsParser.cs

@@ -1179,10 +1179,10 @@ private PreLoginHandshakeStatus ConsumePreLoginHandshake(
 
                         // We must NOT use the response for the FEDAUTHREQUIRED PreLogin option, if the connection string option
                         // was not using the new Authentication keyword or in other words, if Authentication=NotSpecified
-                        // Or AccessToken is not null, mean token based authentication is used.
+                        // Or an access token was supplied (AccessToken/AccessTokenCallback), which means token-based authentication is used.
                         if ((_connHandler.ConnectionOptions != null
                             && _connHandler.ConnectionOptions.Authentication != SqlAuthenticationMethod.NotSpecified)
-                            || _connHandler._accessTokenInBytes != null || _connHandler._accessTokenCallback != null)
+                            || _connHandler.IsAccessTokenProvided)
                         {
                             fedAuthRequired = payload[payloadOffset] == 0x01 ? true : false;
                         }
@@ -1219,7 +1219,7 @@ private PreLoginHandshakeStatus ConsumePreLoginHandshake(
 
                 // Validate Certificate if Trust Server Certificate=false and Encryption forced (EncryptionOptions.ON) from Server.
                 bool shouldValidateServerCert = (_encryptionOption == EncryptionOptions.ON && !trustServerCert) ||
-                    ((_connHandler._accessTokenInBytes != null || _connHandler._accessTokenCallback != null) && !trustServerCert);
+                    (_connHandler.IsAccessTokenProvided && !trustServerCert);
 
                 uint info = (shouldValidateServerCert ? TdsEnums.SNI_SSL_VALIDATE_CERTIFICATE : 0)
                     | TdsEnums.SNI_SSL_USE_SCHANNEL_CACHE;

src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/SqlConnectionOptionsTest.cs

@@ -5,6 +5,9 @@
 using System;
 using Microsoft.Data.SqlClient.Tests.Common;
 using Xunit;
+#if NETFRAMEWORK
+using SqlConnectionInternal = global::Microsoft.Data.SqlClient.Connection.SqlConnectionInternal;
+#endif
 
 namespace Microsoft.Data.SqlClient.UnitTests.Microsoft.Data.SqlClient
 {
@@ -63,6 +66,41 @@ public void TestDefaultTnir(string dataSource, bool? tnirEnabledInConnString, bo
             // Assert
             Assert.Equal(expectedValue, connectionString.TransparentNetworkIPResolution);
         }
+
+        /// <summary>
+        /// TNIR is disabled by default whenever federated authentication is in play, including when
+        /// the token is supplied directly through <c>AccessToken</c> or <c>AccessTokenCallback</c>,
+        /// unless the user explicitly specified the TNIR keyword.
+        /// </summary>
+        [Theory]
+        // Non-Azure endpoint, no explicit TNIR keyword: access token (or callback) disables TNIR.
+        [InlineData("my.test.server", false, false, false)]
+        [InlineData("my.test.server", true, false, true)]
+        // Azure endpoint always disables TNIR when the keyword is absent.
+        [InlineData("test.database.windows.net", false, false, true)]
+        [InlineData("test.database.windows.net", true, false, true)]
+        // An explicit TNIR keyword always wins, regardless of access token or endpoint.
+        [InlineData("my.test.server", true, true, false)]
+        [InlineData("test.database.windows.net", true, true, false)]
+        [InlineData("test.database.windows.net", false, true, false)]
+        public void TestShouldDisableTnirWithAccessToken(
+            string dataSource,
+            bool isAccessTokenProvided,
+            bool tnirExplicitlySpecified,
+            bool expectedValue)
+        {
+            SqlConnectionStringBuilder builder = new() { DataSource = dataSource };
+            if (tnirExplicitlySpecified)
+            {
+                builder.TransparentNetworkIPResolution = true;
+            }
+
+            SqlConnectionOptions connectionOptions = new(builder.ConnectionString);
+
+            Assert.Equal(
+                expectedValue,
+                SqlConnectionInternal.ShouldDisableTnir(connectionOptions, isAccessTokenProvided));
+        }
 #endif
         /// <summary>
         /// Test MSF values when set through connection string and through app context switch.

src/Microsoft.Data.SqlClient/tests/UnitTests/SimulatedServerTests/ConnectionTests.cs

@@ -764,6 +764,65 @@ public void ConnectionTestAccessTokenCallbackCombinations()
             }
         }
 
+        /// <summary>
+        /// Setting one authentication-related property must not silently drop the others from the
+        /// connection pool key. Previously, assigning <see cref="SqlConnection.SspiContextProvider"/>
+        /// rebuilt the pool key with a null <see cref="SqlConnection.AccessTokenCallback"/> (and a
+        /// null <see cref="SqlConnection.AccessToken"/>), so the token was never handed to the
+        /// internal connection even though the public property still reported it as set.
+        /// </summary>
+        [Fact]
+        public void AccessTokenStateIsPreservedInPoolKeyWhenSspiContextProviderIsSet()
+        {
+            Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticationToken>> callback =
+                (ctx, token) => Task.FromResult(new SqlAuthenticationToken("invalid", DateTimeOffset.MaxValue));
+
+            using (SqlConnection conn = new("Data Source=localhost"))
+            {
+                conn.AccessTokenCallback = callback;
+                Assert.Same(callback, conn.PoolGroup.PoolKey.AccessTokenCallback);
+
+                conn.SspiContextProvider = null;
+
+                Assert.Same(callback, conn.AccessTokenCallback);
+                Assert.Same(callback, conn.PoolGroup.PoolKey.AccessTokenCallback);
+            }
+
+            using (SqlConnection conn = new("Data Source=localhost"))
+            {
+                conn.AccessToken = "token";
+                Assert.Equal("token", conn.PoolGroup.PoolKey.AccessToken);
+
+                conn.SspiContextProvider = null;
+
+                Assert.Equal("token", conn.AccessToken);
+                Assert.Equal("token", conn.PoolGroup.PoolKey.AccessToken);
+            }
+        }
+
+        /// <summary>
+        /// <see cref="SqlConnection.AccessToken"/> and <see cref="SqlConnection.AccessTokenCallback"/>
+        /// are mutually exclusive, so neither setter can ever clobber a live value of the other.
+        /// </summary>
+        [Fact]
+        public void AccessTokenAndAccessTokenCallbackAreMutuallyExclusive()
+        {
+            Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticationToken>> callback =
+                (ctx, token) => Task.FromResult(new SqlAuthenticationToken("invalid", DateTimeOffset.MaxValue));
+
+            using (SqlConnection conn = new("Data Source=localhost"))
+            {
+                conn.AccessTokenCallback = callback;
+                Assert.Throws<InvalidOperationException>(() => conn.AccessToken = "token");
+            }
+
+            using (SqlConnection conn = new("Data Source=localhost"))
+            {
+                conn.AccessToken = "token";
+                Assert.Throws<InvalidOperationException>(() => conn.AccessTokenCallback = callback);
+            }
+        }
+
         [Theory]
         [InlineData(9, 0, 2047)] // SQL Server 2005
         [InlineData(10, 0, 2531)] // SQL Server 2008