Skip to content

Commit 9ea0d18

Browse files
github-actions[bot]CopilotCopilotcheenamalhotra
authored
[6.1.7 Cherry-pick] Fix AccessTokenCallback TNIR behavior and token pool keys (#4560)
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Cheena Malhotra <13396919+cheenamalhotra@users.noreply.github.com>
1 parent 596bd92 commit 9ea0d18

8 files changed

Lines changed: 156 additions & 13 deletions

File tree

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -710,7 +710,7 @@ public string AccessToken
710710
}
711711

712712
// Need to call ConnectionString_Set to do proper pool group check
713-
ConnectionString_Set(new SqlConnectionPoolKey(_connectionString, credential: _credential, accessToken: value, accessTokenCallback: null));
713+
ConnectionString_Set(new SqlConnectionPoolKey(_connectionString, credential: _credential, accessToken: value, accessTokenCallback: _accessTokenCallback));
714714
_accessToken = value;
715715
}
716716
}
@@ -733,7 +733,7 @@ public Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticati
733733
CheckAndThrowOnInvalidCombinationOfConnectionOptionAndAccessTokenCallback((SqlConnectionString)ConnectionOptions);
734734
}
735735

736-
ConnectionString_Set(new SqlConnectionPoolKey(_connectionString, credential: _credential, accessToken: null, accessTokenCallback: value));
736+
ConnectionString_Set(new SqlConnectionPoolKey(_connectionString, credential: _credential, accessToken: _accessToken, accessTokenCallback: value));
737737
_accessTokenCallback = value;
738738
}
739739
}

‎src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlInternalConnectionTds.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,9 @@ internal sealed class SqlInternalConnectionTds : SqlInternalConnection, IDisposa
136136
internal byte[] _accessTokenInBytes;
137137
internal readonly Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticationToken>> _accessTokenCallback;
138138

139+
internal bool IsAccessTokenProvided =>
140+
_accessTokenInBytes != null || _accessTokenCallback != null;
141+
139142
private readonly ActiveDirectoryAuthenticationTimeoutRetryHelper _activeDirectoryAuthTimeoutRetryHelper;
140143

141144
internal bool _cleanSQLDNSCaching = false;

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1108,10 +1108,10 @@ private PreLoginHandshakeStatus ConsumePreLoginHandshake(
11081108

11091109
// We must NOT use the response for the FEDAUTHREQUIRED PreLogin option, if the connection string option
11101110
// was not using the new Authentication keyword or in other words, if Authentication=NotSpecified
1111-
// Or AccessToken is not null, mean token based authentication is used.
1111+
// Or an access token was supplied, which means token-based authentication is used.
11121112
if ((_connHandler.ConnectionOptions != null
11131113
&& _connHandler.ConnectionOptions.Authentication != SqlAuthenticationMethod.NotSpecified)
1114-
|| _connHandler._accessTokenInBytes != null || _connHandler._accessTokenCallback != null)
1114+
|| _connHandler.IsAccessTokenProvided)
11151115
{
11161116
fedAuthRequired = payload[payloadOffset] == 0x01 ? true : false;
11171117
}
@@ -1148,7 +1148,7 @@ private PreLoginHandshakeStatus ConsumePreLoginHandshake(
11481148

11491149
// Validate Certificate if Trust Server Certificate=false and Encryption forced (EncryptionOptions.ON) from Server.
11501150
bool shouldValidateServerCert = (_encryptionOption == EncryptionOptions.ON && !trustServerCert) ||
1151-
(_connHandler._accessTokenInBytes != null && !trustServerCert);
1151+
(_connHandler.IsAccessTokenProvided && !trustServerCert);
11521152
uint info = (shouldValidateServerCert ? TdsEnums.SNI_SSL_VALIDATE_CERTIFICATE : 0)
11531153
| TdsEnums.SNI_SSL_USE_SCHANNEL_CACHE;
11541154

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -705,7 +705,7 @@ public string AccessToken
705705

706706
_accessToken = value;
707707
// Need to call ConnectionString_Set to do proper pool group check
708-
ConnectionString_Set(new SqlConnectionPoolKey(_connectionString, _credential, _accessToken, null));
708+
ConnectionString_Set(new SqlConnectionPoolKey(_connectionString, _credential, _accessToken, _accessTokenCallback));
709709
}
710710
}
711711

@@ -727,7 +727,7 @@ public Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticati
727727
CheckAndThrowOnInvalidCombinationOfConnectionOptionAndAccessTokenCallback((SqlConnectionString)ConnectionOptions);
728728
}
729729

730-
ConnectionString_Set(new SqlConnectionPoolKey(_connectionString, _credential, null, value));
730+
ConnectionString_Set(new SqlConnectionPoolKey(_connectionString, _credential, _accessToken, value));
731731
_accessTokenCallback = value;
732732
}
733733
}

‎src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlInternalConnectionTds.cs‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,11 @@ internal sealed class SqlInternalConnectionTds : SqlInternalConnection, IDisposa
137137
internal byte[] _accessTokenInBytes;
138138
internal readonly Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticationToken>> _accessTokenCallback;
139139

140+
internal bool IsAccessTokenProvided =>
141+
_accessTokenInBytes != null || _accessTokenCallback != null;
142+
143+
internal bool? TnirDisabledDuringLogin { get; private set; }
144+
140145
private readonly ActiveDirectoryAuthenticationTimeoutRetryHelper _activeDirectoryAuthTimeoutRetryHelper;
141146

142147
internal bool _cleanSQLDNSCaching = false;
@@ -1574,7 +1579,8 @@ private void LoginNoFailover(ServerInfo serverInfo,
15741579

15751580
ResolveExtendedServerName(serverInfo, !redirectedUserInstance, connectionOptions);
15761581

1577-
bool disableTnir = ShouldDisableTnir(connectionOptions);
1582+
bool disableTnir = ShouldDisableTnir(connectionOptions, IsAccessTokenProvided);
1583+
TnirDisabledDuringLogin = disableTnir;
15781584

15791585
long timeoutUnitInterval = 0;
15801586

@@ -1783,11 +1789,11 @@ private void LoginNoFailover(ServerInfo serverInfo,
17831789
CurrentDataSource = originalServerInfo.UserServerName;
17841790
}
17851791

1786-
private bool ShouldDisableTnir(SqlConnectionString connectionOptions)
1792+
internal static bool ShouldDisableTnir(SqlConnectionString connectionOptions, bool isAccessTokenProvided)
17871793
{
17881794
Boolean isAzureEndPoint = ADP.IsAzureSqlServerEndpoint(connectionOptions.DataSource);
17891795

1790-
Boolean isFedAuthEnabled = this._accessTokenInBytes != null ||
1796+
Boolean isFedAuthEnabled = isAccessTokenProvided ||
17911797
connectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryPassword ||
17921798
connectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryIntegrated ||
17931799
connectionOptions.Authentication == SqlAuthenticationMethod.ActiveDirectoryInteractive ||

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1210,10 +1210,10 @@ private PreLoginHandshakeStatus ConsumePreLoginHandshake(
12101210

12111211
// We must NOT use the response for the FEDAUTHREQUIRED PreLogin option, if the connection string option
12121212
// was not using the new Authentication keyword or in other words, if Authentication=NotSpecified
1213-
// Or AccessToken is not null, mean token based authentication is used.
1213+
// Or an access token was supplied, which means token-based authentication is used.
12141214
if ((_connHandler.ConnectionOptions != null
12151215
&& _connHandler.ConnectionOptions.Authentication != SqlAuthenticationMethod.NotSpecified)
1216-
|| _connHandler._accessTokenInBytes != null || _connHandler._accessTokenCallback != null)
1216+
|| _connHandler.IsAccessTokenProvided)
12171217
{
12181218
fedAuthRequired = payload[payloadOffset] == 0x01 ? true : false;
12191219
}
@@ -1249,7 +1249,7 @@ private PreLoginHandshakeStatus ConsumePreLoginHandshake(
12491249
}
12501250

12511251
// Validate Certificate if Trust Server Certificate=false and Encryption forced (EncryptionOptions.ON) from Server.
1252-
bool shouldValidateServerCert = (_encryptionOption == EncryptionOptions.ON && !trustServerCert) || ((_connHandler._accessTokenInBytes != null || _connHandler._accessTokenCallback != null) && !trustServerCert);
1252+
bool shouldValidateServerCert = (_encryptionOption == EncryptionOptions.ON && !trustServerCert) || (_connHandler.IsAccessTokenProvided && !trustServerCert);
12531253

12541254
uint info = (shouldValidateServerCert ? TdsEnums.SNI_SSL_VALIDATE_CERTIFICATE : 0)
12551255
| TdsEnums.SNI_SSL_USE_SCHANNEL_CACHE;

‎src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/SqlConnectionStringTest.cs‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,39 @@ public void TestDefaultTnir(string dataSource, bool? tnirEnabledInConnString, Tr
6060
// Assert
6161
Assert.Equal(expectedValue, connectionString.TransparentNetworkIPResolution);
6262
}
63+
64+
[Theory]
65+
[InlineData("my.test.server", false, null, false)]
66+
[InlineData("my.test.server", true, null, true)]
67+
[InlineData("test.database.windows.net", false, null, true)]
68+
[InlineData("test.database.windows.net", true, null, true)]
69+
[InlineData("my.test.server", true, true, false)]
70+
[InlineData("test.database.windows.net", true, true, false)]
71+
[InlineData("test.database.windows.net", false, true, false)]
72+
[InlineData("my.test.server", true, false, false)]
73+
[InlineData("my.test.server", false, false, false)]
74+
[InlineData("test.database.windows.net", true, false, false)]
75+
[InlineData("test.database.windows.net", false, false, false)]
76+
public void TestShouldDisableTnirWithCallerSuppliedToken(
77+
string dataSource,
78+
bool isAccessTokenProvided,
79+
bool? tnirInConnectionString,
80+
bool expectedValue)
81+
{
82+
SqlConnectionStringBuilder builder = new() { DataSource = dataSource };
83+
if (tnirInConnectionString.HasValue)
84+
{
85+
builder.TransparentNetworkIPResolution = tnirInConnectionString.Value;
86+
}
87+
88+
SqlConnectionString connectionOptions = new(builder.ConnectionString);
89+
90+
Assert.Equal(
91+
expectedValue,
92+
global::Microsoft.Data.SqlClient.SqlInternalConnectionTds.ShouldDisableTnir(
93+
connectionOptions,
94+
isAccessTokenProvided));
95+
}
6396
#endif
6497

6598
/// <summary>

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

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -676,6 +676,107 @@ public void ConnectionTestAccessTokenCallbackCombinations()
676676
}
677677
}
678678

679+
private static Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticationToken>> CreateStubCallback() =>
680+
(ctx, token) => Task.FromResult(new SqlAuthenticationToken("invalid", DateTimeOffset.MaxValue));
681+
682+
[Theory]
683+
[InlineData(false)]
684+
[InlineData(true)]
685+
public async Task AccessTokenCallbackHonorsPreLoginFedAuthRequired(bool openAsync)
686+
{
687+
using TdsServer server = new(new TdsServerArguments()
688+
{
689+
FedAuthRequiredPreLoginOption = TdsPreLoginFedAuthRequiredOption.FedAuthRequired,
690+
});
691+
server.Start();
692+
693+
string connectionString = new SqlConnectionStringBuilder()
694+
{
695+
DataSource = $"localhost,{server.EndPoint.Port}",
696+
Encrypt = SqlConnectionEncryptOption.Optional,
697+
Pooling = false,
698+
}.ConnectionString;
699+
700+
using SqlConnection connection = new(connectionString)
701+
{
702+
AccessTokenCallback = CreateStubCallback(),
703+
};
704+
705+
if (openAsync)
706+
{
707+
await connection.OpenAsync();
708+
}
709+
else
710+
{
711+
connection.Open();
712+
}
713+
714+
Assert.Equal(ConnectionState.Open, connection.State);
715+
716+
#if NETFRAMEWORK
717+
Assert.True(GetTnirDisabledDuringLogin(connection));
718+
719+
using SqlConnection baseline = new(connectionString);
720+
if (openAsync)
721+
{
722+
await baseline.OpenAsync();
723+
}
724+
else
725+
{
726+
baseline.Open();
727+
}
728+
Assert.False(GetTnirDisabledDuringLogin(baseline));
729+
730+
static bool? GetTnirDisabledDuringLogin(SqlConnection connection) =>
731+
((global::Microsoft.Data.SqlClient.SqlInternalConnectionTds)connection.InnerConnection)
732+
.TnirDisabledDuringLogin;
733+
#endif
734+
}
735+
736+
[Fact]
737+
public void ClearingOneAccessTokenPropertyPreservesTheOtherInPoolKey()
738+
{
739+
Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticationToken>> callback =
740+
CreateStubCallback();
741+
742+
using (SqlConnection connection = new("Data Source=localhost"))
743+
{
744+
connection.AccessTokenCallback = callback;
745+
connection.AccessToken = null;
746+
747+
Assert.Same(callback, connection.AccessTokenCallback);
748+
Assert.Same(callback, ((global::Microsoft.Data.SqlClient.ConnectionPool.SqlConnectionPoolKey)connection.PoolGroup.PoolKey).AccessTokenCallback);
749+
}
750+
751+
using (SqlConnection connection = new("Data Source=localhost"))
752+
{
753+
connection.AccessToken = "token";
754+
connection.AccessTokenCallback = null;
755+
756+
Assert.Equal("token", connection.AccessToken);
757+
Assert.Equal("token", ((global::Microsoft.Data.SqlClient.ConnectionPool.SqlConnectionPoolKey)connection.PoolGroup.PoolKey).AccessToken);
758+
}
759+
}
760+
761+
[Fact]
762+
public void AccessTokenAndAccessTokenCallbackAreMutuallyExclusive()
763+
{
764+
Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticationToken>> callback =
765+
CreateStubCallback();
766+
767+
using (SqlConnection connection = new("Data Source=localhost"))
768+
{
769+
connection.AccessTokenCallback = callback;
770+
Assert.Throws<InvalidOperationException>(() => connection.AccessToken = "token");
771+
}
772+
773+
using (SqlConnection connection = new("Data Source=localhost"))
774+
{
775+
connection.AccessToken = "token";
776+
Assert.Throws<InvalidOperationException>(() => connection.AccessTokenCallback = callback);
777+
}
778+
}
779+
679780
[Theory]
680781
[InlineData(9, 0, 2047)] // SQL Server 2005
681782
[InlineData(10, 0, 2531)] // SQL Server 2008

0 commit comments

Comments
 (0)