Fix AccessTokenCallback TNIR behavior and make SspiContextProvider mutually exclusive with token auth

src/Microsoft.Data.SqlClient/src/Microsoft/Data/Common/AdapterUtil.cs

@@ -1487,6 +1487,12 @@ internal static Exception InvalidMixedUsageOfAccessTokenCallbackAndAuthenticatio
 
         internal static Exception InvalidMixedUsageOfAccessTokenCallbackAndIntegratedSecurity()
             => InvalidOperation(StringsHelper.GetString(Strings.ADP_InvalidMixedUsageOfAccessTokenCallbackAndIntegratedSecurity));
+
+        internal static Exception InvalidMixedUsageOfAccessTokenAndSspiContextProvider()
+            => InvalidOperation(StringsHelper.GetString(Strings.ADP_InvalidMixedUsageOfAccessTokenAndSspiContextProvider));
+
+        internal static Exception InvalidMixedUsageOfSspiContextProviderAndAccessToken()
+            => InvalidOperation(StringsHelper.GetString(Strings.ADP_InvalidMixedUsageOfSspiContextProviderAndAccessToken));
         #endregion
 
         internal static readonly IntPtr s_ptrZero = IntPtr.Zero;

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

@@ -145,6 +145,42 @@ 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>
+        /// <para>
+        /// Use this property wherever the only question is whether the caller supplied a token at
+        /// all, and the two paths are equivalent: the prelogin FEDAUTHREQUIRED handshake, server
+        /// certificate validation, and Transparent Network IP Resolution. Those all key off the
+        /// authentication <em>mode</em>, so treating the two paths differently would be a bug.
+        /// </para>
+        /// <para>
+        /// Do <b>not</b> use this property where the two paths diverge. Login feature-extension
+        /// negotiation must keep testing the fields individually because they select different
+        /// federated authentication library types: <c>_accessTokenCallback</c> requests
+        /// <see cref="TdsEnums.FedAuthLibrary.MSAL"/>, while <c>_accessTokenInBytes</c> requests
+        /// <see cref="TdsEnums.FedAuthLibrary.SecurityToken"/> and carries the token bytes.
+        /// </para>
+        /// </remarks>
+        internal bool IsAccessTokenProvided =>
+            _accessTokenInBytes != null || _accessTokenCallback != null;
+
+        #if NETFRAMEWORK
+        /// <summary>
+        /// The Transparent Network IP Resolution decision applied by the most recent
+        /// <see cref="LoginNoFailover"/> call, or <see langword="null"/> if login has not run.
+        /// </summary>
+        /// <remarks>
+        /// Exposed for tests. <see cref="ShouldDisableTnir"/> is pure and can be tested directly,
+        /// but that leaves the wiring, in particular that <see cref="IsAccessTokenProvided"/> is
+        /// the value fed into it, unverified. This records the decision that was actually used.
+        /// </remarks>
+        internal bool? TnirDisabledDuringLogin { get; private set; }
+        #endif
+
         // @TODO: Should be private and accessed via internal property
         // @TODO: Rename to match naming conventions
         internal bool _cleanSQLDNSCaching = false;
@@ -3115,7 +3151,10 @@ private void LoginNoFailover(
             #if NET
             bool isParallel = connectionOptions.MultiSubnetFailover;
             #else
-            bool disableTnir = ShouldDisableTnir(connectionOptions);
+            bool disableTnir = ShouldDisableTnir(
+                connectionOptions,
+                isAccessTokenProvided: IsAccessTokenProvided);
+            TnirDisabledDuringLogin = disableTnir;
             bool isParallel = connectionOptions.MultiSubnetFailover ||
                               (connectionOptions.TransparentNetworkIPResolution && !disableTnir);
             #endif
@@ -3911,12 +3950,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="global::Microsoft.Data.SqlClient.SqlConnection.AccessToken"/> or
+        /// <see cref="global::Microsoft.Data.SqlClient.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

@@ -262,6 +262,13 @@ private SqlConnection(SqlConnection connection)
 
             _accessToken = connection._accessToken;
             _accessTokenCallback = connection._accessTokenCallback;
+
+            // CopyFrom retains the source PoolGroup, and therefore the source ConnectionPoolKey.
+            // The provider must be copied along with it, otherwise the clone would authenticate
+            // with a provider that its public property does not report, and would let the caller
+            // set AccessToken/AccessTokenCallback without tripping the mutual-exclusivity checks.
+            _sspiContextProvider = connection._sspiContextProvider;
+
             CacheConnectionStringProperties();
         }
 
@@ -763,8 +770,16 @@ 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.
+                // AccessTokenCallback and SspiContextProvider are mutually exclusive with AccessToken
+                // (validated above), so they are always null here when a token is supplied. Passing the
+                // current values through matters only when the token is being cleared.
+                ConnectionString_Set(new ConnectionPoolKey(
+                    _connectionString,
+                    credential: _credential,
+                    accessToken: value,
+                    accessTokenCallback: _accessTokenCallback,
+                    sspiContextProvider: _sspiContextProvider));
                 _accessToken = value;
             }
         }
@@ -787,7 +802,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 +824,21 @@ public SspiContextProvider SspiContextProvider
                     throw ADP.OpenConnectionPropertySet(nameof(SspiContextProvider), InnerConnection.State);
                 }
 
-                ConnectionString_Set(new ConnectionPoolKey(_connectionString, credential: _credential, accessToken: null, accessTokenCallback: null, sspiContextProvider: value));
+                if (value != null)
+                {
+                    // SSPI is an alternative to token-based authentication, so the two are mutually exclusive.
+                    CheckAndThrowOnInvalidCombinationOfConnectionOptionAndSspiContextProvider();
+                }
+
+                // Because SSPI and token authentication are mutually exclusive (validated above), the token
+                // state is always null when a provider is supplied. Passing the current values through
+                // matters only when the provider is being cleared, where any token state must survive.
+                ConnectionString_Set(new ConnectionPoolKey(
+                    _connectionString,
+                    credential: _credential,
+                    accessToken: _accessToken,
+                    accessTokenCallback: _accessTokenCallback,
+                    sspiContextProvider: value));
                 _sspiContextProvider = value;
             }
         }
@@ -1101,7 +1135,12 @@ public SqlCredential Credential
                 _credential = value;
 
                 // Need to call ConnectionString_Set to do proper pool group check
-                ConnectionString_Set(new ConnectionPoolKey(_connectionString, _credential, accessToken: _accessToken, accessTokenCallback: _accessTokenCallback, sspiContextProvider: null));
+                ConnectionString_Set(new ConnectionPoolKey(
+                    _connectionString,
+                    _credential,
+                    accessToken: _accessToken,
+                    accessTokenCallback: _accessTokenCallback,
+                    sspiContextProvider: _sspiContextProvider));
             }
         }
 
@@ -1157,6 +1196,11 @@ private void CheckAndThrowOnInvalidCombinationOfConnectionOptionAndAccessToken(S
             {
                 throw ADP.InvalidMixedUsageOfAccessTokenAndTokenCallback();
             }
+
+            if (_sspiContextProvider != null)
+            {
+                throw ADP.InvalidMixedUsageOfAccessTokenAndSspiContextProvider();
+            }
         }
 
         // CheckAndThrowOnInvalidCombinationOfConnectionOptionAndAccessTokenCallback: check if the usage of AccessTokenCallback has any conflict
@@ -1179,6 +1223,23 @@ private void CheckAndThrowOnInvalidCombinationOfConnectionOptionAndAccessTokenCa
             {
                 throw ADP.InvalidMixedUsageOfAccessTokenAndTokenCallback();
             }
+
+            if (_sspiContextProvider != null)
+            {
+                throw ADP.InvalidMixedUsageOfAccessTokenAndSspiContextProvider();
+            }
+        }
+
+        // CheckAndThrowOnInvalidCombinationOfConnectionOptionAndSspiContextProvider: SSPI is an alternative to
+        //  token-based authentication, so a context provider cannot be combined with AccessToken or
+        //  AccessTokenCallback. If there is any conflict, it throws InvalidOperationException.
+        //  This is to be used by the setter of the SspiContextProvider property.
+        private void CheckAndThrowOnInvalidCombinationOfConnectionOptionAndSspiContextProvider()
+        {
+            if (_accessToken != null || _accessTokenCallback != null)
+            {
+                throw ADP.InvalidMixedUsageOfSspiContextProviderAndAccessToken();
+            }
         }
 
         /// <include file='../../../../../../doc/snippets/Microsoft.Data.SqlClient/SqlConnection.xml' path='docs/members[@name="SqlConnection"]/DbProviderFactory/*' />

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/src/Resources/Strings.resx

@@ -2079,6 +2079,12 @@
   <data name="ADP_InvalidMixedUsageOfAccessTokenCallbackAndIntegratedSecurity" xml:space="preserve">
     <value>Cannot set the AccessTokenCallback property if the 'Integrated Security' connection string keyword has been set to 'true' or 'SSPI'.</value>
   </data>
+  <data name="ADP_InvalidMixedUsageOfAccessTokenAndSspiContextProvider" xml:space="preserve">
+    <value>Cannot set the AccessToken or AccessTokenCallback property if the SspiContextProvider property has been set.</value>
+  </data>
+  <data name="ADP_InvalidMixedUsageOfSspiContextProviderAndAccessToken" xml:space="preserve">
+    <value>Cannot set the SspiContextProvider property if the AccessToken or AccessTokenCallback property has been set.</value>
+  </data>
   <data name="ADP_InvalidMixedUsageOfAuthenticationAndTokenCallback" xml:space="preserve">
     <value>Cannot set the AccessTokenCallback property if 'Authentication=Active Directory Default' has been specified in the connection string.</value>
   </data>

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,48 @@ 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, no token: TNIR stays enabled.
+        [InlineData("my.test.server", false, null, false)]
+        // Non-Azure endpoint, no explicit TNIR keyword, token supplied: TNIR is disabled.
+        [InlineData("my.test.server", true, null, true)]
+        // Azure endpoint always disables TNIR when the keyword is absent.
+        [InlineData("test.database.windows.net", false, null, true)]
+        [InlineData("test.database.windows.net", true, null, true)]
+        // An explicit TNIR keyword always wins, regardless of access token or endpoint. Note that
+        // this method only decides whether the default is overridden; an explicit false is honoured
+        // by the caller through SqlConnectionOptions.TransparentNetworkIPResolution itself.
+        [InlineData("my.test.server", true, true, false)]
+        [InlineData("test.database.windows.net", true, true, false)]
+        [InlineData("test.database.windows.net", false, true, false)]
+        [InlineData("my.test.server", true, false, false)]
+        [InlineData("my.test.server", false, false, false)]
+        [InlineData("test.database.windows.net", true, false, false)]
+        [InlineData("test.database.windows.net", false, false, false)]
+        public void TestShouldDisableTnirWithCallerSuppliedToken(
+            string dataSource,
+            bool isAccessTokenProvided,
+            bool? tnirInConnString,
+            bool expectedValue)
+        {
+            SqlConnectionStringBuilder builder = new() { DataSource = dataSource };
+            if (tnirInConnString.HasValue)
+            {
+                builder.TransparentNetworkIPResolution = tnirInConnString.Value;
+            }
+
+            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.