diff --git a/src/client/Microsoft.Identity.Client/Http/HttpManager.cs b/src/client/Microsoft.Identity.Client/Http/HttpManager.cs index 8cf5939468..268cc09f32 100644 --- a/src/client/Microsoft.Identity.Client/Http/HttpManager.cs +++ b/src/client/Microsoft.Identity.Client/Http/HttpManager.cs @@ -140,17 +140,24 @@ public async Task SendRequestAsync( if (headers != null && headers.Count > 0) { - var correlationId = headers[OAuth2Header.CorrelationId]; - string correlationIdMsg = headers.ContainsKey(OAuth2Header.CorrelationId) ? - $" CorrelationId: {correlationId}" : - string.Empty; + string correlationIdMsg = string.Empty; + string correlationId = null; + + if (headers.ContainsKey(OAuth2Header.CorrelationId)) + { + correlationId = headers[OAuth2Header.CorrelationId]; + correlationIdMsg = $" CorrelationId: {correlationId}"; + } var ex = new MsalServiceException( MsalError.RequestTimeout, msg + correlationIdMsg, timeoutException); - ex.CorrelationId = correlationId; + if (correlationId != null) + { + ex.CorrelationId = correlationId; + } throw ex; } diff --git a/tests/Microsoft.Identity.Test.Unit/CoreTests/HttpTests/HttpManagerTests.cs b/tests/Microsoft.Identity.Test.Unit/CoreTests/HttpTests/HttpManagerTests.cs index 0c72b0d5f3..64ba873a4a 100644 --- a/tests/Microsoft.Identity.Test.Unit/CoreTests/HttpTests/HttpManagerTests.cs +++ b/tests/Microsoft.Identity.Test.Unit/CoreTests/HttpTests/HttpManagerTests.cs @@ -581,6 +581,47 @@ public async Task TestCorrelationIdWithRetryOnTimeoutFailureAsync(bool addCorrel } } + // Regression test for https://github.com/AzureAD/microsoft-authentication-library-for-dotnet/issues/XXXX + // Verifies that timeout with headers but without correlation ID does not throw KeyNotFoundException + [TestMethod] + public async Task TestRetryOnTimeoutWithHeadersButNoCorrelationIdAsync() + { + using (var httpManager = new MockHttpManager()) + { + // Simulate permanent errors (to trigger the maximum number of retries) + const int Num500Errors = 1 + TestDefaultRetryPolicy.DefaultStsMaxRetries; // initial request + maximum number of retries + for (int i = 0; i < Num500Errors; i++) + { + httpManager.AddRequestTimeoutResponseMessageMockHandler(HttpMethod.Post); + } + + // Create headers without correlation ID - this was causing KeyNotFoundException + var headers = new Dictionary + { + ["some-other-header"] = "some-value" + }; + + var exc = await AssertException.TaskThrowsAsync(() => + httpManager.SendRequestAsync( + new Uri(TestConstants.AuthorityHomeTenant + "oauth2/token"), + headers: headers, + body: new FormUrlEncodedContent(new Dictionary()), + method: HttpMethod.Post, + logger: Substitute.For(), + doNotThrow: false, + mtlsCertificate: null, + validateServerCert: null, + cancellationToken: default, + retryPolicy: _stsRetryPolicy)) + .ConfigureAwait(false); + + // Should get timeout error without KeyNotFoundException + Assert.AreEqual(MsalError.RequestTimeout, exc.ErrorCode); + Assert.AreEqual("Request to the endpoint timed out.", exc.Message); + Assert.IsNull(exc.CorrelationId); + } + } + [TestMethod] public async Task TestWithCorrelationId_RetryOnTimeoutFailureAsync() {