[7.0.4] Isolate enclave session cache entries by attestation protocol

doc/snippets/Microsoft.Data.SqlClient/SqlConnectionStringBuilder.xml

@@ -261,6 +261,9 @@ The following example supplies a simple SQL Server connection string in the <xre
       </value>
       <remarks>
         When no value is specified, secure enclaves are disabled on the connection.
+        Enclave sessions are cached separately for each attestation protocol. A session
+        established with <see cref="F:Microsoft.Data.SqlClient.SqlConnectionAttestationProtocol.None" />
+        cannot satisfy a connection configured to require HGS or AAS attestation.
       </remarks>
     </AttestationProtocol>
     <Authentication>

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveSessionCache.cs

@@ -24,7 +24,8 @@ internal class EnclaveSessionCache
         // Retrieves a SqlEnclaveSession from the cache
         internal SqlEnclaveSession GetEnclaveSession(EnclaveSessionParameters enclaveSessionParameters, out long counter)
         {
-            string cacheKey = GenerateCacheKey(enclaveSessionParameters);
+            Tuple<SqlConnectionAttestationProtocol, string, string, string> cacheKey =
+                GenerateCacheKey(enclaveSessionParameters);
             SqlEnclaveSession enclaveSession = enclaveMemoryCache.Get<SqlEnclaveSession>(cacheKey);
             counter = Interlocked.Increment(ref _counter);
             return enclaveSession;
@@ -33,7 +34,8 @@ internal SqlEnclaveSession GetEnclaveSession(EnclaveSessionParameters enclaveSes
         // Invalidates a SqlEnclaveSession entry in the cache
         internal void InvalidateSession(EnclaveSessionParameters enclaveSessionParameters, SqlEnclaveSession enclaveSessionToInvalidate)
         {
-            string cacheKey = GenerateCacheKey(enclaveSessionParameters);
+            Tuple<SqlConnectionAttestationProtocol, string, string, string> cacheKey =
+                GenerateCacheKey(enclaveSessionParameters);
 
             lock (enclaveCacheLock)
             {
@@ -57,7 +59,8 @@ internal void InvalidateSession(EnclaveSessionParameters enclaveSessionParameter
         // Creates a new SqlEnclaveSession and adds it to the cache
         internal SqlEnclaveSession CreateSession(EnclaveSessionParameters enclaveSessionParameters, byte[] sharedSecret, long sessionId, out long counter)
         {
-            string cacheKey = GenerateCacheKey(enclaveSessionParameters);
+            Tuple<SqlConnectionAttestationProtocol, string, string, string> cacheKey =
+                GenerateCacheKey(enclaveSessionParameters);
             SqlEnclaveSession enclaveSession = null;
             lock (enclaveCacheLock)
             {
@@ -70,10 +73,16 @@ internal SqlEnclaveSession CreateSession(EnclaveSessionParameters enclaveSession
             return enclaveSession;
         }
 
-        // Generates the cache key for the enclave session cache
-        private string GenerateCacheKey(EnclaveSessionParameters enclaveSessionParameters)
+        // Keep the attestation policy and field boundaries in the identity while preserving
+        // the existing case-insensitive, null/empty-equivalent string matching.
+        private static Tuple<SqlConnectionAttestationProtocol, string, string, string> GenerateCacheKey(
+            EnclaveSessionParameters enclaveSessionParameters)
         {
-            return (enclaveSessionParameters.ServerName + '+' + enclaveSessionParameters.Database + enclaveSessionParameters.AttestationUrl).ToLowerInvariant();
+            return Tuple.Create(
+                enclaveSessionParameters.AttestationProtocol,
+                (enclaveSessionParameters.ServerName ?? string.Empty).ToLowerInvariant(),
+                (enclaveSessionParameters.Database ?? string.Empty).ToLowerInvariant(),
+                (enclaveSessionParameters.AttestationUrl ?? string.Empty).ToLowerInvariant());
         }
     }
 }

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs

@@ -239,7 +239,8 @@ private EnclaveSessionParameters GetEnclaveSessionParameters() =>
             new EnclaveSessionParameters(
                 _activeConnection.DataSource,
                 _activeConnection.EnclaveAttestationUrl,
-                _activeConnection.Database);
+                _activeConnection.Database,
+                _activeConnection.AttestationProtocol);
 
         // @TODO: Isn't this doing things asynchronously? We should just have a purely asynchronous and a purely synchronous pathway instead of this mix of check this check that and flags.
         private SqlDataReader GetParameterEncryptionDataReader(

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlEnclaveSession.cs

@@ -62,12 +62,15 @@ internal class EnclaveSessionParameters
         internal string ServerName { get; set; }  // The name of the SQL Server instance containing the enclave.
         internal string AttestationUrl { get; set; }  // The endpoint of an attestation service for attesting the enclave.
         internal string Database { get; set; }  //  The database that SqlClient contacts to request an enclave session.
+        internal SqlConnectionAttestationProtocol AttestationProtocol { get; }
 
-        internal EnclaveSessionParameters(string serverName, string attestationUrl, string database)
+        internal EnclaveSessionParameters(string serverName, string attestationUrl, string database,
+            SqlConnectionAttestationProtocol attestationProtocol)
         {
             ServerName = serverName;
             AttestationUrl = attestationUrl;
             Database = database;
+            AttestationProtocol = attestationProtocol;
         }
     }
 }

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

@@ -2483,6 +2483,140 @@ public void TestRetryWhenAEParameterMetadataCacheIsStale(string connectionString
             cmd.ExecuteNonQuery();
         }
 
+        /// <summary>
+        /// Enclave queries reuse a same-policy session across connections and establish a new
+        /// session after invalidation, for both synchronous and asynchronous execution.
+        /// Run with each configured HGS, AAS, and None enclave connection.
+        /// </summary>
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.IsTargetReadyForAeWithKeyStore), nameof(DataTestUtility.EnclaveEnabled))]
+        [ClassData(typeof(AEConnectionStringProviderWithBooleanVariable))]
+        public async Task TestEnclaveSessionReuseAcrossConnections(string connectionString, bool isAsync)
+        {
+            CleanUpTable(connectionString, _tableName);
+            const int customerId = 50;
+            InsertRows(_tableName, 1, GetValues(dataHint: customerId), connectionString);
+
+            object originalSession = null;
+            for (int attempt = 0; attempt < 3; attempt++)
+            {
+                using SqlConnection connection = new(connectionString);
+                if (isAsync)
+                {
+                    await connection.OpenAsync();
+                }
+                else
+                {
+                    connection.Open();
+                }
+
+                using SqlCommand command = new(
+                    $"SELECT CustomerId FROM [{_tableName}] WHERE CustomerId > @CustomerId",
+                    connection, null, SqlCommandColumnEncryptionSetting.Enabled);
+                command.Parameters.Add("@CustomerId", SqlDbType.Int).Value = 0;
+                using SqlDataReader reader = isAsync
+                    ? await command.ExecuteReaderAsync()
+                    : command.ExecuteReader();
+                object session = CommandHelper.GetEnclaveSession(command);
+
+                Assert.True(isAsync ? await reader.ReadAsync() : reader.Read());
+                Assert.Equal(customerId, reader.GetInt32(0));
+                Assert.False(isAsync ? await reader.ReadAsync() : reader.Read());
+
+                if (attempt == 0)
+                {
+                    originalSession = session;
+                }
+                else if (attempt == 1)
+                {
+                    Assert.Same(originalSession, session);
+                    CommandHelper.InvalidateEnclaveSession(command);
+                }
+                else
+                {
+                    Assert.NotSame(originalSession, session);
+                }
+            }
+        }
+
+        /// <summary>
+        /// An enclave session established under one attestation protocol must not be reused by a
+        /// connection that requests a different attestation protocol, even when the server, database
+        /// and attestation URL are identical. The HGS connection must perform its own attestation.
+        /// </summary>
+        [ConditionalTheory(typeof(ApiShould), nameof(IsHgsVbsEnclaveSetup))]
+        [InlineData(false, false)]
+        [InlineData(false, true)]
+        [InlineData(true, false)]
+        [InlineData(true, true)]
+        public async Task TestEnclaveSessionIsNotSharedAcrossAttestationProtocols(bool noneFirst, bool isAsync)
+        {
+            string hgsConnectionString = DataTestUtility.TCPConnectionStringHGSVBS;
+
+            // Derive the None connection from the HGS one so the attestation protocol is the only
+            // difference; server, database and attestation URL stay identical.
+            SqlConnectionStringBuilder noneBuilder = new(hgsConnectionString)
+            {
+                AttestationProtocol = SqlConnectionAttestationProtocol.None
+            };
+            string noneConnectionString = noneBuilder.ConnectionString;
+
+            const int customerId = 51;
+            CleanUpTable(hgsConnectionString, _tableName);
+            InsertRows(_tableName, 1, GetValues(dataHint: customerId), hgsConnectionString);
+
+            // Start from an empty cache for both protocols.
+            await ExecuteEnclaveQueryAsync(noneConnectionString, customerId, isAsync, invalidateSession: true);
+            await ExecuteEnclaveQueryAsync(hgsConnectionString, customerId, isAsync, invalidateSession: true);
+            await ExecuteEnclaveQueryAsync(noneConnectionString, customerId, isAsync, invalidateSession: true);
+
+            string firstConnectionString = noneFirst ? noneConnectionString : hgsConnectionString;
+            string secondConnectionString = noneFirst ? hgsConnectionString : noneConnectionString;
+
+            object firstSession = await ExecuteEnclaveQueryAsync(firstConnectionString, customerId, isAsync);
+            object secondSession = await ExecuteEnclaveQueryAsync(secondConnectionString, customerId, isAsync);
+
+            Assert.NotSame(firstSession, secondSession);
+
+            // Each protocol continues to reuse its own session.
+            Assert.Same(firstSession, await ExecuteEnclaveQueryAsync(firstConnectionString, customerId, isAsync));
+            Assert.Same(secondSession, await ExecuteEnclaveQueryAsync(secondConnectionString, customerId, isAsync));
+        }
+
+        public static bool IsHgsVbsEnclaveSetup() =>
+            DataTestUtility.IsTargetReadyForAeWithKeyStore() &&
+            DataTestUtility.EnclaveEnabled &&
+            !string.IsNullOrEmpty(DataTestUtility.TCPConnectionStringHGSVBS);
+
+        private async Task<object> ExecuteEnclaveQueryAsync(string connectionString, int expectedCustomerId, bool isAsync, bool invalidateSession = false)
+        {
+            using SqlConnection connection = new(connectionString);
+            if (isAsync)
+            {
+                await connection.OpenAsync();
+            }
+            else
+            {
+                connection.Open();
+            }
+
+            using SqlCommand command = new(
+                $"SELECT CustomerId FROM [{_tableName}] WHERE CustomerId > @CustomerId",
+                connection, null, SqlCommandColumnEncryptionSetting.Enabled);
+            command.Parameters.Add("@CustomerId", SqlDbType.Int).Value = 0;
+            using SqlDataReader reader = isAsync ? await command.ExecuteReaderAsync() : command.ExecuteReader();
+            object session = CommandHelper.GetEnclaveSession(command);
+
+            Assert.True(isAsync ? await reader.ReadAsync() : reader.Read());
+            Assert.Equal(expectedCustomerId, reader.GetInt32(0));
+            Assert.False(isAsync ? await reader.ReadAsync() : reader.Read());
+
+            if (invalidateSession)
+            {
+                CommandHelper.InvalidateEnclaveSession(command);
+            }
+            return session;
+        }
+
         [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringSetupForAE), nameof(DataTestUtility.EnclaveEnabled))]
         [ClassData(typeof(AEConnectionStringProvider))]
         public void TestRetryWhenAEEnclaveCacheIsStale(string connectionString)

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/Common/SystemDataInternals/CommandHelper.cs

@@ -13,6 +13,7 @@ internal static class CommandHelper
         private static MethodInfo s_completePendingReadWithSuccess = s_sqlCommand.GetMethod("CompletePendingReadWithSuccess", BindingFlags.NonPublic | BindingFlags.Instance);
         private static MethodInfo s_completePendingReadWithFailure = s_sqlCommand.GetMethod("CompletePendingReadWithFailure", BindingFlags.NonPublic | BindingFlags.Instance);
         private static MethodInfo s_invalidateEnclaveSession = s_sqlCommand.GetMethod("InvalidateEnclaveSession", BindingFlags.NonPublic | BindingFlags.Instance);
+        private static readonly FieldInfo s_enclavePackage = s_sqlCommand.GetField("enclavePackage", BindingFlags.NonPublic | BindingFlags.Instance);
 #if DEBUG
         private static FieldInfo s_forceRetryableEnclaveQueryExecutionExceptionDuringGenerateEnclavePackage =
             s_sqlCommand.GetField(@"_forceRetryableEnclaveQueryExecutionExceptionDuringGenerateEnclavePackage", BindingFlags.NonPublic | BindingFlags.Static);
@@ -41,6 +42,23 @@ internal static void InvalidateEnclaveSession(SqlCommand command)
             s_invalidateEnclaveSession.Invoke(command, null);
         }
 
+        /// <summary>
+        /// Gets the session used for a command's enclave package without exposing its key material.
+        /// Call while the command still retains its encryption state.
+        /// </summary>
+        /// <param name="command">An executed command requiring enclave computations.</param>
+        /// <returns>The session object, for reference-identity assertions.</returns>
+        internal static object GetEnclaveSession(SqlCommand command)
+        {
+            object package = s_enclavePackage.GetValue(command)
+                ?? throw new InvalidOperationException("The command did not generate an enclave package.");
+            PropertyInfo sessionProperty = package.GetType().GetProperty(
+                "EnclaveSession", BindingFlags.NonPublic | BindingFlags.Instance)
+                ?? throw new MissingMemberException(package.GetType().FullName, "EnclaveSession");
+            return sessionProperty.GetValue(package)
+                ?? throw new InvalidOperationException("The enclave package has no session.");
+        }
+
 #if DEBUG
         internal static void ForceThrowDuringGenerateEnclavePackage(SqlCommand command)
         {