Cleanup | Deduplicate enclave attestation session caching logic

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

@@ -17,8 +17,8 @@
       <param name="enclaveAttestationInfo">
         The information the provider uses to attest the enclave and generate a symmetric key for the session. The format of this information is specific to the enclave attestation protocol.
       </param>
-      <param name="clientDiffieHellmanKey">
-        A Diffie-Hellman algorithm object that encapsulates a client-side key pair.
+      <param name="attestationParameters">
+        A set of parameters for the attestation process, including the client's Diffie-Hellman key pair.
       </param>
       <param name="enclaveSessionParameters">
         The set of parameters required for an enclave session.

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/AzureAttestationBasedEnclaveProvider.cs

@@ -1,4 +1,4 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
@@ -79,56 +79,35 @@ internal override SqlEnclaveAttestationParameters GetAttestationParameters(strin
             return new SqlEnclaveAttestationParameters(AzureBasedAttestationProtocolId, attestationParam, clientDHKey);
         }
 
-        // When overridden in a derived class, performs enclave attestation, generates a symmetric key for the session, creates a an enclave session and stores the session information in the cache.
-        internal override void CreateEnclaveSession(byte[] attestationInfo, ECDiffieHellman clientDHKey, EnclaveSessionParameters enclaveSessionParameters, byte[] customData, int customDataLength, out SqlEnclaveSession sqlEnclaveSession, out long counter)
+        protected override SqlEnclaveSession CreateEnclaveSessionCore(
+            byte[] enclaveAttestationInfo,
+            SqlEnclaveAttestationParameters attestationParameters,
+            EnclaveSessionParameters enclaveSessionParameters,
+            byte[] customData,
+            int customDataLength)
         {
-            sqlEnclaveSession = null;
-            counter = 0;
-            try
+            if (!string.IsNullOrEmpty(enclaveSessionParameters.AttestationUrl) && customData != null && customDataLength > 0)
             {
-                ThreadRetryCache.Remove(Thread.CurrentThread.ManagedThreadId.ToString());
-                sqlEnclaveSession = GetEnclaveSessionFromCache(enclaveSessionParameters, out counter);
-                if (sqlEnclaveSession == null)
-                {
-                    if (!string.IsNullOrEmpty(enclaveSessionParameters.AttestationUrl) && customData != null && customDataLength > 0)
-                    {
-                        byte[] nonce = customData;
+                byte[] nonce = customData;
 
-                        IdentityModelEventSource.ShowPII = true;
+                IdentityModelEventSource.ShowPII = true;
 
-                        // Deserialize the payload
-                        AzureAttestationInfo attestInfo = new AzureAttestationInfo(attestationInfo);
+                // Deserialize the payload
+                AzureAttestationInfo attestInfo = new AzureAttestationInfo(enclaveAttestationInfo);
 
-                        // Validate the attestation info
-                        VerifyAzureAttestationInfo(enclaveSessionParameters.AttestationUrl, attestInfo.EnclaveType, attestInfo.AttestationToken.AttestationToken, attestInfo.Identity, nonce);
+                // Validate the attestation info
+                VerifyAzureAttestationInfo(enclaveSessionParameters.AttestationUrl, attestInfo.EnclaveType, attestInfo.AttestationToken.AttestationToken, attestInfo.Identity, nonce);
 
-                        // Set up shared secret and validate signature
-                        byte[] sharedSecret = GetSharedSecret(attestInfo.Identity, nonce, attestInfo.EnclaveType, attestInfo.EnclaveDHInfo, clientDHKey);
+                // Set up shared secret and validate signature
+                byte[] sharedSecret = GetSharedSecret(attestInfo.Identity, nonce, attestInfo.EnclaveType, attestInfo.EnclaveDHInfo, attestationParameters.ClientDiffieHellmanKey);
 
-                        // add session to cache
-                        sqlEnclaveSession = AddEnclaveSessionToCache(enclaveSessionParameters, sharedSecret, attestInfo.SessionId, out counter);
-                    }
-                    else
-                    {
-                        throw SQL.AttestationFailed(Strings.FailToCreateEnclaveSession);
-                    }
-                }
+                return new SqlEnclaveSession(sharedSecret, attestInfo.SessionId);
             }
-            finally
+            else
             {
-                // As per current design, we want to minimize the number of create session calls. To achieve this we block all the GetEnclaveSession calls until the first call to
-                // GetEnclaveSession -> GetAttestationParameters -> CreateEnclaveSession completes or the event timeout happen.
-                // Case 1: When the first request successfully creates the session, then all outstanding GetEnclaveSession will use the current session.
-                // Case 2: When the first request unable to create the enclave session (may be due to some error or the first request doesn't require enclave computation) then in those case we set the event timeout to 0.
-                UpdateEnclaveSessionLockStatus(sqlEnclaveSession);
+                throw SQL.AttestationFailed(Strings.FailToCreateEnclaveSession);
             }
         }
-
-        // When overridden in a derived class, looks up and evicts an enclave session from the enclave session cache, if the provider implements session caching.
-        internal override void InvalidateEnclaveSession(EnclaveSessionParameters enclaveSessionParameters, SqlEnclaveSession enclaveSessionToInvalidate)
-        {
-            InvalidateEnclaveSessionHelper(enclaveSessionParameters, enclaveSessionToInvalidate);
-        }
         #endregion
 
         #region Internal Class

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveDelegate.Crypto.cs

@@ -1,4 +1,4 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
@@ -47,7 +47,7 @@ internal void CreateEnclaveSession(SqlConnectionAttestationProtocol attestationP
 
                 sqlColumnEncryptionEnclaveProvider.CreateEnclaveSession(
                     attestationInfo,
-                    attestationParameters.ClientDiffieHellmanKey,
+                    attestationParameters,
                     enclaveSessionParameters,
                     customData,
                     customDataLength,

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveProviderBase.cs

@@ -1,4 +1,4 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
@@ -88,6 +88,9 @@ internal abstract class EnclaveProviderBase : SqlColumnEncryptionEnclaveProvider
         #endregion
 
         #region protected methods
+        // When overridden in a derived class, performs enclave attestation, generates a symmetric key for the session and creates an enclave session.
+        protected abstract SqlEnclaveSession CreateEnclaveSessionCore(byte[] enclaveAttestationInfo, SqlEnclaveAttestationParameters attestationParameters, EnclaveSessionParameters enclaveSessionParameters, byte[] customData, int customDataLength);
+
         // Helper method to get the enclave session from the cache if present
         protected void GetEnclaveSessionHelper(EnclaveSessionParameters enclaveSessionParameters, bool shouldGenerateNonce, bool isRetry, out SqlEnclaveSession sqlEnclaveSession, out long counter, out byte[] customData, out int customDataLength)
         {
@@ -173,8 +176,39 @@ protected void GetEnclaveSessionHelper(EnclaveSessionParameters enclaveSessionPa
             }
         }
 
+        // Calls CreateEnclaveSessionCore to create an enclave session and stores the session information in the cache.
+        internal override void CreateEnclaveSession(byte[] attestationInfo, SqlEnclaveAttestationParameters attestationParameters, EnclaveSessionParameters enclaveSessionParameters, byte[] customData, int customDataLength, out SqlEnclaveSession sqlEnclaveSession, out long counter)
+        {
+            sqlEnclaveSession = null;
+            counter = 0;
+            try
+            {
+                ThreadRetryCache.Remove(Thread.CurrentThread.ManagedThreadId.ToString());
+                sqlEnclaveSession = SessionCache.GetEnclaveSession(enclaveSessionParameters, out counter);
+                if (sqlEnclaveSession == null)
+                {
+                    // Add session to cache
+                    sqlEnclaveSession = CreateEnclaveSessionCore(attestationInfo, attestationParameters, enclaveSessionParameters, customData, customDataLength);
+                    SessionCache.CreateSession(enclaveSessionParameters, sqlEnclaveSession, out counter);
+                }
+            }
+            finally
+            {
+                // As per current design, we want to minimize the number of create session calls. To achieve this we block all the GetEnclaveSession calls until the first call to
+                // GetEnclaveSession -> GetAttestationParameters -> CreateEnclaveSession completes or the event timeout happen.
+                // Case 1: When the first request successfully creates the session, then all outstanding GetEnclaveSession will use the current session.
+                // Case 2: When the first request unable to create the enclave session (may be due to some error or the first request doesn't require enclave computation) then in those case we set the event timeout to 0.
+                UpdateEnclaveSessionLockStatus(sqlEnclaveSession);
+            }
+        }
+
+        internal override void InvalidateEnclaveSession(EnclaveSessionParameters enclaveSessionParameters, SqlEnclaveSession enclaveSessionToInvalidate)
+        {
+            SessionCache.InvalidateSession(enclaveSessionParameters, enclaveSessionToInvalidate);
+        }
+
         // Reset the session lock status
-        protected void UpdateEnclaveSessionLockStatus(SqlEnclaveSession sqlEnclaveSession)
+        private void UpdateEnclaveSessionLockStatus(SqlEnclaveSession sqlEnclaveSession)
         {
             // As per current design, we want to minimize the number of create session calls. To achieve this we block all the GetEnclaveSession calls until the first call to
             // GetEnclaveSession -> GetAttestationParameters -> CreateEnclaveSession completes or the event timeout happens.
@@ -193,24 +227,6 @@ protected void UpdateEnclaveSessionLockStatus(SqlEnclaveSession sqlEnclaveSessio
                 }
             }
         }
-
-        // Helper method to remove the enclave session from the cache
-        protected void InvalidateEnclaveSessionHelper(EnclaveSessionParameters enclaveSessionParameters, SqlEnclaveSession enclaveSessionToInvalidate)
-        {
-            SessionCache.InvalidateSession(enclaveSessionParameters, enclaveSessionToInvalidate);
-        }
-
-        // Helper method for getting the enclave session from the session cache
-        protected SqlEnclaveSession GetEnclaveSessionFromCache(EnclaveSessionParameters enclaveSessionParameters, out long counter)
-        {
-            return SessionCache.GetEnclaveSession(enclaveSessionParameters, out counter);
-        }
-
-        // Helper method for adding the enclave session to the session cache
-        protected SqlEnclaveSession AddEnclaveSessionToCache(EnclaveSessionParameters enclaveSessionParameters, byte[] sharedSecret, long sessionId, out long counter)
-        {
-            return SessionCache.CreateSession(enclaveSessionParameters, sharedSecret, sessionId, out counter);
-        }
     }
     #endregion
 }

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

@@ -1,4 +1,4 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
@@ -54,20 +54,16 @@ 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)
+        // Adds a new SqlEnclaveSession to the cache
+        internal void CreateSession(EnclaveSessionParameters enclaveSessionParameters, SqlEnclaveSession enclaveSession, out long counter)
         {
             string cacheKey = GenerateCacheKey(enclaveSessionParameters);
-            SqlEnclaveSession enclaveSession = null;
             lock (enclaveCacheLock)
             {
-                enclaveSession = new SqlEnclaveSession(sharedSecret, sessionId);
                 enclaveMemoryCache.Set<SqlEnclaveSession>(cacheKey, enclaveSession,
                     absoluteExpirationRelativeToNow: s_enclaveCacheTimeout);
                 counter = Interlocked.Increment(ref _counter);
             }
-
-            return enclaveSession;
         }
 
         // Generates the cache key for the enclave session cache