From 6cec1ebe0a124d12253d1823884ca370e9c47025 Mon Sep 17 00:00:00 2001
From: Lawrence Cheung <31262254+lcheunglci@users.noreply.github.com>
Date: Mon, 20 Mar 2023 23:47:59 -0400
Subject: [PATCH] Fix | TDS RPC error on large queries in
 SqlCommand.ExecuteReaderAsync (#1936)

---
 .../src/Microsoft/Data/SqlClient/TdsParser.cs |   6 +-
 .../ManualTests/AlwaysEncrypted/ApiShould.cs  | 123 ++++++++++++++++++
 2 files changed, 126 insertions(+), 3 deletions(-)

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 df5e789e8c..776b4f0a9b 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
@@ -9907,10 +9907,10 @@ internal Task TdsExecuteRPC(SqlCommand cmd, _SqlRPC[] rpcArray, int timeout, boo
 
                             // Options
                             WriteShort((short)rpcext.options, stateObj);
-                        }
 
-                        byte[] enclavePackage = cmd.enclavePackage != null ? cmd.enclavePackage.EnclavePackageBytes : null;
-                        WriteEnclaveInfo(stateObj, enclavePackage);
+                            byte[] enclavePackage = cmd.enclavePackage != null ? cmd.enclavePackage.EnclavePackageBytes : null;
+                            WriteEnclaveInfo(stateObj, enclavePackage);
+                        }
 
                         // Stream out parameters
                         SqlParameter[] parameters = rpcext.parameters;
diff --git a/src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ApiShould.cs b/src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ApiShould.cs
index 2bec916210..855cf590b0 100644
--- a/src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ApiShould.cs
+++ b/src/Microsoft.Data.SqlClient/tests/ManualTests/AlwaysEncrypted/ApiShould.cs
@@ -8,6 +8,7 @@
 using System.IO;
 using System.Linq;
 using System.Reflection;
+using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider;
@@ -689,6 +690,59 @@ public void TestExecuteReader(string connection)
             });
         }
 
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringSetupForAE))]
+        [ClassData(typeof(AEConnectionStringProvider))]
+        public async void TestExecuteReaderAsyncWithLargeQuery(string connectionString)
+        {
+            string randomName = DataTestUtility.GetUniqueName(Guid.NewGuid().ToString().Replace("-", ""), false);
+            if (randomName.Length > 50)
+            {
+                randomName = randomName.Substring(0, 50);
+            }
+            string tableName = $"VeryLong_{randomName}_TestTableName";
+            int columnsCount = 50;
+
+            // Arrange - drops the table with long name and re-creates it with 52 columns (ID, name, ColumnName0..49)
+            try
+            {
+                CreateTable(connectionString, tableName, columnsCount);
+                string name = "nobody";
+
+                using (SqlConnection connection = new SqlConnection(connectionString))
+                {
+                    await connection.OpenAsync();
+                    // This creates a "select top 100" query that has over 40k characters
+                    using (SqlCommand sqlCommand = new SqlCommand(GenerateSelectQuery(tableName, columnsCount, 10, "WHERE Name = @FirstName AND ID = @CustomerId"),
+                        connection,
+                        transaction: null,
+                        columnEncryptionSetting: SqlCommandColumnEncryptionSetting.Enabled))
+                    {
+                        sqlCommand.Parameters.Add(@"CustomerId", SqlDbType.Int);
+                        sqlCommand.Parameters.Add(@"FirstName", SqlDbType.VarChar, name.Length);
+
+                        sqlCommand.Parameters[0].Value = 0;
+                        sqlCommand.Parameters[1].Value = name;
+
+                        // Act and Assert
+                        // Test that execute reader async does not throw an exception.
+                        // The table is empty so there should be no results; however, the bug previously found is that it causes a TDS RPC exception on enclave.
+                        using (SqlDataReader sqlDataReader = await sqlCommand.ExecuteReaderAsync())
+                        {
+                            Assert.False(sqlDataReader.HasRows, "The table should be empty");
+                        }
+                    }
+                }
+            }
+            catch
+            {
+                throw;
+            }
+            finally
+            {
+                DropTableIfExists(connectionString, tableName);
+            }
+        }
+
         [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringSetupForAE))]
         [ClassData(typeof(AEConnectionStringProviderWithCommandBehaviorSet1))]
         public void TestExecuteReaderWithCommandBehavior(string connection, CommandBehavior commandBehavior)
@@ -2788,6 +2842,75 @@ private void CleanUpTable(string connString, string tableName)
             }
         }
 
+        private static void CreateTable(string connString, string tableName, int columnsCount)
+            => DataTestUtility.RunNonQuery(connString, GenerateCreateQuery(tableName, columnsCount));
+        /// <summary>
+        /// Drops the table if the specified table exists
+        /// </summary>
+        /// <param name="connString">The connection string to the database</param>
+        /// <param name="tableName">The name of the table to be dropped</param>
+        private static void DropTableIfExists(string connString, string tableName)
+        {
+            using var sqlConnection = new SqlConnection(connString);
+            sqlConnection.Open();
+            DataTestUtility.DropTable(sqlConnection, tableName);
+        }
+
+        /// <summary>
+        /// Generates the query for creating a table with the number of bit columns specified.
+        /// </summary>
+        /// <param name="tableName">The name of the table</param>
+        /// <param name="columnsCount">The number of columns for the table</param>
+        /// <returns></returns>
+        private static string GenerateCreateQuery(string tableName, int columnsCount)
+        {
+            StringBuilder builder = new StringBuilder();
+            builder.Append(string.Format("CREATE TABLE [dbo].[{0}]", tableName));
+            builder.Append('(');
+            builder.AppendLine("[ID][bigint] NOT NULL,");
+            builder.AppendLine("[Name] [varchar] (200) NOT NULL");
+            for (int i = 0; i < columnsCount; i++)
+            {
+                builder.Append(',');
+                builder.Append($"[ColumnName{i}][bit] NULL");
+            }
+            builder.Append(");");
+
+            return builder.ToString();
+        }
+
+        /// <summary>
+        /// Generates the large query with the select top 100 of all the columns repeated multiple times.
+        /// </summary>
+        /// <param name="tableName">The name of the table</param>
+        /// <param name="columnsCount">The number of columns to be explicitly included</param>
+        /// <param name="repeat">The number of times the select query is repeated</param>
+        /// <param name="where">A where clause for additional filters</param>
+        /// <returns></returns>
+        private static string GenerateSelectQuery(string tableName, int columnsCount, int repeat = 10, string where = "")
+        {
+            StringBuilder builder = new StringBuilder();
+            builder.AppendLine($"SELECT TOP 100");
+            builder.AppendLine($"[{tableName}].[ID],");
+            builder.AppendLine($"[{tableName}].[Name]");
+            for (int i = 0; i < columnsCount; i++)
+            {
+                builder.Append(",");
+                builder.AppendLine($"[{tableName}].[ColumnName{i}]");
+            }
+
+            string extra = string.IsNullOrEmpty(where) ? $"(NOLOCK) [{tableName}]" : where;
+            builder.AppendLine($"FROM [{tableName}] {extra};");
+
+            StringBuilder builder2 = new StringBuilder();
+            for (int i = 0; i < repeat; i++)
+            {
+                builder2.AppendLine(builder.ToString());
+            }
+
+            return builder2.ToString();
+        }
+
         /// <summary>
         /// An helper method to test the cancellation of the command using cancellationToken to async SqlCommand APIs.
         /// </summary>