Merge pull request #4 from stevejgordon/ResponseHeaderBug

Check if response header exists before trying to add it
stevejgordon · Oct 31, 2017 · 442744a · 442744a
2 parents 015ace9 + 791583a
commit 442744a
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 9 deletions.
diff --git a/src/CorrelationId/CorrelationIdMiddleware.cs b/src/CorrelationId/CorrelationIdMiddleware.cs
@@ -18,14 +18,8 @@ public class CorrelationIdMiddleware
         /// <param name="options">The configuration options.</param>
         public CorrelationIdMiddleware(RequestDelegate next, IOptions<CorrelationIdOptions> options)
         {
-            if (options == null)
-            {
-                throw new ArgumentNullException(nameof(options));
-            }
-
             _next = next ?? throw new ArgumentNullException(nameof(next));
-
-            _options = options.Value;
+            _options = options?.Value ?? throw new ArgumentNullException(nameof(options));
         }
 
         /// <summary>
@@ -45,7 +39,11 @@ public Task Invoke(HttpContext context)
                 // apply the correlation ID to the response header for client side tracking
                 context.Response.OnStarting(() =>
                 {
-                    context.Response.Headers.Add(_options.Header, new[] { context.TraceIdentifier });
+                    if (!context.Response.Headers.ContainsKey(_options.Header))
+                    {
+                        context.Response.Headers.Add(_options.Header, new[] { context.TraceIdentifier });
+                    }
+
                     return Task.CompletedTask;
                 });
             }

diff --git a/test/CorrelationId.Tests/CorrelationIdMiddlewareTests.cs b/test/CorrelationId.Tests/CorrelationIdMiddlewareTests.cs
@@ -1,3 +1,4 @@
+using System;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.TestHost;
 using System.Collections.Generic;
@@ -26,6 +27,23 @@ public async Task ReturnsCorrelationIdInResponseHeader_WhenOptionSetToTrue()
             Assert.NotNull(header);            
         }
 
+        [Fact]
+        public async Task DoesNotThrowException_WhenOptionSetToTrue_IfHeaderIsAlreadySet()
+        {
+            var builder = new WebHostBuilder()
+                .Configure(app =>
+                {
+                    app.UseCorrelationId();
+                    app.UseCorrelationId(); // header will already be set on this second use of the middleware
+                });
+
+            var server = new TestServer(builder);
+
+            await server.CreateClient().GetAsync("");
+
+            // Note: This test will only fail if the middleware is causing an exception by trying to set a response header which already exists.
+        }
+
         [Fact]
         public async Task DoesNotReturnCorrelationIdInResponseHeader_WhenOptionSetToFalse()
         {
@@ -38,7 +56,7 @@ public async Task DoesNotReturnCorrelationIdInResponseHeader_WhenOptionSetToFals
 
             var response = await server.CreateClient().GetAsync("");
 
-            var headerExists = response.Headers.TryGetValues(options.Header, out IEnumerable<string> headerValue);
+            var headerExists = response.Headers.TryGetValues(options.Header, out IEnumerable<string> _);
 
             Assert.False(headerExists);
         }