feat(provisioning-device, prov-amqp, prov-mqtt, prov-https): Add supp…

…ort for timespan timeouts to provisioning device client (#2041)

As discussed in #2036, AMQP provisioning device clients have no way to configure the timeout given to the AMQP library for operations like opening links. This adds overloads to the existing provisioning device client's registerAsync methods that allow for users to configure this timespan timeout.

e2e/test/provisioning/ProvisioningE2ETests.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Net;
+using System.Net.Sockets;
 using System.Security.Cryptography;
 using System.Security.Cryptography.X509Certificates;
 using System.Text;
@@ -432,6 +433,62 @@ public async Task ProvisioningDeviceClient_ValidRegistrationId_MqttWsWithProxy_S
             await ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(Client.TransportType.Mqtt, AttestationMechanismType.SymmetricKey, EnrollmentType.Individual, true, s_proxyServerAddress).ConfigureAwait(false);
         }
 
+        [LoggedTestMethod]
+        public async Task ProvisioningDeviceClient_ValidRegistrationId_TimeSpanTimeoutRespected_Mqtt()
+        {
+            try
+            {
+                await ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(Client.TransportType.Mqtt_Tcp_Only, AttestationMechanismType.SymmetricKey, EnrollmentType.Individual, TimeSpan.Zero).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                return; // expected exception was thrown, so exit the test
+            }
+
+            throw new AssertFailedException("Expected an OperationCanceledException to be thrown since the timeout was set to TimeSpan.Zero");
+        }
+
+        [LoggedTestMethod]
+        public async Task ProvisioningDeviceClient_ValidRegistrationId_TimeSpanTimeoutRespected_Https()
+        {
+            try
+            {
+                await ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(Client.TransportType.Http1, AttestationMechanismType.SymmetricKey, EnrollmentType.Individual, TimeSpan.Zero).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                return; // expected exception was thrown, so exit the test
+            }
+
+            throw new AssertFailedException("Expected an OperationCanceledException to be thrown since the timeout was set to TimeSpan.Zero");
+        }
+
+        [LoggedTestMethod]
+        public async Task ProvisioningDeviceClient_ValidRegistrationId_TimeSpanTimeoutRespected_Amqps()
+        {
+            try
+            {
+                await ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(Client.TransportType.Amqp_Tcp_Only, AttestationMechanismType.SymmetricKey, EnrollmentType.Individual, TimeSpan.Zero).ConfigureAwait(false);
+            }
+            catch (ProvisioningTransportException ex) when (ex.InnerException is SocketException && ((SocketException) ex.InnerException).SocketErrorCode == SocketError.TimedOut)
+            {
+                // The expected exception is a bit different in AMQP compared to MQTT/HTTPS
+                return; // expected exception was thrown, so exit the test
+            }
+
+            throw new AssertFailedException("Expected an OperationCanceledException to be thrown since the timeout was set to TimeSpan.Zero");
+        }
+
+        public async Task ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(
+            Client.TransportType transportType,
+            AttestationMechanismType attestationType,
+            EnrollmentType? enrollmentType,
+            TimeSpan timeout)
+        {
+            //Default reprovisioning settings: Hashed allocation, no reprovision policy, hub names, or custom allocation policy
+            await ProvisioningDeviceClientValidRegistrationIdRegisterOkAsync(transportType, attestationType, enrollmentType, false, null, AllocationPolicy.Hashed, null, null, null, timeout, s_proxyServerAddress).ConfigureAwait(false);
+        }
+
         public async Task ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(
             Client.TransportType transportType,
             AttestationMechanismType attestationType,
@@ -440,7 +497,7 @@ public async Task ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(
             string proxyServerAddress = null)
         {
             //Default reprovisioning settings: Hashed allocation, no reprovision policy, hub names, or custom allocation policy
-            await ProvisioningDeviceClientValidRegistrationIdRegisterOkAsync(transportType, attestationType, enrollmentType, setCustomProxy, null, AllocationPolicy.Hashed, null, null, null, s_proxyServerAddress).ConfigureAwait(false);
+            await ProvisioningDeviceClientValidRegistrationIdRegisterOkAsync(transportType, attestationType, enrollmentType, setCustomProxy, null, AllocationPolicy.Hashed, null, null, null, TimeSpan.MaxValue, proxyServerAddress).ConfigureAwait(false);
         }
 
         public async Task ProvisioningDeviceClient_ValidRegistrationId_Register_Ok(
@@ -463,7 +520,8 @@ await ProvisioningDeviceClientValidRegistrationIdRegisterOkAsync(
                     null,
                     iothubs,
                     capabilities,
-                    s_proxyServerAddress)
+                    TimeSpan.MaxValue,
+                    proxyServerAddress)
                 .ConfigureAwait(false);
         }
 
@@ -477,6 +535,7 @@ private async Task ProvisioningDeviceClientValidRegistrationIdRegisterOkAsync(
             CustomAllocationDefinition customAllocationDefinition,
             ICollection<string> iothubs,
             DeviceCapabilities deviceCapabilities,
+            TimeSpan timeout,
             string proxyServerAddress = null)
         {
             string groupId = _idPrefix + AttestationTypeToString(attestationType) + "-" + Guid.NewGuid();
@@ -516,7 +575,14 @@ private async Task ProvisioningDeviceClientValidRegistrationIdRegisterOkAsync(
             {
                 try
                 {
-                    result = await provClient.RegisterAsync(cts.Token).ConfigureAwait(false);
+                    if (timeout != TimeSpan.MaxValue)
+                    {
+                        result = await provClient.RegisterAsync(timeout).ConfigureAwait(false);
+                    }
+                    else
+                    { 
+                        result = await provClient.RegisterAsync(cts.Token).ConfigureAwait(false);
+                    }
                     break;
                 }
                 // Catching all ProvisioningTransportException as the status code is not the same for Mqtt, Amqp and Http.

provisioning/device/src/ProvisioningDeviceClient.cs

@@ -3,6 +3,7 @@
 
 using Microsoft.Azure.Devices.Provisioning.Client.Transport;
 using Microsoft.Azure.Devices.Shared;
+using System;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -63,52 +64,87 @@ private ProvisioningDeviceClient(
         /// <summary>
         /// Registers the current device using the Device Provisioning Service and assigns it to an IoT Hub.
         /// </summary>
+        /// <param name="timeout">The maximum amount of time to allow this operation to run for before timing out.</param>
+        /// <remarks>
+        /// Due to the AMQP library used by this library uses not accepting cancellation tokens, this overload and <see cref="RegisterAsync(ProvisioningRegistrationAdditionalData, TimeSpan)"/>
+        /// are the only overloads for this method that allow for a specified timeout to be respected in the middle of an AMQP operation such as opening
+        /// the AMQP connection. MQTT and HTTPS connections do not share that same limitation, though.
+        /// </remarks>
         /// <returns>The registration result.</returns>
-        public Task<DeviceRegistrationResult> RegisterAsync()
+        public Task<DeviceRegistrationResult> RegisterAsync(TimeSpan timeout)
         {
-            return RegisterAsync(CancellationToken.None);
+            return RegisterAsync(null, timeout);
         }
 
         /// <summary>
         /// Registers the current device using the Device Provisioning Service and assigns it to an IoT Hub.
         /// </summary>
-        /// <param name="data">The optional additional data.</param>
+        /// <param name="data">
+        /// The optional additional data that is passed through to the custom allocation policy webhook if 
+        /// a custom allocation policy webhook is setup for this enrollment.
+        /// </param>
+        /// <param name="timeout">The maximum amount of time to allow this operation to run for before timing out.</param>
+        /// <remarks>
+        /// Due to the AMQP library used by this library uses not accepting cancellation tokens, this overload and <see cref="RegisterAsync(TimeSpan)"/>
+        /// are the only overloads for this method that allow for a specified timeout to be respected in the middle of an AMQP operation such as opening
+        /// the AMQP connection. MQTT and HTTPS connections do not share that same limitation, though.
+        /// </remarks>
         /// <returns>The registration result.</returns>
-        public Task<DeviceRegistrationResult> RegisterAsync(ProvisioningRegistrationAdditionalData data)
+        public Task<DeviceRegistrationResult> RegisterAsync(ProvisioningRegistrationAdditionalData data, TimeSpan timeout)
         {
-            return RegisterAsync(data, CancellationToken.None);
+            Logging.RegisterAsync(this, _globalDeviceEndpoint, _idScope, _transport, _security);
+
+            var request = new ProvisioningTransportRegisterMessage(_globalDeviceEndpoint, _idScope, _security, data?.JsonData)
+            {
+                ProductInfo = ProductInfo,
+            };
+
+            return _transport.RegisterAsync(request, timeout);
         }
 
         /// <summary>
         /// Registers the current device using the Device Provisioning Service and assigns it to an IoT Hub.
         /// </summary>
         /// <param name="cancellationToken">The cancellation token.</param>
+        /// <remarks>
+        /// Due to the AMQP library used by this library uses not accepting cancellation tokens, the provided cancellation token will only be checked 
+        /// for cancellation in between AMQP operations, and not during. In order to have a timeout for this operation that is checked during AMQP operations
+        /// (such as opening the connection), you must use <see cref="RegisterAsync(TimeSpan)"/> instead. MQTT and HTTPS connections do not have the same 
+        /// behavior as AMQP connections in this regard. MQTT and HTTPS connections will check this cancellation token for cancellation during their protocol level operations.
+        /// </remarks>
         /// <returns>The registration result.</returns>
-        public Task<DeviceRegistrationResult> RegisterAsync(CancellationToken cancellationToken)
+        public Task<DeviceRegistrationResult> RegisterAsync(CancellationToken cancellationToken = default)
         {
             Logging.RegisterAsync(this, _globalDeviceEndpoint, _idScope, _transport, _security);
 
-            var request = new ProvisioningTransportRegisterMessage(_globalDeviceEndpoint, _idScope, _security)
-            {
-                ProductInfo = ProductInfo
-            };
-            return _transport.RegisterAsync(request, cancellationToken);
+            return RegisterAsync(null, cancellationToken);
         }
 
         /// <summary>
         /// Registers the current device using the Device Provisioning Service and assigns it to an IoT Hub.
         /// </summary>
-        /// <param name="data">The custom content.</param>
+        /// <param name="data">
+        /// The optional additional data that is passed through to the custom allocation policy webhook if 
+        /// a custom allocation policy webhook is setup for this enrollment.
+        /// </param>
         /// <param name="cancellationToken">The cancellation token.</param>
+        /// <remarks>
+        /// Due to the AMQP library used by this library uses not accepting cancellation tokens, the provided cancellation token will only be checked 
+        /// for cancellation in between AMQP operations, and not during. In order to have a timeout for this operation that is checked during AMQP operations
+        /// (such as opening the connection), you must use <see cref="RegisterAsync(ProvisioningRegistrationAdditionalData, TimeSpan)">this overload</see> instead. 
+        /// MQTT and HTTPS connections do not have the same behavior as AMQP connections in this regard. MQTT and HTTPS connections will check this cancellation 
+        /// token for cancellation during their protocol level operations.
+        /// </remarks>
         /// <returns>The registration result.</returns>
-        public Task<DeviceRegistrationResult> RegisterAsync(ProvisioningRegistrationAdditionalData data, CancellationToken cancellationToken)
+        public Task<DeviceRegistrationResult> RegisterAsync(ProvisioningRegistrationAdditionalData data, CancellationToken cancellationToken = default)
         {
             Logging.RegisterAsync(this, _globalDeviceEndpoint, _idScope, _transport, _security);
 
             var request = new ProvisioningTransportRegisterMessage(_globalDeviceEndpoint, _idScope, _security, data?.JsonData)
             {
                 ProductInfo = ProductInfo,
             };
+
             return _transport.RegisterAsync(request, cancellationToken);
         }
     }

provisioning/device/src/ProvisioningTransportHandler.cs

@@ -85,6 +85,19 @@ public virtual Task<DeviceRegistrationResult> RegisterAsync(
             return _innerHandler.RegisterAsync(message, cancellationToken);
         }
 
+        /// <summary>
+        /// Registers a device described by the message.
+        /// </summary>
+        /// <param name="message">The provisioning message.</param>
+        /// <param name="timeout">The maximum amount of time to allow this operation to run for before timing out.</param>
+        /// <returns>The registration result.</returns>
+        public virtual Task<DeviceRegistrationResult> RegisterAsync(
+            ProvisioningTransportRegisterMessage message,
+            TimeSpan timeout)
+        {
+            return _innerHandler.RegisterAsync(message, timeout);
+        }
+
         /// <summary>
         /// Releases the unmanaged resources and disposes of the managed resources used by the invoker.
         /// </summary>