EnclaveDelegate.Crypto GetEnclaveProvider appears to not be thread safe

[dna495]

As a result of issue #1422 we are restarting app service every 8 hours to prevent the null enclave session. Occasionally we are getting the following exception when multiple async calls are made at once.

System.InvalidOperationException: Operations that change non-concurrent collections must have exclusive access. A concurrent update was performed on this collection and corrupted its state. The collection's state is no longer correct.
   at System.Collections.Generic.Dictionary`2.FindEntry(TKey key)
   at System.Collections.Generic.Dictionary`2.TryGetValue(TKey key, TValue& value)
   at Microsoft.Data.SqlClient.EnclaveDelegate.GetEnclaveProvider(SqlConnectionAttestationProtocol attestationProtocol, String enclaveType)
   at Microsoft.Data.SqlClient.SqlCommand.TryFetchInputParameterEncryptionInfo(Int32 timeout, Boolean isAsync, Boolean asyncWrite, Boolean& inputParameterEncryptionNeeded, Task& task, ReadOnlyDictionary`2& describeParameterEncryptionRpcOriginalRpcMap)
   at Microsoft.Data.SqlClient.SqlCommand.PrepareForTransparentEncryption(CommandBehavior cmdBehavior, Boolean returnStream, Boolean isAsync, Int32 timeout, TaskCompletionSource`1 completion, Task& returnTask, Boolean asyncWrite, Boolean& usedCache, Boolean inRetry)
   at Microsoft.Data.SqlClient.SqlCommand.RunExecuteReader(CommandBehavior cmdBehavior, RunBehavior runBehavior, Boolean returnStream, TaskCompletionSource`1 completion, Int32 timeout, Task& task, Boolean& usedCache, Boolean asyncWrite, Boolean inRetry, String method)
   at Microsoft.Data.SqlClient.SqlCommand.BeginExecuteReaderInternal(CommandBehavior behavior, AsyncCallback callback, Object stateObject, Int32 timeout, Boolean inRetry, Boolean asyncWrite)
   at Microsoft.Data.SqlClient.SqlCommand.BeginExecuteReaderAsyncCallback(AsyncCallback callback, Object stateObject)
   at System.Threading.Tasks.TaskFactory`1.FromAsyncImpl(Func`3 beginMethod, Func`2 endFunction, Action`1 endAction, Object state, TaskCreationOptions creationOptions)
   at Microsoft.Data.SqlClient.SqlCommand.InternalExecuteReaderAsync(CommandBehavior behavior, CancellationToken cancellationToken)

To reproduce (it doesn't occur every time, but if you run program a few times it will occur)

using Azure.Identity;
using Microsoft.Data.SqlClient;
using Microsoft.Data.SqlClient.AlwaysEncrypted.AzureKeyVaultProvider;
using System;
using System.Collections.Generic;
using System.Data;
using System.Diagnostics;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;

namespace EncryptionTest
{
    class Program
    {
        static readonly string server = "";
        static readonly string user = "";
        static readonly string password = "";
        static readonly string attestationUrl = "";
        static readonly string catalog = "";
        static readonly string s_connectionString = $"Server={server};Initial Catalog={catalog};Connection Timeout=30;UID={user};PWD={password}; Column Encryption Setting = Enabled;Attestation Protocol = AAS; Enclave Attestation Url = {attestationUrl};";

        static Action<int> sleepAction = (int d) => Thread.Sleep(d);
        private static object CurrentUserDbContext(int sleep)
        {
            sleepAction(sleep);
            return "";
        }

        public static async Task RunQueryAsync(int sleep)
        {
            Debug.Print($"RunQuery: {DateTime.Now}");

            using (SqlConnection sqlConnection = new SqlConnection(s_connectionString))
            {
                if (sqlConnection.State != ConnectionState.Open)
                {
                    sqlConnection.Open();
                }
                using (SqlCommand sqlCommand = new SqlCommand("SP_SET_SESSION_CONTEXT", sqlConnection))
                {
                    sqlCommand.CommandType = CommandType.StoredProcedure;
                    sqlCommand.Parameters.AddWithValue("@@key", "UserContext");
                    sqlCommand.Parameters.AddWithValue("@@value", JsonSerializer.Serialize(CurrentUserDbContext(sleep)));
                    //Thread.Sleep(sleep);
                    int rowAffected = await sqlCommand.ExecuteNonQueryAsync();
                    //Thread.Sleep(sleep);
                    Debug.Print($"Param: {rowAffected} - Time: {DateTime.Now}");
                }

            }
        }

        static void Main(string[] args)
        {
            SqlColumnEncryptionAzureKeyVaultProvider akvProvider = new SqlColumnEncryptionAzureKeyVaultProvider(new DefaultAzureCredential());
            SqlConnection.RegisterColumnEncryptionKeyStoreProviders(customProviders: new Dictionary<string, SqlColumnEncryptionKeyStoreProvider>(capacity: 1, comparer: StringComparer.OrdinalIgnoreCase)
            {
                    { SqlColumnEncryptionAzureKeyVaultProvider.ProviderName, akvProvider}
            });

            int sleepTime = 500;

            Parallel.Invoke(
                async () => await RunQueryAsync(sleepTime),
                async () => await RunQueryAsync(sleepTime),
                async () => await RunQueryAsync(sleepTime),
                async () => await RunQueryAsync(sleepTime),
                async () => await RunQueryAsync(sleepTime),
                async () => await RunQueryAsync(sleepTime),
                async () => await RunQueryAsync(sleepTime),
                async () => await RunQueryAsync(sleepTime));

        }
    }
}

Expected behavior

s_enclaveProviders should be thread safe and not fail in GetEnclaveProvider

Further technical details

Microsoft.Data.SqlClient version: 4.0.0
.NET target: Core 3.1
SQL Server version: Azure SQL
Attestation: AAS

[dna495]

thoughts?
dna495@106e80b

[Wraith2]

thoughts? dna495@106e80b

Are there any other places that the dictionary is mutated? if so those will need to be locked as well. Otherwise looks ok to me.

[johnnypham]

Changing it to a ConcurrentDictionary seems to fix it. Let me know what you think.