diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParser.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParser.cs index f9b1306e42..1074047644 100644 --- a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParser.cs +++ b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParser.cs @@ -7223,16 +7223,34 @@ internal TdsOperationStatus TryReadSqlValue(SqlBuffer value, { if (stateObj is not null) { - // call to decrypt column keys has failed. The data wont be decrypted. - // Not setting the value to false, forces the driver to look for column value. - // Packet received from Key Vault will throws invalid token header. - if (stateObj.HasPendingData) + // Throwing an exception here circumvents the normal pending data checks and cleanup processes, + // so we need to ensure the appropriate state. Increment the _nextColumnDataToRead index because + // we already read the encrypted column data; Otherwise we'll double count and attempt to drain a + // corresponding number of bytes a second time. We don't want the rest of the pending data to + // interfere with future operations, so we must drain it. Set HasPendingData to false to indicate + // that we successfully drained the data. + + // The SqlDataReader also maintains a state called dataReady. We need to set that to false if we've + // drained the data off the connection. Otherwise, a consumer that catches the exception may + // continue to use the reader and will timeout waiting to read data that doesn't exist. + + // Order matters here. Must increment column before draining data. + // Update state objects after draining data. + + if (stateObj._readerState != null) { - // Drain the pending data now if setting the HasPendingData to false. - // SqlDataReader.TryCloseInternal can not drain if HasPendingData = false. - DrainData(stateObj); + stateObj._readerState._nextColumnDataToRead++; } + + DrainData(stateObj); + + if (stateObj._readerState != null) + { + stateObj._readerState._dataReady = false; + } + stateObj.HasPendingData = false; + } throw SQL.ColumnDecryptionFailed(columnName, null, e); } diff --git a/src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ColumnDecryptErrorTests.cs b/src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ColumnDecryptErrorTests.cs index 954d94ee7f..d7435b95af 100644 --- a/src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ColumnDecryptErrorTests.cs +++ b/src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ColumnDecryptErrorTests.cs @@ -6,7 +6,6 @@ using System.Collections; using System.Collections.Generic; using Microsoft.Data.SqlClient.ManualTesting.Tests.AlwaysEncrypted.Setup; -using Microsoft.Data.SqlClient.ManualTesting.Tests.SystemDataInternals; using Xunit; namespace Microsoft.Data.SqlClient.ManualTesting.Tests.AlwaysEncrypted diff --git a/src/Microsoft.Data.SqlClient/tests/ManualTests/Microsoft.Data.SqlClient.ManualTesting.Tests.csproj b/src/Microsoft.Data.SqlClient/tests/ManualTests/Microsoft.Data.SqlClient.ManualTesting.Tests.csproj index 6de3141afa..c6cd8c17df 100644 --- a/src/Microsoft.Data.SqlClient/tests/ManualTests/Microsoft.Data.SqlClient.ManualTesting.Tests.csproj +++ b/src/Microsoft.Data.SqlClient/tests/ManualTests/Microsoft.Data.SqlClient.ManualTesting.Tests.csproj @@ -47,6 +47,7 @@ + @@ -74,7 +75,6 @@ -