Fix IMDS retry policy to address review feedback: keep first request as fast probe

Co-authored-by: antkmsft <41349689+antkmsft@users.noreply.github.com>

sdk/identity/azure-identity/src/managed_identity_source.cpp

@@ -118,6 +118,24 @@ void ValidateArcKeyFile(std::string const& fileName)
     throw AuthenticationException("Failed to get file size for '" + fileName + "'.");
   }
 }
+
+// Create IMDS-specific retry options that handle HTTP 410 responses
+// Note: This is a compromise solution. The ideal implementation would apply
+// extended retry duration only for HTTP 410 responses, which requires
+// Azure Core support for conditional retry behavior.
+Azure::Core::Credentials::TokenCredentialOptions CreateImdsRetryOptions(
+    Azure::Core::Credentials::TokenCredentialOptions const& options)
+{
+  using Azure::Core::Http::HttpStatusCode;
+  
+  auto imdsOptions = options;
+  
+  // Add HTTP 410 (Gone) to the retryable status codes for IMDS
+  // According to Azure docs, IMDS returns 410 for the first 70 seconds when not ready
+  imdsOptions.Retry.StatusCodes.insert(HttpStatusCode::Gone);
+  
+  return imdsOptions;
+}
 } // namespace
 
 Azure::Core::Url ManagedIdentitySource::ParseEndpointUrl(
@@ -539,7 +557,7 @@ ImdsManagedIdentitySource::ImdsManagedIdentitySource(
     std::string const& resourceId,
     Azure::Core::Url const& imdsUrl,
     Azure::Core::Credentials::TokenCredentialOptions const& options)
-    : ManagedIdentitySource(clientId, std::string(), options),
+    : ManagedIdentitySource(clientId, std::string(), CreateImdsRetryOptions(options)),
       m_request(Azure::Core::Http::HttpMethod::Get, imdsUrl)
 {
   {

sdk/identity/azure-identity/test/ut/managed_identity_credential_test.cpp

@@ -3193,4 +3193,215 @@ namespace Azure { namespace Identity { namespace Test {
     Logger::SetListener(nullptr);
   }
 
+  TEST(ManagedIdentityCredential, ImdsHttp410Retry)
+  {
+    // Test that IMDS properly retries HTTP 410 responses and eventually succeeds
+    using Azure::Core::Diagnostics::Logger;
+    using LogMsgVec = std::vector<std::pair<Logger::Level, std::string>>;
+    LogMsgVec log;
+    Logger::SetLevel(Logger::Level::Verbose);
+    Logger::SetListener([&](auto lvl, auto msg) { log.push_back(std::make_pair(lvl, msg)); });
+
+    try
+    {
+      auto const actual = CredentialTestHelper::SimulateTokenRequest(
+          [&](auto transport) {
+            TokenCredentialOptions options;
+            options.Transport.Transport = transport;
+
+            CredentialTestHelper::EnvironmentOverride const env({
+                {"MSI_ENDPOINT", ""},
+                {"MSI_SECRET", ""},
+                {"IDENTITY_ENDPOINT", ""},
+                {"IMDS_ENDPOINT", ""},
+                {"IDENTITY_HEADER", ""},
+                {"IDENTITY_SERVER_THUMBPRINT", ""},
+            });
+
+            return std::make_unique<ManagedIdentityCredential>(options);
+          },
+          {{"https://azure.com/.default"}},
+          std::vector<CredentialTestHelper::TokenRequestSimulationServerResponse>{
+              {HttpStatusCode::Gone, "{\"error\":\"not_ready\"}", {}},
+              {HttpStatusCode::Gone, "{\"error\":\"not_ready\"}", {}},
+              {HttpStatusCode::Ok, "{\"expires_in\":3600, \"access_token\":\"ACCESSTOKEN1\"}", {}}});
+
+      // Should make 3 requests: 2 HTTP 410 retries + 1 successful
+      EXPECT_EQ(actual.Requests.size(), 3U);
+      EXPECT_EQ(actual.Responses.size(), 1U);
+
+      // All requests should be to the IMDS endpoint
+      for (const auto& request : actual.Requests)
+      {
+        EXPECT_EQ(request.HttpMethod, HttpMethod::Get);
+        EXPECT_TRUE(request.AbsoluteUrl.find("169.254.169.254") != std::string::npos);
+        EXPECT_NE(request.Headers.find("Metadata"), request.Headers.end());
+        EXPECT_EQ(request.Headers.at("Metadata"), "true");
+      }
+
+      // Should get successful token response
+      EXPECT_EQ(actual.Responses.at(0).AccessToken.Token, "ACCESSTOKEN1");
+    }
+    catch (...)
+    {
+      Logger::SetListener(nullptr);
+      throw;
+    }
+
+    Logger::SetListener(nullptr);
+  }
+
+  TEST(ManagedIdentityCredential, ImdsRetryDuration)
+  {
+    // Test that IMDS retry policy includes HTTP 410 as retryable status code
+    // Note: This test validates HTTP 410 is retryable but doesn't test the full 70+ second
+    // requirement which would need extended retry duration (requires Azure Core support)
+    using Azure::Core::Diagnostics::Logger;
+    using LogMsgVec = std::vector<std::pair<Logger::Level, std::string>>;
+    LogMsgVec log;
+    Logger::SetLevel(Logger::Level::Verbose);
+    Logger::SetListener([&](auto lvl, auto msg) { log.push_back(std::make_pair(lvl, msg)); });
+
+    try
+    {
+      // Create 4 HTTP 410 responses (3 retries + initial attempt) followed by success  
+      std::vector<CredentialTestHelper::TokenRequestSimulationServerResponse> responses;
+      for (int i = 0; i < 4; ++i)
+      {
+        responses.push_back({HttpStatusCode::Gone, "{\"error\":\"not_ready\"}", {}});
+      }
+      responses.push_back({HttpStatusCode::Ok, "{\"expires_in\":3600, \"access_token\":\"ACCESSTOKEN1\"}", {}});
+
+      auto const actual = CredentialTestHelper::SimulateTokenRequest(
+          [&](auto transport) {
+            TokenCredentialOptions options;
+            options.Transport.Transport = transport;
+
+            CredentialTestHelper::EnvironmentOverride const env({
+                {"MSI_ENDPOINT", ""},
+                {"MSI_SECRET", ""},
+                {"IDENTITY_ENDPOINT", ""},
+                {"IMDS_ENDPOINT", ""},
+                {"IDENTITY_HEADER", ""},
+                {"IDENTITY_SERVER_THUMBPRINT", ""},
+            });
+
+            return std::make_unique<ManagedIdentityCredential>(options);
+          },
+          {{"https://azure.com/.default"}},
+          responses);
+
+      // Should make 7 requests (6 retries + initial) and then succeed on the 8th
+      EXPECT_EQ(actual.Requests.size(), 7U);
+      EXPECT_EQ(actual.Responses.size(), 1U);
+
+      // Verify we get a successful token response after all the retries
+      EXPECT_EQ(actual.Responses.at(0).AccessToken.Token, "ACCESSTOKEN1");
+    }
+    catch (...)
+    {
+      Logger::SetListener(nullptr);
+      throw;
+    }
+
+    Logger::SetListener(nullptr);
+  }
+
+  TEST(ManagedIdentityCredential, ImdsHttp410RetryExhaustion)
+  {
+    // Test that IMDS eventually fails when all retries return HTTP 410
+    using Azure::Core::Diagnostics::Logger;
+    using LogMsgVec = std::vector<std::pair<Logger::Level, std::string>>;
+    LogMsgVec log;
+    Logger::SetLevel(Logger::Level::Verbose);
+    Logger::SetListener([&](auto lvl, auto msg) { log.push_back(std::make_pair(lvl, msg)); });
+
+    try
+    {
+      // Create 5 HTTP 410 responses (more than the 3 max retries + initial attempt)
+      std::vector<CredentialTestHelper::TokenRequestSimulationServerResponse> responses;
+      for (int i = 0; i < 5; ++i)
+      {
+        responses.push_back({HttpStatusCode::Gone, "{\"error\":\"not_ready\"}", {}});
+      }
+
+      EXPECT_THROW(
+          CredentialTestHelper::SimulateTokenRequest(
+              [&](auto transport) {
+                TokenCredentialOptions options;
+                options.Transport.Transport = transport;
+
+                CredentialTestHelper::EnvironmentOverride const env({
+                    {"MSI_ENDPOINT", ""},
+                    {"MSI_SECRET", ""},
+                    {"IDENTITY_ENDPOINT", ""},
+                    {"IMDS_ENDPOINT", ""},
+                    {"IDENTITY_HEADER", ""},
+                    {"IDENTITY_SERVER_THUMBPRINT", ""},
+                });
+
+                return std::make_unique<ManagedIdentityCredential>(options);
+              },
+              {{"https://azure.com/.default"}},
+              responses),
+          Azure::Identity::AuthenticationException);
+    }
+    catch (...)
+    {
+      Logger::SetListener(nullptr);
+      throw;
+    }
+
+    Logger::SetListener(nullptr);
+  }
+
+  TEST(ManagedIdentityCredential, ImdsRetryOtherStatusCodes)
+  {
+    // Test that other retryable status codes still work correctly with IMDS
+    using Azure::Core::Diagnostics::Logger;
+    using LogMsgVec = std::vector<std::pair<Logger::Level, std::string>>;
+    LogMsgVec log;
+    Logger::SetLevel(Logger::Level::Verbose);
+    Logger::SetListener([&](auto lvl, auto msg) { log.push_back(std::make_pair(lvl, msg)); });
+
+    try
+    {
+      auto const actual = CredentialTestHelper::SimulateTokenRequest(
+          [&](auto transport) {
+            TokenCredentialOptions options;
+            options.Transport.Transport = transport;
+
+            CredentialTestHelper::EnvironmentOverride const env({
+                {"MSI_ENDPOINT", ""},
+                {"MSI_SECRET", ""},
+                {"IDENTITY_ENDPOINT", ""},
+                {"IMDS_ENDPOINT", ""},
+                {"IDENTITY_HEADER", ""},
+                {"IDENTITY_SERVER_THUMBPRINT", ""},
+            });
+
+            return std::make_unique<ManagedIdentityCredential>(options);
+          },
+          {{"https://azure.com/.default"}},
+          std::vector<CredentialTestHelper::TokenRequestSimulationServerResponse>{
+              {HttpStatusCode::TooManyRequests, "", {}},
+              {HttpStatusCode::InternalServerError, "", {}},
+              {HttpStatusCode::Ok, "{\"expires_in\":3600, \"access_token\":\"ACCESSTOKEN1\"}", {}}});
+
+      // Should make 3 requests: 2 retries + 1 successful
+      EXPECT_EQ(actual.Requests.size(), 3U);
+      EXPECT_EQ(actual.Responses.size(), 1U);
+
+      // Should get successful token response
+      EXPECT_EQ(actual.Responses.at(0).AccessToken.Token, "ACCESSTOKEN1");
+    }
+    catch (...)
+    {
+      Logger::SetListener(nullptr);
+      throw;
+    }
+
+    Logger::SetListener(nullptr);
+  }
+
 }}} // namespace Azure::Identity::Test