Fix | Fix Async issue with Azure Attestation Service

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

@@ -21,21 +21,29 @@ the enclave attestation protocol as well as the logic for creating and caching e
             <param name="clientDiffieHellmanKey">A Diffie-Hellman algorithm object that encapsulates a client-side key pair.</param>
             <param name="attestationUrl">The endpoint of an attestation service for attesting the enclave.</param>
             <param name="servername">The name of the SQL Server instance containing the enclave.</param>
+            <param name="customData">The set of extra data needed for attestating the enclave.</param>
+            <param name="customDataLength">The length of the extra data needed for attestating the enclave.</param>
             <param name="sqlEnclaveSession">The requested enclave session or <see langword="null" /> if the provider doesn't implement session caching.</param>
             <param name="counter">A counter that the enclave provider is expected to increment each time SqlClient retrieves the session from the cache. The purpose of this field is to prevent replay attacks.</param>
             <summary>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.</summary>
             <remarks>To be added.</remarks>
         </CreateEnclaveSession>
         <GetAttestationParameters>
+            <param name="attestationUrl">The endpoint of an attestation service for attesting the enclave.</param>
+            <param name="customData">A set of extra data needed for attestating the enclave.</param>
+            <param name="customDataLength">The length of the extra data needed for attestating the enclave.</param>
             <summary>Gets the information that SqlClient subsequently uses to initiate the process of attesting the enclave and to establish a secure session with the enclave.</summary>
             <returns>The information SqlClient subsequently uses to initiate the process of attesting the enclave and to establish a secure session with the enclave.</returns>
             <remarks>To be added.</remarks>
         </GetAttestationParameters>
         <GetEnclaveSession>
             <param name="serverName">The name of the SQL Server instance containing the enclave.</param>
             <param name="attestationUrl">The endpoint of an attestation service, SqlClient contacts to attest the enclave.</param>
+            <param name="generateCustomData"><see langword="true" /> to indicate that a set of extra data needs to be generated for attestation; otherwise, <see langword="false" />.</param>
             <param name="sqlEnclaveSession">When this method returns, the requested enclave session or <see langword="null" /> if the provider doesn't implement session caching. This parameter is treated as uninitialized.</param>
             <param name="counter">A counter that the enclave provider is expected to increment each time SqlClient retrieves the session from the cache. The purpose of this field is to prevent replay attacks.</param>
+            <param name="customData">A set of extra data needed for attestating the enclave.</param>
+            <param name="customDataLength">The length of the extra data needed for attestating the enclave.</param>
             <summary>When overridden in a derived class, looks up an existing enclave session information in the enclave session cache. If the enclave provider doesn't implement enclave session caching, this method is expected to return <see langword="null" /> in the <paramref name="sqlEnclaveSession" /> parameter.
             </summary>
             <remarks>To be added.</remarks>

src/Microsoft.Data.SqlClient/netcore/ref/Microsoft.Data.SqlClient.NetCoreApp.cs

@@ -185,11 +185,11 @@ public abstract partial class SqlColumnEncryptionEnclaveProvider
         /// <include file='../../../../doc/snippets/Microsoft.Data.SqlClient/SqlColumnEncryptionEnclaveProvider.xml' path='docs/members[@name="SqlColumnEncryptionEnclaveProvider"]/ctor/*'/>
         protected SqlColumnEncryptionEnclaveProvider() { }
         /// <include file='../../../../doc/snippets/Microsoft.Data.SqlClient/SqlColumnEncryptionEnclaveProvider.xml' path='docs/members[@name="SqlColumnEncryptionEnclaveProvider"]/CreateEnclaveSession/*'/>
-        public abstract void CreateEnclaveSession(byte[] enclaveAttestationInfo, System.Security.Cryptography.ECDiffieHellmanCng clientDiffieHellmanKey, string attestationUrl, string servername, out Microsoft.Data.SqlClient.SqlEnclaveSession sqlEnclaveSession, out long counter);
+        public abstract void CreateEnclaveSession(byte[] enclaveAttestationInfo, System.Security.Cryptography.ECDiffieHellmanCng clientDiffieHellmanKey, string attestationUrl, string servername, byte[] customData, int customDataLength, out Microsoft.Data.SqlClient.SqlEnclaveSession sqlEnclaveSession, out long counter);
         /// <include file='../../../../doc/snippets/Microsoft.Data.SqlClient/SqlColumnEncryptionEnclaveProvider.xml' path='docs/members[@name="SqlColumnEncryptionEnclaveProvider"]/GetAttestationParameters/*'/>
-        public abstract Microsoft.Data.SqlClient.SqlEnclaveAttestationParameters GetAttestationParameters();
+        public abstract Microsoft.Data.SqlClient.SqlEnclaveAttestationParameters GetAttestationParameters(string attestationUrl, byte[] customData, int customDataLength);
         /// <include file='../../../../doc/snippets/Microsoft.Data.SqlClient/SqlColumnEncryptionEnclaveProvider.xml' path='docs/members[@name="SqlColumnEncryptionEnclaveProvider"]/GetEnclaveSession/*'/>
-        public abstract void GetEnclaveSession(string serverName, string attestationUrl, out Microsoft.Data.SqlClient.SqlEnclaveSession sqlEnclaveSession, out long counter);
+        public abstract void GetEnclaveSession(string serverName, string attestationUrl, bool generateCustomData, out Microsoft.Data.SqlClient.SqlEnclaveSession sqlEnclaveSession, out long counter, out byte[] customData, out int customDataLength);
         /// <include file='../../../../doc/snippets/Microsoft.Data.SqlClient/SqlColumnEncryptionEnclaveProvider.xml' path='docs/members[@name="SqlColumnEncryptionEnclaveProvider"]/InvalidateEnclaveSession/*'/>
         public abstract void InvalidateEnclaveSession(string serverName, string enclaveAttestationUrl, Microsoft.Data.SqlClient.SqlEnclaveSession enclaveSession);
     }

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

@@ -66,35 +66,35 @@ internal class AzureAttestationEnclaveProvider : EnclaveProviderBase
         #region Public methods
         // When overridden in a derived class, looks up an existing enclave session information in the enclave session cache.
         // If the enclave provider doesn't implement enclave session caching, this method is expected to return null in the sqlEnclaveSession parameter.
-        public override void GetEnclaveSession(string servername, string attestationUrl, out SqlEnclaveSession sqlEnclaveSession, out long counter)
+        public override void GetEnclaveSession(string servername, string attestationUrl, bool generateCustomData, out SqlEnclaveSession sqlEnclaveSession, out long counter, out byte[] customData, out int customDataLength)
         {
-            GetEnclaveSessionHelper(servername, attestationUrl, true, out sqlEnclaveSession, out counter);
+            GetEnclaveSessionHelper(servername, attestationUrl, generateCustomData, out sqlEnclaveSession, out counter, out customData, out customDataLength);
         }
 
         // Gets the information that SqlClient subsequently uses to initiate the process of attesting the enclave and to establish a secure session with the enclave.
-        public override SqlEnclaveAttestationParameters GetAttestationParameters()
+        public override SqlEnclaveAttestationParameters GetAttestationParameters(string attestationUrl, byte[] customData, int customDataLength)
         {
             ECDiffieHellmanCng clientDHKey = new ECDiffieHellmanCng(DiffieHellmanKeySize);
             clientDHKey.KeyDerivationFunction = ECDiffieHellmanKeyDerivationFunction.Hash;
             clientDHKey.HashAlgorithm = CngAlgorithm.Sha256;
-            byte[] attestationParam = PrepareAttestationParameters();
+            byte[] attestationParam = PrepareAttestationParameters(attestationUrl, customData, customDataLength);
             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.
-        public override void CreateEnclaveSession(byte[] attestationInfo, ECDiffieHellmanCng clientDHKey, string attestationUrl, string servername, out SqlEnclaveSession sqlEnclaveSession, out long counter)
+        public override void CreateEnclaveSession(byte[] attestationInfo, ECDiffieHellmanCng clientDHKey, string attestationUrl, string servername, byte[] customData, int customDataLength, out SqlEnclaveSession sqlEnclaveSession, out long counter)
         {
             sqlEnclaveSession = null;
             counter = 0;
             try
             {
-                AttestationInfoCacheItem attestationInfoCacheItem = AttestationInfoCache.Remove(Thread.CurrentThread.ManagedThreadId.ToString()) as AttestationInfoCacheItem;
+                ThreadRetryCache.Remove(Thread.CurrentThread.ManagedThreadId.ToString());
                 sqlEnclaveSession = GetEnclaveSessionFromCache(servername, attestationUrl, out counter);
                 if (sqlEnclaveSession == null)
                 {
-                    if (attestationInfoCacheItem != null)
+                    if (!string.IsNullOrEmpty(attestationUrl) && customData != null && customDataLength > 0)
                     {
-                        byte[] nonce = attestationInfoCacheItem.AttestNonce;
+                        byte[] nonce = customData;
 
                         IdentityModelEventSource.ShowPII = true;
 
@@ -241,19 +241,18 @@ public AzureAttestationToken(byte[] payload)
         // Attestation Url
         // Size of nonce
         // Nonce value
-        internal byte[] PrepareAttestationParameters()
+        internal byte[] PrepareAttestationParameters(string attestationUrl, byte[] attestNonce, int attestNonceLength)
         {
-            AttestationInfoCacheItem attestationInfoCacheItem = AttestationInfoCache[Thread.CurrentThread.ManagedThreadId.ToString()] as AttestationInfoCacheItem;
-            if (attestationInfoCacheItem != null)
+            if (!string.IsNullOrEmpty(attestationUrl) && attestNonce != null && attestNonceLength > 0)
             {
                 // In c# strings are not null terminated, so adding the null termination before serializing it
-                string attestationUrlLocal = attestationInfoCacheItem.AttestationUrl + char.MinValue;
+                string attestationUrlLocal = attestationUrl + char.MinValue;
                 byte[] serializedAttestationUrl = Encoding.Unicode.GetBytes(attestationUrlLocal);
                 byte[] serializedAttestationUrlLength = BitConverter.GetBytes(serializedAttestationUrl.Length);
 
                 // serializing nonce
-                byte[] serializedNonce = attestationInfoCacheItem.AttestNonce;
-                byte[] serializedNonceLength = BitConverter.GetBytes(attestationInfoCacheItem.AttestNonce.Length);
+                byte[] serializedNonce = attestNonce;
+                byte[] serializedNonceLength = BitConverter.GetBytes(attestNonceLength);
 
                 // Computing the total length of the data
                 int totalDataSize = serializedAttestationUrl.Length + serializedAttestationUrlLength.Length + serializedNonce.Length + serializedNonceLength.Length;

src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/EnclaveDelegate.NetCoreApp.cs

@@ -66,25 +66,33 @@ internal byte[] GetSerializedAttestationParameters(SqlEnclaveAttestationParamete
         /// <param name="attestationUrl">attestation url for attestation service endpoint</param>
         /// <param name="attestationInfo">attestation info from SQL Server</param>
         /// <param name="attestationParameters">attestation parameters</param>
+        /// <param name="customData">A set of extra data needed for attestating the enclave.</param>
+        /// <param name="customDataLength">The length of the extra data needed for attestating the enclave.</param>
         internal void CreateEnclaveSession(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType, string serverName, string attestationUrl,
-            byte[] attestationInfo, SqlEnclaveAttestationParameters attestationParameters)
+            byte[] attestationInfo, SqlEnclaveAttestationParameters attestationParameters, byte[] customData, int customDataLength)
         {
 
             lock (_lock)
             {
                 SqlColumnEncryptionEnclaveProvider sqlColumnEncryptionEnclaveProvider = GetEnclaveProvider(attestationProtocol, enclaveType);
                 long counter;
                 SqlEnclaveSession sqlEnclaveSession = null;
-                sqlColumnEncryptionEnclaveProvider.GetEnclaveSession(serverName, attestationUrl, out sqlEnclaveSession, out counter);
+                byte[] dummyCustomData = null;
+                int dummyCustomDataLength;
+
+                sqlColumnEncryptionEnclaveProvider.GetEnclaveSession(serverName, attestationUrl, false, out sqlEnclaveSession, out counter, out dummyCustomData, out dummyCustomDataLength);
 
                 if (sqlEnclaveSession != null)
                 {
                     return;
                 }
 
-                sqlColumnEncryptionEnclaveProvider.CreateEnclaveSession(attestationInfo, attestationParameters.ClientDiffieHellmanKey, attestationUrl, serverName, out sqlEnclaveSession, out counter);
+                sqlColumnEncryptionEnclaveProvider.CreateEnclaveSession(attestationInfo, attestationParameters.ClientDiffieHellmanKey, attestationUrl, serverName, customData, customDataLength, out sqlEnclaveSession, out counter);
 
-                if (sqlEnclaveSession == null) throw SQL.NullEnclaveSessionReturnedFromProvider(enclaveType, attestationUrl);
+                if (sqlEnclaveSession == null) 
+                {
+                    throw SQL.NullEnclaveSessionReturnedFromProvider(enclaveType, attestationUrl);
+                }
             }
         }
 
@@ -103,9 +111,12 @@ internal EnclavePackage GenerateEnclavePackage(SqlConnectionAttestationProtocol
 
             SqlEnclaveSession sqlEnclaveSession = null;
             long counter;
+            byte[] dummyCustomData = null;
+            int dummyCustomDataLength;
+
             try
             {
-                GetEnclaveSession(attestationProtocol, enclaveType, serverName, enclaveAttestationUrl, out sqlEnclaveSession, out counter, throwIfNull: true);
+                GetEnclaveSession(attestationProtocol, enclaveType, serverName, enclaveAttestationUrl, false, out sqlEnclaveSession, out counter, out dummyCustomData, out dummyCustomDataLength, throwIfNull: true);
             }
             catch (Exception e)
             {
@@ -129,10 +140,10 @@ internal void InvalidateEnclaveSession(SqlConnectionAttestationProtocol attestat
         }
 
 
-        internal SqlEnclaveAttestationParameters GetAttestationParameters(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType)
+        internal SqlEnclaveAttestationParameters GetAttestationParameters(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType, string attestationUrl, byte[] customData, int customDataLength)
         {
             SqlColumnEncryptionEnclaveProvider sqlColumnEncryptionEnclaveProvider = GetEnclaveProvider(attestationProtocol, enclaveType);
-            return sqlColumnEncryptionEnclaveProvider.GetAttestationParameters();
+            return sqlColumnEncryptionEnclaveProvider.GetAttestationParameters(attestationUrl, customData, customDataLength);
         }
 
         private SqlColumnEncryptionEnclaveProvider GetEnclaveProvider(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType)
@@ -183,16 +194,16 @@ private string ConvertAttestationProtocolToString(SqlConnectionAttestationProtoc
             }
         }
 
-        internal void GetEnclaveSession(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType, string serverName, string enclaveAttestationUrl, out SqlEnclaveSession sqlEnclaveSession)
+        internal void GetEnclaveSession(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType, string serverName, string enclaveAttestationUrl, bool generateCustomData, out SqlEnclaveSession sqlEnclaveSession, out byte[] customData, out int customDataLength)
         {
             long counter;
-            GetEnclaveSession(attestationProtocol, enclaveType, serverName, enclaveAttestationUrl, out sqlEnclaveSession, out counter, throwIfNull: false);
+            GetEnclaveSession(attestationProtocol, enclaveType, serverName, enclaveAttestationUrl, generateCustomData, out sqlEnclaveSession, out counter, out customData, out customDataLength, throwIfNull: false);
         }
 
-        private void GetEnclaveSession(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType, string serverName, string enclaveAttestationUrl, out SqlEnclaveSession sqlEnclaveSession, out long counter, bool throwIfNull)
+        private void GetEnclaveSession(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType, string serverName, string enclaveAttestationUrl, bool generateCustomData, out SqlEnclaveSession sqlEnclaveSession, out long counter, out byte[] customData, out int customDataLength, bool throwIfNull)
         {
             SqlColumnEncryptionEnclaveProvider sqlColumnEncryptionEnclaveProvider = GetEnclaveProvider(attestationProtocol, enclaveType);
-            sqlColumnEncryptionEnclaveProvider.GetEnclaveSession(serverName, enclaveAttestationUrl, out sqlEnclaveSession, out counter);
+            sqlColumnEncryptionEnclaveProvider.GetEnclaveSession(serverName, enclaveAttestationUrl, generateCustomData, out sqlEnclaveSession, out counter, out customData, out customDataLength);
 
             if (throwIfNull && sqlEnclaveSession == null)
             {

src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/EnclaveDelegate.NetStandard.cs

@@ -27,13 +27,15 @@ internal byte[] GetSerializedAttestationParameters(
         /// <param name="attestationUrl">attestation url for attestation service endpoint</param>
         /// <param name="attestationInfo">attestation info from SQL Server</param>
         /// <param name="attestationParameters">attestation parameters</param>
+        /// <param name="customData">A set of extra data needed for attestating the enclave.</param>
+        /// <param name="customDataLength">The length of the extra data needed for attestating the enclave.</param>
         internal void CreateEnclaveSession(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType, string serverName, string attestationUrl,
-            byte[] attestationInfo, SqlEnclaveAttestationParameters attestationParameters)
+            byte[] attestationInfo, SqlEnclaveAttestationParameters attestationParameters, byte[] customData, int customDataLength)
         {
             throw new PlatformNotSupportedException();
         }
 
-        internal void GetEnclaveSession(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType, string serverName, string enclaveAttestationUrl, out SqlEnclaveSession sqlEnclaveSession)
+        internal void GetEnclaveSession(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType, string serverName, string enclaveAttestationUrl, bool generateCustomData, out SqlEnclaveSession sqlEnclaveSession, out byte[] customData, out int customDataLength)
         {
             throw new PlatformNotSupportedException();
         }
@@ -48,7 +50,7 @@ internal void InvalidateEnclaveSession(SqlConnectionAttestationProtocol attestat
             throw new PlatformNotSupportedException();
         }
 
-        internal SqlEnclaveAttestationParameters GetAttestationParameters(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType)
+        internal SqlEnclaveAttestationParameters GetAttestationParameters(SqlConnectionAttestationProtocol attestationProtocol, string enclaveType, string attestationUrl, byte[] customData, int customDataLength)
         {
             throw new PlatformNotSupportedException();
         }