CosmosClient: Fixes NullReferenceException when AzMetadata.Compute is…

… null (#3242) * add null check for azMetadata compute and add test to confirm * also add null checking for VMId * switch to string null or whitespace check
Azure · Jun 2, 2022 · 7f17fe1 · 7f17fe1
1 parent 1264fb4
commit 7f17fe1
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 1 deletion.
diff --git a/Microsoft.Azure.Cosmos/src/Telemetry/VmMetadataApiHandler.cs b/Microsoft.Azure.Cosmos/src/Telemetry/VmMetadataApiHandler.cs
@@ -103,7 +103,7 @@ internal static async Task<AzureVMMetadata> ProcessResponseAsync(HttpResponseMes
         /// <returns>machine id</returns>
         internal static string GetMachineId()
         {
-            if (VmMetadataApiHandler.azMetadata != null)
+            if (!String.IsNullOrWhiteSpace(VmMetadataApiHandler.azMetadata?.Compute?.VMId))
             {
                 return VmMetadataApiHandler.azMetadata.Compute.VMId;
             }

diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/VmMetadataApiHandlerTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/VmMetadataApiHandlerTest.cs
@@ -59,6 +59,30 @@ static Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellati
             Assert.AreEqual("vmId:d0cb93eb-214b-4c2b-bd3d-cc93e90d9efd", VmMetadataApiHandler.GetMachineId());
         }
 
+        [TestMethod]
+        public async Task GetVMMachineIdWithNullComputeTest()
+        {
+            static Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, CancellationToken cancellationToken)
+            {
+                object jsonObject = JsonConvert.DeserializeObject("{\"compute\":null}");
+                string payload = JsonConvert.SerializeObject(jsonObject);
+                HttpResponseMessage response = new HttpResponseMessage(HttpStatusCode.OK)
+                {
+                    Content = new StringContent(payload, Encoding.UTF8, "application/json")
+                };
+                return Task.FromResult(response);
+            }
+
+            HttpMessageHandler messageHandler = new MockMessageHandler(sendFunc);
+            CosmosHttpClient cosmoshttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));
+
+            VmMetadataApiHandler.TryInitialize(cosmoshttpClient);
+
+            await Task.Delay(2000);
+            Assert.IsNull(VmMetadataApiHandler.GetMachineInfo());
+            Assert.IsNotNull(VmMetadataApiHandler.GetMachineId());
+        }
+
         [TestMethod]
         public async Task GetHashedMachineNameAsMachineIdTest()
         {