Skip to content

Commit

Permalink
add env variable
Browse files Browse the repository at this point in the history
  • Loading branch information
@sourabh1007
sourabh1007 committed Jul 3, 2024
1 parent 5287ffb commit db04319
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 21 deletions.
22 changes: 14 additions & 8 deletions Microsoft.Azure.Cosmos/src/Telemetry/VmMetadataApiHandler.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,28 +34,34 @@ internal static class VmMetadataApiHandler

private static bool isInitialized = false;
private static AzureVMMetadata azMetadata = null;

internal static void TryInitialize(CosmosHttpClient httpClient)
{
if (VmMetadataApiHandler.isInitialized)
bool isVMMetadataAccessDisabled = ConfigurationManager.GetEnvironmentVariable<bool>("COSMOS_DISABLE_VM_METADATA_ACCESS", false);
if (System.Diagnostics.Debugger.IsAttached || isVMMetadataAccessDisabled)
{
return;
}

lock (VmMetadataApiHandler.lockObject)
else
{
if (VmMetadataApiHandler.isInitialized)
{
return;
}

DefaultTrace.TraceInformation("Initializing VM Metadata API ");
lock (VmMetadataApiHandler.lockObject)
{
if (VmMetadataApiHandler.isInitialized)
{
return;
}

DefaultTrace.TraceInformation("Initializing VM Metadata API ");

VmMetadataApiHandler.isInitialized = true;
VmMetadataApiHandler.isInitialized = true;

_ = Task.Run(() => MetadataApiCallAsync(httpClient), default);
_ = Task.Run(() => MetadataApiCallAsync(httpClient), default);
}
}

}

private static async Task MetadataApiCallAsync(CosmosHttpClient httpClient)
Expand Down
49 changes: 36 additions & 13 deletions Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/VmMetadataApiHandlerTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,14 @@ namespace Microsoft.Azure.Cosmos
using Microsoft.Azure.Cosmos.Telemetry.Models;
using Microsoft.Azure.Cosmos.Tests;
using Microsoft.VisualStudio.TestTools.UnitTesting;
using Moq;
using Newtonsoft.Json;
using Util;

[TestClass]
public class VmMetadataApiHandlerTest
{
[TestInitialize]
public void Intialize()
public void Initialize()
{
var isInitializedField = typeof(VmMetadataApiHandler).GetField("isInitialized",
BindingFlags.Static |
Expand All @@ -39,8 +38,21 @@ public void Intialize()
}

[TestMethod]
public async Task GetVmIdAsMachineIdTest()
[DataRow(true, false)]
[DataRow(false, false)]
[DataRow(false, true)]
public async Task GetVmIdAsMachineIdTest(bool isVmMetadataAccessDisabled, bool isDebuggerAttached)
{
if (isVmMetadataAccessDisabled)
{
Environment.SetEnvironmentVariable("COSMOS_DISABLE_VM_METADATA_ACCESS", "true");
}

if (isDebuggerAttached)
{
System.Diagnostics.Debugger.Launch();
}

static Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, CancellationToken cancellationToken)
{
object jsonObject = JsonConvert.DeserializeObject("{\"compute\":{\"azEnvironment\":\"AzurePublicCloud\",\"customData\":\"\",\"isHostCompatibilityLayerVm\":\"false\",\"licenseType\":\"\",\"location\":\"eastus\",\"name\":\"sourabh-testing\",\"offer\":\"UbuntuServer\",\"osProfile\":{\"adminUsername\":\"azureuser\",\"computerName\":\"sourabh-testing\"},\"osType\":\"Linux\",\"placementGroupId\":\"\",\"plan\":{\"name\":\"\",\"product\":\"\",\"publisher\":\"\"},\"platformFaultDomain\":\"0\",\"platformUpdateDomain\":\"0\",\"provider\":\"Microsoft.Compute\",\"publicKeys\":[{\"keyData\":\"ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC5uCeOAm3ehmhI+2PbMoMl17Eo\r\nqfHKCycSaBJsv9qxlmBOuFheSJc1XknJleXUSsuTO016/d1PyWpevnqOZNRksWoa\r\nJvQ23sDTxcK+X2OP3QlCUeX4cMjPXqlL8z1UYzU4Bx3fFvf8fs67G3N72sxWBw5P\r\nZyuXyhBm0NCe/2NYMKgEDT4ma8XszO0ikbhoPKbMbgHAQk/ktWQHNcqYOPQKEWqp\r\nEK1R0rjS2nmtovfScP/ZGXcvOpJ1/NDBo4dh1K+OxOGM/4PSH/F448J5Zy4eAyEk\r\nscys+IpeIOTOlRUy/703SNIX0LEWlnYqbyL9c1ypcYLQqF76fKkDfzzFI/OWVlGw\r\nhj/S9uP8iMsR+fhGIbn6MAa7O4DWPWLuedSp7KDYyjY09gqNJsfuaAJN4LiC6bPy\r\nhknm0PVLK3ux7EUOt+cZrHCdIFWbdOtxiPNIl1tkv9kV5aE5Aj2gJm4MeB9uXYhS\r\nOuksboBc0wyUGrl9+XZJ1+NlZOf7IjVi86CieK8= generated-by-azure\r\n\",\"path\":\"/home/azureuser/.ssh/authorized_keys\"}],\"publisher\":\"Canonical\",\"resourceGroupName\":\"sourabh-telemetry-sdk\",\"resourceId\":\"/subscriptions/8fba6d4f-7c37-4d13-9063-fd58ad2b86e2/resourceGroups/sourabh-telemetry-sdk/providers/Microsoft.Compute/virtualMachines/sourabh-testing\",\"securityProfile\":{\"secureBootEnabled\":\"false\",\"virtualTpmEnabled\":\"false\"},\"sku\":\"18.04-LTS\",\"storageProfile\":{\"dataDisks\":[],\"imageReference\":{\"id\":\"\",\"offer\":\"UbuntuServer\",\"publisher\":\"Canonical\",\"sku\":\"18.04-LTS\",\"version\":\"latest\"},\"osDisk\":{\"caching\":\"ReadWrite\",\"createOption\":\"FromImage\",\"diffDiskSettings\":{\"option\":\"\"},\"diskSizeGB\":\"30\",\"encryptionSettings\":{\"enabled\":\"false\"},\"image\":{\"uri\":\"\"},\"managedDisk\":{\"id\":\"/subscriptions/8fba6d4f-7c37-4d13-9063-fd58ad2b86e2/resourceGroups/sourabh-telemetry-sdk/providers/Microsoft.Compute/disks/sourabh-testing_OsDisk_1_9a54abfc5ba149c6a106bd9e5b558c2a\",\"storageAccountType\":\"Premium_LRS\"},\"name\":\"sourabh-testing_OsDisk_1_9a54abfc5ba149c6a106bd9e5b558c2a\",\"osType\":\"Linux\",\"vhd\":{\"uri\":\"\"},\"writeAcceleratorEnabled\":\"false\"}},\"subscriptionId\":\"8fba6d4f-7c37-4d13-9063-fd58ad2b86e2\",\"tags\":\"azsecpack:nonprod;platformsettings.host_environment.service.platform_optedin_for_rootcerts:true\",\"tagsList\":[{\"name\":\"azsecpack\",\"value\":\"nonprod\"},{\"name\":\"platformsettings.host_environment.service.platform_optedin_for_rootcerts\",\"value\":\"true\"}],\"version\":\"18.04.202103250\",\"vmId\":\"d0cb93eb-214b-4c2b-bd3d-cc93e90d9efd\",\"vmScaleSetName\":\"\",\"vmSize\":\"Standard_D2s_v3\",\"zone\":\"1\"},\"network\":{\"interface\":[{\"ipv4\":{\"ipAddress\":[{\"privateIpAddress\":\"10.0.7.5\",\"publicIpAddress\":\"\"}],\"subnet\":[{\"address\":\"10.0.7.0\",\"prefix\":\"24\"}]},\"ipv6\":{\"ipAddress\":[]},\"macAddress\":\"000D3A8F8BA0\"}]}}");
Expand All @@ -53,13 +65,24 @@ static Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellati
}

HttpMessageHandler messageHandler = new MockMessageHandler(sendFunc);
CosmosHttpClient cosmoshttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));
CosmosHttpClient cosmosHttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));

VmMetadataApiHandler.TryInitialize(cosmoshttpClient);
VmMetadataApiHandler.TryInitialize(cosmosHttpClient);

await Task.Delay(2000);
Assert.AreEqual($"{VmMetadataApiHandler.VmIdPrefix}{"d0cb93eb-214b-4c2b-bd3d-cc93e90d9efd"}", VmMetadataApiHandler.GetMachineId());
Assert.AreEqual(VmMetadataApiHandler.GetMachineRegion(), "eastus");

if (isVmMetadataAccessDisabled || isDebuggerAttached)
{
Assert.AreNotEqual($"{VmMetadataApiHandler.VmIdPrefix}{"d0cb93eb-214b-4c2b-bd3d-cc93e90d9efd"}", VmMetadataApiHandler.GetMachineId());
Assert.IsNull(VmMetadataApiHandler.GetMachineRegion(), VmMetadataApiHandler.GetMachineRegion());

Environment.SetEnvironmentVariable("COSMOS_DISABLE_VM_METADATA_ACCESS", "false");
}
else
{
Assert.AreEqual($"{VmMetadataApiHandler.VmIdPrefix}{"d0cb93eb-214b-4c2b-bd3d-cc93e90d9efd"}", VmMetadataApiHandler.GetMachineId());
Assert.AreEqual(VmMetadataApiHandler.GetMachineRegion(), "eastus");
}
}

[TestMethod]
Expand All @@ -77,9 +100,9 @@ static Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellati
}

HttpMessageHandler messageHandler = new MockMessageHandler(sendFunc);
CosmosHttpClient cosmoshttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));
CosmosHttpClient cosmosHttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));

VmMetadataApiHandler.TryInitialize(cosmoshttpClient);
VmMetadataApiHandler.TryInitialize(cosmosHttpClient);

await Task.Delay(2000);
Assert.IsNull(VmMetadataApiHandler.GetMachineInfo());
Expand All @@ -95,9 +118,9 @@ public async Task GetHashedMachineNameAsMachineIdTest()
static Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, CancellationToken cancellationToken) { throw new Exception("error while making API call"); };

HttpMessageHandler messageHandler = new MockMessageHandler(sendFunc);
CosmosHttpClient cosmoshttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));
CosmosHttpClient cosmosHttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));

VmMetadataApiHandler.TryInitialize(cosmoshttpClient);
VmMetadataApiHandler.TryInitialize(cosmosHttpClient);

await Task.Delay(2000);
Assert.AreEqual(expectedMachineId, VmMetadataApiHandler.GetMachineId());
Expand Down Expand Up @@ -139,7 +162,7 @@ public void CatchMetadataApiCallExceptionTest()
static Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, CancellationToken cancellationToken) { throw new Exception("error while making API call"); };

HttpMessageHandler messageHandler = new MockMessageHandler(sendFunc);
CosmosHttpClient cosmoshttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));
CosmosHttpClient cosmosHttpClient = MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(messageHandler));

string expectedMsg = "Azure Environment metadata information not available.";
ManualResetEvent manualResetEvent = new ManualResetEvent(false);
Expand All @@ -155,7 +178,7 @@ void TraceHandler(string message)
DefaultTrace.TraceSource.Listeners.Add(new TestTraceListener { Callback = TraceHandler });
DefaultTrace.InitEventListener();

VmMetadataApiHandler.TryInitialize(cosmoshttpClient);
VmMetadataApiHandler.TryInitialize(cosmosHttpClient);

int timeout = 30000;
Assert.IsTrue(manualResetEvent.WaitOne(timeout));
Expand Down

0 comments on commit db04319

Please sign in to comment.