Ignore response transform OCEs that occur during error handling (dotnet#2452)

Author: MihaZupan

src/ReverseProxy/Forwarder/HttpForwarder.cs

@@ -5,7 +5,6 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
-using System.Linq;
 using System.Net;
 using System.Net.Http;
 using System.Threading;
@@ -15,7 +14,6 @@
 #if NET8_0_OR_GREATER
 using Microsoft.AspNetCore.Http.Timeouts;
 #endif
-using Microsoft.AspNetCore.Server.Kestrel.Core.Features;
 using Microsoft.AspNetCore.WebUtilities;
 using Microsoft.Extensions.Logging;
 using Microsoft.Net.Http.Headers;
@@ -643,7 +641,16 @@ private async ValueTask<ForwarderError> HandleRequestFailureAsync(HttpContext co
             {
                 var error = HandleRequestBodyFailure(context, requestBodyCopyResult, requestBodyException!, requestException,
                     timedOut: requestCancellationSource.IsCancellationRequested);
-                await transformer.TransformResponseAsync(context, proxyResponse: null, requestCancellationSource.Token);
+
+                try
+                {
+                    await transformer.TransformResponseAsync(context, proxyResponse: null, requestCancellationSource.Token);
+                }
+                catch (OperationCanceledException)
+                {
+                    // We're about to report a more specific error, so ignore OCEs that occur here.
+                }
+
                 return error;
             }
         }
@@ -669,7 +676,16 @@ async ValueTask<ForwarderError> ReportErrorAsync(ForwarderError error, int statu
                 await requestContent.ConsumptionTask;
             }
 
-            await transformer.TransformResponseAsync(context, null, requestCancellationSource.Token);
+            try
+            {
+                await transformer.TransformResponseAsync(context, proxyResponse: null, requestCancellationSource.Token);
+            }
+            catch (OperationCanceledException)
+            {
+                // We may have manually cancelled the request CTS as part of error handling.
+                // We're about to report a more specific error, so ignore OCEs that occur here.
+            }
+
             return error;
         }
     }

src/ReverseProxy/SessionAffinity/BaseEncryptedSessionAffinityPolicy.cs

@@ -26,13 +26,19 @@ protected BaseEncryptedSessionAffinityPolicy(IDataProtectionProvider dataProtect
 
     public abstract string Name { get; }
 
-    public virtual void AffinitizeResponse(HttpContext context, ClusterState cluster, SessionAffinityConfig config, DestinationState destination)
+    public void AffinitizeResponse(HttpContext context, ClusterState cluster, SessionAffinityConfig config, DestinationState destination)
     {
         if (!config.Enabled.GetValueOrDefault())
         {
             throw new InvalidOperationException($"Session affinity is disabled for cluster.");
         }
 
+        if (context.RequestAborted.IsCancellationRequested)
+        {
+            // Avoid wasting time if the client is already gone.
+            return;
+        }
+
         // Affinity key is set on the response only if it's a new affinity.
         if (!context.Items.ContainsKey(AffinityKeyId))
         {

src/ReverseProxy/SessionAffinity/BaseHashCookieSessionAffinityPolicy.cs

@@ -31,6 +31,12 @@ public void AffinitizeResponse(HttpContext context, ClusterState cluster, Sessio
             throw new InvalidOperationException("Session affinity is disabled for cluster.");
         }
 
+        if (context.RequestAborted.IsCancellationRequested)
+        {
+            // Avoid wasting time if the client is already gone.
+            return;
+        }
+
         // Affinity key is set on the response only if it's a new affinity.
         if (!context.Items.ContainsKey(AffinityKeyId))
         {

src/ReverseProxy/SessionAffinity/ISessionAffinityPolicy.cs

@@ -65,8 +65,6 @@ ValueTask<AffinityResult> FindAffinitizedDestinationsAsync(HttpContext context,
     /// <param name="cancellationToken">The token to monitor for cancellation requests.</param>
     ValueTask AffinitizeResponseAsync(HttpContext context, ClusterState cluster, SessionAffinityConfig config, DestinationState destination, CancellationToken cancellationToken)
     {
-        cancellationToken.ThrowIfCancellationRequested();
-
         AffinitizeResponse(context, cluster, config, destination);
         return default;
     }

test/ReverseProxy.Tests/Forwarder/HttpForwarderTests.cs

@@ -2696,6 +2696,77 @@ public async Task RequestFailure_ResponseTransformsAreCalled(bool failureInReque
         events.AssertContainProxyStages(failureInRequestTransform ? Array.Empty<ForwarderStage>() : new [] { ForwarderStage.SendAsyncStart });
     }
 
+    [Theory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public async Task RequestFailure_CancellationExceptionInResponseTransformIsIgnored(bool throwOce)
+    {
+        var events = TestEventListener.Collect();
+
+        var httpContext = new DefaultHttpContext();
+        httpContext.Request.Method = "GET";
+        httpContext.Request.Host = new HostString("example.com:3456");
+
+        var destinationPrefix = "https://localhost:123/";
+        var sut = CreateProxy();
+        var client = MockHttpHandler.CreateClient(
+            (HttpRequestMessage request, CancellationToken cancellationToken) =>
+            {
+                throw new Exception();
+            });
+
+        var responseTransformWithNullResponseCalled = false;
+
+        var transformer = new DelegateHttpTransforms
+        {
+            OnResponse = (context, response) =>
+            {
+                if (response is null)
+                {
+                    responseTransformWithNullResponseCalled = true;
+
+                    throw throwOce
+                        ? new OperationCanceledException("Foo")
+                        : new InvalidOperationException("Bar");
+                }
+
+                return new ValueTask<bool>(true);
+            }
+        };
+
+        var proxyError = ForwarderError.None;
+        Exception exceptionThrownBySendAsync = null;
+
+        try
+        {
+            proxyError = await sut.SendAsync(httpContext, destinationPrefix, client, ForwarderRequestConfig.Empty, transformer);
+        }
+        catch (Exception ex)
+        {
+            exceptionThrownBySendAsync = ex;
+        }
+
+        Assert.True(responseTransformWithNullResponseCalled);
+
+        if (throwOce)
+        {
+            Assert.Null(exceptionThrownBySendAsync);
+            Assert.Equal(ForwarderError.Request, proxyError);
+        }
+        else
+        {
+            Assert.NotNull(exceptionThrownBySendAsync);
+        }
+
+        Assert.Equal(StatusCodes.Status502BadGateway, httpContext.Response.StatusCode);
+        var errorFeature = httpContext.Features.Get<IForwarderErrorFeature>();
+        Assert.Equal(ForwarderError.Request, errorFeature.Error);
+        Assert.IsType<Exception>(errorFeature.Exception);
+
+        AssertProxyStartFailedStop(events, destinationPrefix, httpContext.Response.StatusCode, errorFeature.Error);
+        events.AssertContainProxyStages([ForwarderStage.SendAsyncStart]);
+    }
+
     public enum CancellationScenario
     {
         RequestAborted,