From 00304a0556491e42442bdf497152b21d8e51f9f8 Mon Sep 17 00:00:00 2001 From: Ahson Khan Date: Mon, 22 Jul 2024 18:27:16 -0700 Subject: [PATCH] Add support for customizing the IMDS endpoint within ManagedIdentityCredential using an env var. (#5834) * Add support for customizing the IMDS endpoint within ManagedIdentityCredential using an env var. * Clean up the impl. * Add imds as an opt-out for spell checl. * Address PR feedback. --- .vscode/cspell.json | 1 + sdk/identity/azure-identity/CHANGELOG.md | 1 + .../src/managed_identity_source.cpp | 19 ++- .../src/private/managed_identity_source.hpp | 1 + .../ut/managed_identity_credential_test.cpp | 124 ++++++++++++++++++ 5 files changed, 142 insertions(+), 4 deletions(-) diff --git a/.vscode/cspell.json b/.vscode/cspell.json index 9b4028a90..6a04ec4a7 100644 --- a/.vscode/cspell.json +++ b/.vscode/cspell.json @@ -130,6 +130,7 @@ "hlocal", "HLOCAL", "HRESULT", + "imds", "Imds", "IMDS", "immutability", diff --git a/sdk/identity/azure-identity/CHANGELOG.md b/sdk/identity/azure-identity/CHANGELOG.md index 648500e47..fbb5983ff 100644 --- a/sdk/identity/azure-identity/CHANGELOG.md +++ b/sdk/identity/azure-identity/CHANGELOG.md @@ -5,6 +5,7 @@ ### Features Added - Added `ClientAssertionCredential` to enable applications to authenticate with custom client assertions. +- Added support for customizing the IMDS endpoint within `ManagedIdentityCredential`. ### Breaking Changes diff --git a/sdk/identity/azure-identity/src/managed_identity_source.cpp b/sdk/identity/azure-identity/src/managed_identity_source.cpp index 1b7d6467e..e3f56bd9b 100644 --- a/sdk/identity/azure-identity/src/managed_identity_source.cpp +++ b/sdk/identity/azure-identity/src/managed_identity_source.cpp @@ -449,16 +449,27 @@ std::unique_ptr ImdsManagedIdentitySource::Create( credName + " will be created" + WithSourceMessage("Azure Instance Metadata Service") + ".\nSuccessful creation does not guarantee further successful token retrieval."); - return std::unique_ptr(new ImdsManagedIdentitySource(clientId, options)); + std::string imdsHost = "http://169.254.169.254"; + std::string customImdsHost = Environment::GetVariable("AZURE_IMDS_CUSTOM_AUTHORITY_HOST"); + if (!customImdsHost.empty()) + { + IdentityLog::Write( + IdentityLog::Level::Informational, "Custom IMDS host is set to: " + customImdsHost); + imdsHost = customImdsHost; + } + Azure::Core::Url imdsUrl(imdsHost); + imdsUrl.AppendPath("/metadata/identity/oauth2/token"); + + return std::unique_ptr( + new ImdsManagedIdentitySource(clientId, imdsUrl, options)); } ImdsManagedIdentitySource::ImdsManagedIdentitySource( std::string const& clientId, + Azure::Core::Url const& imdsUrl, Azure::Core::Credentials::TokenCredentialOptions const& options) : ManagedIdentitySource(clientId, std::string(), options), - m_request( - Azure::Core::Http::HttpMethod::Get, - Azure::Core::Url("http://169.254.169.254/metadata/identity/oauth2/token")) + m_request(Azure::Core::Http::HttpMethod::Get, imdsUrl) { { using Azure::Core::Url; diff --git a/sdk/identity/azure-identity/src/private/managed_identity_source.hpp b/sdk/identity/azure-identity/src/private/managed_identity_source.hpp index 2b5e9e96f..cfeb06714 100644 --- a/sdk/identity/azure-identity/src/private/managed_identity_source.hpp +++ b/sdk/identity/azure-identity/src/private/managed_identity_source.hpp @@ -176,6 +176,7 @@ namespace Azure { namespace Identity { namespace _detail { explicit ImdsManagedIdentitySource( std::string const& clientId, + Core::Url const& imdsUrl, Core::Credentials::TokenCredentialOptions const& options); public: diff --git a/sdk/identity/azure-identity/test/ut/managed_identity_credential_test.cpp b/sdk/identity/azure-identity/test/ut/managed_identity_credential_test.cpp index edad7e6f4..5c0a8da69 100644 --- a/sdk/identity/azure-identity/test/ut/managed_identity_credential_test.cpp +++ b/sdk/identity/azure-identity/test/ut/managed_identity_credential_test.cpp @@ -1847,3 +1847,127 @@ TEST(ManagedIdentityCredential, ImdsCreation) EXPECT_GE(response2.AccessToken.ExpiresOn, response2.EarliestExpiration + 3600s); EXPECT_LE(response2.AccessToken.ExpiresOn, response2.LatestExpiration + 3600s); } + +TEST(ManagedIdentityCredential, ImdsCustomHost) +{ + auto const actual1 = CredentialTestHelper::SimulateTokenRequest( + [](auto transport) { + TokenCredentialOptions options; + options.Transport.Transport = transport; + + CredentialTestHelper::EnvironmentOverride const env({ + {"MSI_ENDPOINT", ""}, + {"MSI_SECRET", ""}, + {"IDENTITY_ENDPOINT", "https://visualstudio.com/"}, + {"IMDS_ENDPOINT", ""}, + {"IDENTITY_HEADER", ""}, + {"IDENTITY_SERVER_THUMBPRINT", ""}, + {"AZURE_IMDS_CUSTOM_AUTHORITY_HOST", ""}, + }); + + return std::make_unique( + "fedcba98-7654-3210-0123-456789abcdef", options); + }, + {{"https://azure.com/.default"}}, + {"{\"expires_in\":3600, \"access_token\":\"ACCESSTOKEN1\"}"}); + + auto const actual2 = CredentialTestHelper::SimulateTokenRequest( + [](auto transport) { + TokenCredentialOptions options; + options.Transport.Transport = transport; + + CredentialTestHelper::EnvironmentOverride const env({ + {"MSI_ENDPOINT", ""}, + {"MSI_SECRET", ""}, + {"IDENTITY_ENDPOINT", ""}, + {"IMDS_ENDPOINT", "https://xbox.com/"}, + {"IDENTITY_HEADER", ""}, + {"IDENTITY_SERVER_THUMBPRINT", ""}, + {"AZURE_IMDS_CUSTOM_AUTHORITY_HOST", "https://custom.imds.endpoint/"}, + }); + + return std::make_unique( + "01234567-89ab-cdef-fedc-ba9876543210", options); + }, + {{"https://outlook.com/.default"}}, + {"{\"expires_in\":7200, \"access_token\":\"ACCESSTOKEN2\"}"}); + + auto const actual3 = CredentialTestHelper::SimulateTokenRequest( + [](auto transport) { + TokenCredentialOptions options; + options.Transport.Transport = transport; + + CredentialTestHelper::EnvironmentOverride const env({ + {"MSI_ENDPOINT", ""}, + {"MSI_SECRET", ""}, + {"IDENTITY_ENDPOINT", ""}, + {"IMDS_ENDPOINT", "https://xbox.com/"}, + {"IDENTITY_HEADER", ""}, + {"IDENTITY_SERVER_THUMBPRINT", ""}, + {"AZURE_IMDS_CUSTOM_AUTHORITY_HOST", "http://localhost:59202"}, + }); + + return std::make_unique( + "01234567-89ab-cdef-fedc-ba9876543210", options); + }, + {{"https://outlook.com/.default"}}, + {"{\"expires_in\":7200, \"access_token\":\"ACCESSTOKEN2\"}"}); + + EXPECT_EQ(actual1.Requests.size(), 1U); + EXPECT_EQ(actual1.Responses.size(), 1U); + + EXPECT_EQ(actual2.Requests.size(), 1U); + EXPECT_EQ(actual2.Responses.size(), 1U); + + auto const& request1 = actual1.Requests.at(0); + auto const& response1 = actual1.Responses.at(0); + + auto const& request2 = actual2.Requests.at(0); + auto const& response2 = actual2.Responses.at(0); + + EXPECT_EQ(request1.HttpMethod, HttpMethod::Get); + EXPECT_EQ(request2.HttpMethod, HttpMethod::Get); + + EXPECT_EQ( + request1.AbsoluteUrl, + "http://169.254.169.254/metadata/identity/oauth2/token" + "?api-version=2018-02-01" + "&client_id=fedcba98-7654-3210-0123-456789abcdef" + "&resource=https%3A%2F%2Fazure.com"); // cspell:disable-line + + EXPECT_EQ( + request2.AbsoluteUrl, + "https://custom.imds.endpoint/metadata/identity/oauth2/token" + "?api-version=2018-02-01" + "&client_id=01234567-89ab-cdef-fedc-ba9876543210" + "&resource=https%3A%2F%2Foutlook.com"); // cspell:disable-line + + auto const& request3 = actual3.Requests.at(0); + EXPECT_EQ( + request3.AbsoluteUrl, + "http://localhost:59202/metadata/identity/oauth2/token" + "?api-version=2018-02-01" + "&client_id=01234567-89ab-cdef-fedc-ba9876543210" + "&resource=https%3A%2F%2Foutlook.com"); // cspell:disable-line + + EXPECT_TRUE(request1.Body.empty()); + EXPECT_TRUE(request2.Body.empty()); + + { + EXPECT_NE(request1.Headers.find("Metadata"), request1.Headers.end()); + EXPECT_EQ(request1.Headers.at("Metadata"), "true"); + + EXPECT_NE(request2.Headers.find("Metadata"), request2.Headers.end()); + EXPECT_EQ(request2.Headers.at("Metadata"), "true"); + } + + EXPECT_EQ(response1.AccessToken.Token, "ACCESSTOKEN1"); + EXPECT_EQ(response2.AccessToken.Token, "ACCESSTOKEN2"); + + using namespace std::chrono_literals; + EXPECT_GE(response1.AccessToken.ExpiresOn, response1.EarliestExpiration + 3600s); + EXPECT_LE(response1.AccessToken.ExpiresOn, response1.LatestExpiration + 3600s); + + EXPECT_GE(response2.AccessToken.ExpiresOn, response2.EarliestExpiration + 3600s); + EXPECT_LE(response2.AccessToken.ExpiresOn, response2.LatestExpiration + 3600s); +}