diff --git a/src/client/Microsoft.Identity.Client/Http/HttpManager.cs b/src/client/Microsoft.Identity.Client/Http/HttpManager.cs index 8cf5939468..805786df54 100644 --- a/src/client/Microsoft.Identity.Client/Http/HttpManager.cs +++ b/src/client/Microsoft.Identity.Client/Http/HttpManager.cs @@ -138,19 +138,16 @@ public async Task SendRequestAsync( //If the correlation id is available, include it in the exception message string msg = MsalErrorMessage.RequestTimeOut; - if (headers != null && headers.Count > 0) + if (headers != null && headers.Count > 0 && + headers.TryGetValue(OAuth2Header.CorrelationId, out var correlationId)) { - var correlationId = headers[OAuth2Header.CorrelationId]; - string correlationIdMsg = headers.ContainsKey(OAuth2Header.CorrelationId) ? - $" CorrelationId: {correlationId}" : - string.Empty; - var ex = new MsalServiceException( MsalError.RequestTimeout, - msg + correlationIdMsg, - timeoutException); - - ex.CorrelationId = correlationId; + msg + $" CorrelationId: {correlationId}", + timeoutException) + { + 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() {