From d8d21ae1a64638c265390106fbccbd1921044a2b Mon Sep 17 00:00:00 2001 From: Lawrence Cheung <31262254+lcheunglci@users.noreply.github.com> Date: Mon, 10 Apr 2023 20:40:00 -0400 Subject: [PATCH] Fix | TDS RPC error on large queries in SqlCommand.ExecuteReaderAsync (#1936) (#1986) --- .../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 2fd15a0def..920c81983b 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 @@ -9868,10 +9868,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 5096d15af4..588781a84a 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.ManualTesting.Tests.AlwaysEncrypted.Setup; @@ -728,6 +729,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) @@ -2552,6 +2606,75 @@ private void CleanUpTable(string connString, string tableName) } } + private static void CreateTable(string connString, string tableName, int columnsCount) + => DataTestUtility.RunNonQuery(connString, GenerateCreateQuery(tableName, columnsCount)); + /// + /// Drops the table if the specified table exists + /// + /// The connection string to the database + /// The name of the table to be dropped + private static void DropTableIfExists(string connString, string tableName) + { + using var sqlConnection = new SqlConnection(connString); + sqlConnection.Open(); + DataTestUtility.DropTable(sqlConnection, tableName); + } + + /// + /// Generates the query for creating a table with the number of bit columns specified. + /// + /// The name of the table + /// The number of columns for the table + /// + 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(); + } + + /// + /// Generates the large query with the select top 100 of all the columns repeated multiple times. + /// + /// The name of the table + /// The number of columns to be explicitly included + /// The number of times the select query is repeated + /// A where clause for additional filters + /// + 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(); + } + /// /// An helper method to test the cancellation of the command using cancellationToken to async SqlCommand APIs. ///