From 02f14d27a431948dbab47f5dc06e970a99e0e8d3 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 | Fix decrypt failure to drain data (#2618)

---
 .../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 7e04f6b9c5..3c73cb89a9 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
@@ -6139,6 +6139,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 7a9bbfdfd3..49b6eacb8e 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
@@ -6931,6 +6931,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.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);