From 39484926cf1b660c186d4e4e13c9276267c82f23 Mon Sep 17 00:00:00 2001
From: Aris Rellegue <134557572+arellegue@users.noreply.github.com>
Date: Thu, 11 Jul 2024 10:53:48 -0700
Subject: [PATCH] Fix decrypt failure to drain data

---
 .../netcore/src/Microsoft/Data/SqlClient/TdsParser.cs       | 6 ++++++
 .../netfx/src/Microsoft/Data/SqlClient/TdsParser.cs         | 6 ++++++
 2 files changed, 12 insertions(+)

diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParser.cs b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParser.cs
index 17327f16e1..dae197f5f7 100644
--- a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParser.cs
+++ b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParser.cs
@@ -6090,6 +6090,12 @@ internal bool TryReadSqlValue(SqlBuffer value, SqlMetaDataPriv md, int length, T
                                 // 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)
+                                {
+                                    // Drain the pending data now if setting the HasPendingData to false.
+                                    // SqlDataReader.TryCloseInternal can not drain if HasPendingData = false.
+                                    DrainData(stateObj);
+                                }
                                 stateObj.HasPendingData = false;
                             }
                             throw SQL.ColumnDecryptionFailed(columnName, null, e);
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 e143b13dfc..c94f714973 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
@@ -6910,6 +6910,12 @@ internal bool TryReadSqlValue(SqlBuffer value,
                                 // 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._pendingData)
+                                {
+                                    // Drain the pending data now if setting the HasPendingData to false.
+                                    // SqlDataReader.TryCloseInternal can not drain if HasPendingData = false.
+                                    DrainData(stateObj);
+                                }
                                 stateObj._pendingData = false;
                             }
                             throw SQL.ColumnDecryptionFailed(columnName, null, e);