update controller to middleware

Azure · Jan 26, 2024 · ae1cc9c · ae1cc9c
1 parent 165501d
commit ae1cc9c
Showing 1 changed file with 16 additions and 21 deletions.
diff --git a/...FaultInjector/FaultInjectingController.cs → ...FaultInjector/FaultInjectingMiddleware.cs b/...FaultInjector/FaultInjectingController.cs → ...FaultInjector/FaultInjectingMiddleware.cs
@@ -20,23 +20,20 @@ public class FaultInjectingMiddleware
     {
         private readonly ILogger<FaultInjectingMiddleware> _logger;
         private readonly HttpClient _httpClient;
-        private readonly RequestDelegate _next;
-        public FaultInjectingMiddleware(RequestDelegate next, IHttpClientFactory httpClientFactory, ILogger<FaultInjectingMiddleware> logger)
+        public FaultInjectingMiddleware(RequestDelegate _, IHttpClientFactory httpClientFactory, ILogger<FaultInjectingMiddleware> logger)
         {
             _logger = logger;
             _httpClient = httpClientFactory.CreateClient("upstream");
-            _next = next;
         }
 
-        public async Task InvokeAsync(HttpContext context, 
-            CancellationToken cancellationToken)
+        public async Task InvokeAsync(HttpContext context)
         {
             string faultHeaderValue = context.Request.Headers[Utils.ResponseSelectionHeader];
             string upstreamBaseUri = context.Request.Headers[Utils.UpstreamBaseUriHeader];
 
             if (ValidateOrReadFaultMode(faultHeaderValue, out var fault))
             {
-                await ProxyResponse(context, upstreamBaseUri, fault, cancellationToken);
+                await ProxyResponse(context, upstreamBaseUri, fault, context.RequestAborted);
             }
             else
             {
@@ -87,21 +84,18 @@ private async Task<UpstreamResponse> SendUpstreamRequest(HttpRequest request, st
                     }
                 }
 
-                using (var upstreamResponseMessage = await _httpClient.SendAsync(upstreamRequest))
-                {
-                    var headers = new List<KeyValuePair<string, IEnumerable<string>>>();
-                    // Must skip "Transfer-Encoding" header, since if it's set manually Kestrel requires you to implement
-                    // your own chunking.
-                    headers.AddRange(upstreamResponseMessage.Headers.Where(header => !string.Equals(header.Key, "Transfer-Encoding", StringComparison.OrdinalIgnoreCase)));
-                    headers.AddRange(upstreamResponseMessage.Content.Headers);
-
-                    var upstreamResponse = await UpstreamResponseFromHttpContent(upstreamResponseMessage.Content, cancellationToken);
-                    upstreamResponse.StatusCode = (int)upstreamResponseMessage.StatusCode;
-                    upstreamResponse.Headers = headers.Select(h => new KeyValuePair<string, StringValues>(h.Key, h.Value.ToArray()));
-                    _logger.LogInformation("Finished reading response body ({length})", upstreamResponse.ContentLength);
-
-                    return upstreamResponse;
-                }
+                var upstreamResponseMessage = await _httpClient.SendAsync(upstreamRequest);
+                var headers = new List<KeyValuePair<string, IEnumerable<string>>>();
+                // Must skip "Transfer-Encoding" header, since if it's set manually Kestrel requires you to implement
+                // your own chunking.
+                headers.AddRange(upstreamResponseMessage.Headers.Where(header => !string.Equals(header.Key, "Transfer-Encoding", StringComparison.OrdinalIgnoreCase)));
+                headers.AddRange(upstreamResponseMessage.Content.Headers);
+
+                var upstreamResponse = await UpstreamResponseFromHttpContent(upstreamResponseMessage.Content, cancellationToken);
+                upstreamResponse.StatusCode = (int)upstreamResponseMessage.StatusCode;
+                upstreamResponse.Headers = headers.Select(h => new KeyValuePair<string, StringValues>(h.Key, h.Value.ToArray()));
+
+                return upstreamResponse;
             }
         }
 
@@ -125,6 +119,7 @@ private async Task<MemoryStream> BufferContentAsync(HttpContent content, Cancell
 
             _logger.LogWarning("Response does not have content length (is chunked or malformed) and is being buffered");
             byte[] contentBytes = await content.ReadAsByteArrayAsync(cancellationToken);
+            _logger.LogInformation("Finished buffering response body ({length})", contentBytes.Length);
             return new MemoryStream(contentBytes);
         }