Throw exception if IoT device module invokes an IoT edge module API

iothub/device/src/Edge/EdgeModuleClientHelper.cs

@@ -7,6 +7,7 @@
 using System.Diagnostics;
 using System.Runtime.InteropServices;
 using System.Security.Cryptography.X509Certificates;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Azure.Devices.Client.HsmAuthentication;
 using static System.Runtime.InteropServices.RuntimeInformation;
@@ -84,7 +85,8 @@ internal static IotHubConnectionCredentials CreateIotHubConnectionCredentialsFro
 
         internal static async Task<ICertificateValidator> CreateCertificateValidatorFromEnvironmentAsync(
             ITrustBundleProvider trustBundleProvider, 
-            IotHubClientOptions options)
+            IotHubClientOptions options,
+            CancellationToken cancellationToken)
         {
             Debug.Assert(options != null);
 
@@ -122,7 +124,7 @@ internal static async Task<ICertificateValidator> CreateCertificateValidatorFrom
             if (!string.IsNullOrEmpty(gateway))
             {
                 IList<X509Certificate2> certs = await trustBundleProvider
-                    .GetTrustBundleAsync(new Uri(edgeWorkloadUri), EdgeHsmApiVersion)
+                    .GetTrustBundleAsync(new Uri(edgeWorkloadUri), EdgeHsmApiVersion, cancellationToken)
                     .ConfigureAwait(false);
                 certificateValidator = CreateCertificateValidator(certs, options);
             }

iothub/device/src/Edge/ITrustBundleProvider.cs

@@ -4,12 +4,13 @@
 using System;
 using System.Collections.Generic;
 using System.Security.Cryptography.X509Certificates;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Microsoft.Azure.Devices.Client
 {
     internal interface ITrustBundleProvider
     {
-        Task<IList<X509Certificate2>> GetTrustBundleAsync(Uri providerUri, string defaultApiVersion);
+        Task<IList<X509Certificate2>> GetTrustBundleAsync(Uri providerUri, string defaultApiVersion, CancellationToken cancellationToken);
     }
 }

iothub/device/src/Edge/TrustBundleProvider.cs

@@ -7,6 +7,7 @@
 using System.Net.Http;
 using System.Security.Cryptography.X509Certificates;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Azure.Devices.Client.HsmAuthentication;
 using Microsoft.Azure.Devices.Client.HsmAuthentication.GeneratedCode;
@@ -17,7 +18,7 @@ internal sealed class TrustBundleProvider : ITrustBundleProvider
     {
         private static readonly IIotHubClientRetryPolicy s_retryPolicy = new IotHubClientExponentialBackoffRetryPolicy(3, TimeSpan.FromSeconds(30));
 
-        public async Task<IList<X509Certificate2>> GetTrustBundleAsync(Uri providerUri, string apiVersion)
+        public async Task<IList<X509Certificate2>> GetTrustBundleAsync(Uri providerUri, string apiVersion, CancellationToken cancellationToken)
         {
             try
             {
@@ -26,7 +27,7 @@ public async Task<IList<X509Certificate2>> GetTrustBundleAsync(Uri providerUri,
                 {
                     BaseUrl = HttpClientHelper.GetBaseUri(providerUri)
                 };
-                TrustBundleResponse response = await GetTrustBundleWithRetryAsync(hsmHttpClient, apiVersion).ConfigureAwait(false);
+                TrustBundleResponse response = await GetTrustBundleWithRetryAsync(hsmHttpClient, apiVersion, cancellationToken).ConfigureAwait(false);
 
                 IList<X509Certificate2> certs = ParseCertificates(response.Certificate);
                 return certs;
@@ -45,13 +46,15 @@ public async Task<IList<X509Certificate2>> GetTrustBundleAsync(Uri providerUri,
 
         private static async Task<TrustBundleResponse> GetTrustBundleWithRetryAsync(
             HttpHsmClient hsmHttpClient,
-            string apiVersion)
+            string apiVersion,
+            CancellationToken cancellationToken)
         {
             var transientRetryPolicy = new RetryHandler(s_retryPolicy);
             return await transientRetryPolicy
                 .RunWithRetryAsync(
                     () => hsmHttpClient.TrustBundleAsync(apiVersion),
-                    (Exception ex) => ex is SwaggerException se && se.StatusCode >= 500)
+                    (Exception ex) => ex is SwaggerException se && se.StatusCode >= 500,
+                    cancellationToken)
                 .ConfigureAwait(false);
         }
 

iothub/device/src/IotHubModuleClient.cs

@@ -21,6 +21,9 @@ public class IotHubModuleClient : IotHubBaseClient
         private const string ModuleMethodUriFormat = "/twins/{0}/modules/{1}/methods?" + ClientApiVersionHelper.ApiVersionQueryStringLatest;
         private const string DeviceMethodUriFormat = "/twins/{0}/methods?" + ClientApiVersionHelper.ApiVersionQueryStringLatest;
 
+        private const string IotDeviceModuleMethodInvokeErrorMessage = "This API call is relevant only for IoT Edge modules. Please make sure your client is initialized correctly with a gateway hostname. " +
+            "For subscribing to IoT device module direct method invocations, see SetDirectMethodCallbackAsync(...).";
+
         /// <summary>
         /// Creates a disposable client from the specified connection string.
         /// </summary>
@@ -93,29 +96,34 @@ internal IotHubModuleClient(IotHubConnectionCredentials iotHubConnectionCredenti
             if (Logging.IsEnabled)
                 Logging.CreateClient(
                     this,
-                    $"HostName={IotHubConnectionCredentials.HostName};DeviceId={IotHubConnectionCredentials.DeviceId};ModuleId={IotHubConnectionCredentials.ModuleId}",
+                    $"HostName={IotHubConnectionCredentials.HostName};DeviceId={IotHubConnectionCredentials.DeviceId};ModuleId={IotHubConnectionCredentials.ModuleId};isEdgeModule={IotHubConnectionCredentials.IsEdgeModule}",
                     _clientOptions);
         }
 
         /// <summary>
         /// Creates a disposable <c>IotHubModuleClient</c> instance in an IoT Edge deployment based on environment variables.
         /// </summary>
         /// <param name="options">The options that allow configuration of the module client instance during initialization.</param>
+        /// <param name="cancellationToken">A cancellation token to cancel the operation.</param>
         /// <returns>A disposable client instance.</returns>
         /// <exception cref="InvalidOperationException">The required environmental variables were missing. Check the exception thrown for additional details.</exception>
+        /// <exception cref="OperationCanceledException">The operation has been canceled.</exception>
         /// <example>
         /// <code language="csharp">
         /// await using var client = await IotHubModuleClient.CreateFromEnvironmentAsync(new IotHubClientOptions(new IotHubClientMqttSettings(IotHubClientTransportProtocol.WebSocket)));
         /// </code>
         /// </example>
-        public static async Task<IotHubModuleClient> CreateFromEnvironmentAsync(IotHubClientOptions options = default)
+        public static async Task<IotHubModuleClient> CreateFromEnvironmentAsync(IotHubClientOptions options = default, CancellationToken cancellationToken = default)
         {
             IotHubClientOptions clientOptions = options != null
                 ? options.Clone()
                 : new();
 
             IotHubConnectionCredentials iotHubConnectionCredentials = EdgeModuleClientHelper.CreateIotHubConnectionCredentialsFromEnvironment();
-            ICertificateValidator certificateValidator = await EdgeModuleClientHelper.CreateCertificateValidatorFromEnvironmentAsync(new TrustBundleProvider(), clientOptions);
+            ICertificateValidator certificateValidator = await EdgeModuleClientHelper.CreateCertificateValidatorFromEnvironmentAsync(
+                new TrustBundleProvider(),
+                clientOptions,
+                cancellationToken);
 
             return new IotHubModuleClient(iotHubConnectionCredentials, options, certificateValidator);
         }
@@ -236,6 +244,7 @@ public async Task SendMessagesToRouteAsync(string outputName, IEnumerable<Teleme
         /// <returns>The result of the method invocation.</returns>
         /// <exception cref="ArgumentNullException"><paramref name="deviceId"/> or <paramref name="methodRequest"/> is null.</exception>
         /// <exception cref="InvalidOperationException">The client instance is not already open.</exception>
+        /// <exception cref="InvalidOperationException">An IoT device module is used to invoke this API.</exception>
         /// <exception cref="OperationCanceledException">The operation has been canceled.</exception>
         /// <exception cref="IotHubClientException">An error occured when communicating with IoT hub service.</exception>
         /// <exception cref="ObjectDisposedException">The client has been disposed.</exception>
@@ -248,6 +257,12 @@ public Task<DirectMethodResponse> InvokeMethodAsync(string deviceId, EdgeModuleD
         {
             Argument.AssertNotNullOrWhiteSpace(deviceId, nameof(deviceId));
             Argument.AssertNotNull(methodRequest, nameof(methodRequest));
+
+            if (!IotHubConnectionCredentials.IsEdgeModule)
+            {
+                throw new InvalidOperationException(IotDeviceModuleMethodInvokeErrorMessage);
+            }
+
             cancellationToken.ThrowIfCancellationRequested();
 
             return InvokeMethodAsync(GetDeviceMethodUri(deviceId), methodRequest, cancellationToken);
@@ -270,6 +285,7 @@ public Task<DirectMethodResponse> InvokeMethodAsync(string deviceId, EdgeModuleD
         /// <returns>The result of the method invocation.</returns>
         /// <exception cref="ArgumentNullException"><paramref name="deviceId"/>, <paramref name="moduleId"/> or <paramref name="methodRequest"/> is null.</exception>
         /// <exception cref="InvalidOperationException">The client instance is not already open.</exception>
+        /// <exception cref="InvalidOperationException">An IoT device module is used to invoke this API.</exception>
         /// <exception cref="OperationCanceledException">The operation has been canceled.</exception>
         /// <exception cref="IotHubClientException">An error occured when communicating with IoT hub service.</exception>
         /// <exception cref="ObjectDisposedException">The client has been disposed.</exception>
@@ -283,6 +299,12 @@ public Task<DirectMethodResponse> InvokeMethodAsync(string deviceId, string modu
             Argument.AssertNotNullOrWhiteSpace(deviceId, nameof(deviceId));
             Argument.AssertNotNullOrWhiteSpace(moduleId, nameof(moduleId));
             Argument.AssertNotNull(methodRequest, nameof(methodRequest));
+
+            if (!IotHubConnectionCredentials.IsEdgeModule)
+            {
+                throw new InvalidOperationException(IotDeviceModuleMethodInvokeErrorMessage);
+            }
+
             cancellationToken.ThrowIfCancellationRequested();
 
             return InvokeMethodAsync(GetModuleMethodUri(deviceId, moduleId), methodRequest, cancellationToken);

iothub/device/tests/Edge/EdgeModuleClientHelperTest.cs

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Security.Cryptography.X509Certificates;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -186,11 +187,11 @@ public async Task TestCreate_FromEnvironment_SetAmqpTransportSettings_ShouldCrea
             var options = new IotHubClientOptions(settings);
             var trustBundle = new Mock<ITrustBundleProvider>();
             trustBundle
-                .Setup(x => x.GetTrustBundleAsync(It.IsAny<Uri>(), It.IsAny<string>()))
+                .Setup(x => x.GetTrustBundleAsync(It.IsAny<Uri>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
                 .Returns(() => Task.FromResult<IList<X509Certificate2>>(new List<X509Certificate2>(0)));
             IotHubConnectionCredentials creds = EdgeModuleClientHelper.CreateIotHubConnectionCredentialsFromEnvironment();
             ICertificateValidator certValidator = await EdgeModuleClientHelper
-                .CreateCertificateValidatorFromEnvironmentAsync(trustBundle.Object, options)
+                .CreateCertificateValidatorFromEnvironmentAsync(trustBundle.Object, options, CancellationToken.None)
                 .ConfigureAwait(false);
 
             // act
@@ -223,11 +224,11 @@ public async Task TestCreate_FromEnvironment_SetMqttTransportSettings_ShouldCrea
             var options = new IotHubClientOptions(settings);
             var trustBundle = new Mock<ITrustBundleProvider>();
             trustBundle
-                .Setup(x => x.GetTrustBundleAsync(It.IsAny<Uri>(), It.IsAny<string>()))
+                .Setup(x => x.GetTrustBundleAsync(It.IsAny<Uri>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
                 .Returns(() => Task.FromResult<IList<X509Certificate2>>(new List<X509Certificate2>(0)));
             IotHubConnectionCredentials creds = EdgeModuleClientHelper.CreateIotHubConnectionCredentialsFromEnvironment();
             ICertificateValidator certValidator = await EdgeModuleClientHelper
-                .CreateCertificateValidatorFromEnvironmentAsync(trustBundle.Object, options)
+                .CreateCertificateValidatorFromEnvironmentAsync(trustBundle.Object, options, CancellationToken.None)
                 .ConfigureAwait(false);
 
             // act

iothub/device/tests/IotHubModuleClientDisposeTests.cs

@@ -27,7 +27,7 @@ public static async Task ClassInitializeAsync(TestContext context)
             string testSharedAccessKey = Convert.ToBase64String(rndBytes);
             var csBuilder = new IotHubConnectionString(
                 "contoso.azure-devices.net",
-                null,
+                "my-gateway",
                 "deviceId",
                 "moduleId",
                 null,