From 05473b8ea97e6c70b1b1575cb59228585189f569 Mon Sep 17 00:00:00 2001 From: Lawrence Cheung <31262254+lcheunglci@users.noreply.github.com> Date: Fri, 10 Mar 2023 11:10:21 -0500 Subject: [PATCH] Backport 3.1.x Fix | TDS RPC error on large queries in SqlCommand.ExecuteReaderAsync (#1939) --- .../src/Microsoft/Data/SqlClient/TdsParser.cs | 6 +- .../ManualTests/AlwaysEncrypted/ApiShould.cs | 138 ++++++++++++++++++ 2 files changed, 141 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 31a44e3221..58cfc9af79 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 @@ -9860,10 +9860,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 136bd36b42..fadf2ed254 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; @@ -753,6 +754,55 @@ public void TestExecuteReader(string connection) }); } + [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringSetupForAE))] + [ClassData(typeof(AEConnectionStringProvider))] + public async void TestExecuteReaderAsyncWithLargeQuery(string connection) + { + string tableName = "VeryLong_01234567890123456789012345678901234567890123456789_TestTableName"; + int columnsCount = 50; + + // Arrange - drops the table with long name and re-creates it with 52 columns (ID, name, ColumnName0..49) + try + { + DropTableIfExists(connection, tableName); + CreateTable(connection, tableName, columnsCount); + string name = "nobody"; + + using (SqlConnection sqlConnection = new SqlConnection(connection)) + { + await sqlConnection.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"), + sqlConnection, + 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 (Exception ex) + { + Assert.False(true, $"The following exception was thrown: {ex.Message}"); + } + finally + { + DropTableIfExists(connection, tableName); + } + } + [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringSetupForAE))] [ClassData(typeof(AEConnectionStringProviderWithCommandBehaviorSet1))] public void TestExecuteReaderWithCommandBehavior(string connection, CommandBehavior commandBehavior) @@ -2820,6 +2870,94 @@ private void CleanUpTable(string connString, string tableName) } } + + /// + /// Creates a table with the specified number of bit columns. + /// + /// The connection string to the database + /// The table name + /// The number of bit columns + private void CreateTable(string connString, string tableName, int columnsCount) + { + using (var sqlConnection = new SqlConnection(connString)) + { + sqlConnection.Open(); + + SqlCommand cmd = new SqlCommand(GenerateCreateQuery(tableName, columnsCount), sqlConnection); + cmd.ExecuteNonQuery(); + } + } + + /// + /// Drops the table if the specified table exists + /// + /// The connection string to the database + /// The name of the table to be dropped + private void DropTableIfExists(string connString, string tableName) + { + using (var sqlConnection = new SqlConnection(connString)) + { + sqlConnection.Open(); + SqlCommand cmd = new SqlCommand($"DROP TABLE IF EXISTS {tableName};", sqlConnection); + cmd.ExecuteNonQuery(); + } + } + + /// + /// 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 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 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. ///