Test | Unit Test for decrypt failure to drain data fix

src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/AKVTests.cs

@@ -35,7 +35,7 @@ public void TestEncryptDecryptWithAKV()
                 AttestationProtocol = SqlConnectionAttestationProtocol.NotSpecified,
                 EnclaveAttestationUrl = ""
             };
-            using SqlConnection sqlConnection = new (builder.ConnectionString);
+            using SqlConnection sqlConnection = new(builder.ConnectionString);
 
             sqlConnection.Open();
             Customer customer = new(45, "Microsoft", "Corporation");
@@ -48,7 +48,7 @@ public void TestEncryptDecryptWithAKV()
             }
 
             // Test INPUT parameter on an encrypted parameter
-            using SqlCommand sqlCommand = new ($"SELECT CustomerId, FirstName, LastName FROM [{_akvTableName}] WHERE FirstName = @firstName",
+            using SqlCommand sqlCommand = new($"SELECT CustomerId, FirstName, LastName FROM [{_akvTableName}] WHERE FirstName = @firstName",
                                                             sqlConnection);
             SqlParameter customerFirstParam = sqlCommand.Parameters.AddWithValue(@"firstName", @"Microsoft");
             customerFirstParam.Direction = System.Data.ParameterDirection.Input;
@@ -58,11 +58,82 @@ public void TestEncryptDecryptWithAKV()
             DatabaseHelper.ValidateResultSet(sqlDataReader);
         }
 
+        /*
+         This unit test is going to assess an issue where a failed decryption leaves a connection in a bad state  
+         when it is returned to the connection pool. If a subsequent connection is retried it will result in an "Internal connection fatal error",
+         which causes that connection to be doomed, preventing it from being returned to the pool. 
+         Consequently, retrying a third connection will encounter the same decryption error, leading to a repetitive failure cycle.
+
+         The purpose of this unit test is to simulate a decryption error and verify that the connection remains usable when returned to the pool. 
+         It aims to confirm that three consecutive connections will consistently fail with the "Failed to decrypt column" error.
+        */
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringSetupForAE), nameof(DataTestUtility.IsAKVSetupAvailable))]
+        public void ForcedColumnDecryptErrorTestShouldFail()
+        {
+            SqlConnectionStringBuilder builder = new(DataTestUtility.TCPConnectionStringHGSVBS)
+            {
+                ColumnEncryptionSetting = SqlConnectionColumnEncryptionSetting.Enabled,
+                AttestationProtocol = SqlConnectionAttestationProtocol.NotSpecified,
+                EnclaveAttestationUrl = ""
+            };
+
+            // Setup record to query
+            using (SqlConnection sqlConnection = new(builder.ConnectionString))
+            {
+                sqlConnection.Open();
+                Customer customer = new(88, "Microsoft2", "Corporation2");
+
+                using (SqlTransaction sqlTransaction = sqlConnection.BeginTransaction())
+                {
+                    DatabaseHelper.InsertCustomerData(sqlConnection, sqlTransaction, _akvTableName, customer);
+                    sqlTransaction.Commit();
+                }
+            }
+
+            // Setup Empty key store provider
+            Dictionary<String, SqlColumnEncryptionKeyStoreProvider> emptyKeyStoreProviders = new()
+            {
+                { "AZURE_KEY_VAULT", new EmptyKeyStoreProvider() }
+            };
+
+            // Three consecutive connections should fail with "Failed to decrypt column" error. This proves that an error in decryption
+            // does not leave the connection in a bad state.
+            // In each try, when a "Failed to decrypt error" is thrown, the connection's TDS Parser state object buffer is drained of any
+            // pending data so it does not interfere with future operations. In addition, the TDS parser state object's reader.DataReady flag
+            // is set to false so that the calling function that catches the exception will not continue to use the reader. Otherwise, it will 
+            // timeout waiting to read data that doesn't exist. Also, the TDS Parser state object HasPendingData flag is also set to false
+            // to indicate that the buffer has been cleared and to avoid it getting cleared again in SqlDataReader.TryCloseInternal function.
+            // Finally, after successfully handling the decryption error, the connection is then returned back to the connection pool without 
+            // an error. A proof that the connection's state object is clean is in the second connection being able to throw the same error.
+            // The third connection is for making sure we test 3 times as the minimum number of connections to reproduce the issue previously.
+            for (int i = 0; i < 3; i++)
+            {
+                using (SqlConnection sqlConnection = new SqlConnection(builder.ConnectionString))
+                {
+                    sqlConnection.Open();
+                    // Setup connection using the empty key store provider thereby forcing a decryption error.
+                    sqlConnection.RegisterColumnEncryptionKeyStoreProvidersOnConnection(emptyKeyStoreProviders);
+
+                    using SqlCommand sqlCommand = new($"SELECT FirstName FROM [{_akvTableName}] WHERE FirstName = @firstName", sqlConnection);
+                    SqlParameter customerFirstParam = sqlCommand.Parameters.AddWithValue(@"firstName", @"Microsoft2");
+                    customerFirstParam.Direction = System.Data.ParameterDirection.Input;
+                    customerFirstParam.ForceColumnEncryption = true;
+
+                    using SqlDataReader sqlDataReader = sqlCommand.ExecuteReader();
+                    while (sqlDataReader.Read())
+                    {
+                        var error = Assert.Throws<SqlException>(() => DatabaseHelper.CompareResults(sqlDataReader, new string[] { @"string" }, 1));
+                        Assert.Contains("Failed to decrypt column", error.Message);
+                    }
+                }
+            }
+        }
+
         [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.IsAKVSetupAvailable))]
         [PlatformSpecific(TestPlatforms.Windows)]
         public void TestRoundTripWithAKVAndCertStoreProvider()
         {
-            using SQLSetupStrategyCertStoreProvider certStoreFixture = new ();
+            using SQLSetupStrategyCertStoreProvider certStoreFixture = new();
             byte[] plainTextColumnEncryptionKey = ColumnEncryptionKey.GenerateRandomBytes(ColumnEncryptionKey.KeySizeInBytes);
             byte[] encryptedColumnEncryptionKeyUsingAKV = _fixture.AkvStoreProvider.EncryptColumnEncryptionKey(DataTestUtility.AKVUrl, @"RSA_OAEP", plainTextColumnEncryptionKey);
             byte[] columnEncryptionKeyReturnedAKV2Cert = certStoreFixture.CertStoreProvider.DecryptColumnEncryptionKey(certStoreFixture.CspColumnMasterKey.KeyPath, @"RSA_OAEP", encryptedColumnEncryptionKeyUsingAKV);
@@ -120,5 +191,18 @@ public void TestLocalCekCacheIsScopedToProvider()
             Exception ex = Assert.Throws<SqlException>(() => sqlCommand.ExecuteReader());
             Assert.StartsWith("The current credential is not configured to acquire tokens for tenant", ex.InnerException.Message);
         }
+
+        private class EmptyKeyStoreProvider : SqlColumnEncryptionKeyStoreProvider
+        {
+            public override byte[] DecryptColumnEncryptionKey(string masterKeyPath, string encryptionAlgorithm, byte[] encryptedColumnEncryptionKey)
+            {
+                return new byte[32];
+            }
+
+            public override byte[] EncryptColumnEncryptionKey(string masterKeyPath, string encryptionAlgorithm, byte[] columnEncryptionKey)
+            {
+                return new byte[32];
+            }
+        }
     }
 }