Feature/header behaviour

src/Stott.Security.Optimizely/Features/CustomHeaders/Service/CustomHeaderService.cs

@@ -43,7 +43,8 @@ public async Task<IList<HeaderDto>> GetCompiledHeaders(string? appId, string? ho
                 {
                     Key = header.HeaderName,
                     Value = header.HeaderValue ?? string.Empty,
-                    IsRemoval = header.Behavior == CustomHeaderBehavior.Remove
+                    IsRemoval = header.Behavior == CustomHeaderBehavior.Remove,
+                    IsReplacement = header.Behavior == CustomHeaderBehavior.Add
                 }).ToList();
     }
 

src/Stott.Security.Optimizely/Features/Header/HeaderCompilationService.cs

@@ -104,26 +104,49 @@ private async Task<List<HeaderDto>> ModifyHeadersForRequest(List<HeaderDto> head
                     newValue = newValue.Replace(CspConstants.Sources.Nonce, nonceValue);
                 }
 
-                updatedHeaders.Add(new HeaderDto { Key = header.Key, Value = newValue, IsRemoval = header.IsRemoval });
+                updatedHeaders.Add(Clone(header, newValue));
             }
             else if (header.Key == CspConstants.HeaderNames.ReportingEndpoints)
             {
-                updatedHeaders.Add(new HeaderDto { Key = header.Key, Value = header.Value?.Replace(CspConstants.InternalReportingPlaceholder, cspReportUrlResolver.GetReportToPath()), IsRemoval = header.IsRemoval });
+                var reportValue = header.Value?.Replace(CspConstants.InternalReportingPlaceholder, cspReportUrlResolver.GetReportToPath());
+                updatedHeaders.Add(Clone(header, reportValue));
             }
             else if (header.Key == CspConstants.HeaderNames.StrictTransportSecurity)
             {
                 // HSTS should only be sent over HTTPS
                 if (isHttps)
                 {
-                    updatedHeaders.Add(new HeaderDto { Key = header.Key, Value = header.Value, IsRemoval = header.IsRemoval });
+                    updatedHeaders.Add(Clone(header));
                 }
             }
             else
             {
-                updatedHeaders.Add(new HeaderDto { Key = header.Key, Value = header.Value, IsRemoval = header.IsRemoval });
+                updatedHeaders.Add(Clone(header));
             }
         }
 
         return updatedHeaders;
     }
+
+    private static HeaderDto Clone(HeaderDto header)
+    {
+        return new HeaderDto
+        {
+            Key = header.Key,
+            Value = header.Value,   
+            IsRemoval = header.IsRemoval,
+            IsReplacement = header.IsReplacement
+        };
+    }
+
+    private static HeaderDto Clone(HeaderDto header, string? newValue)
+    {
+        return new HeaderDto
+        {
+            Key = header.Key,
+            Value = newValue,
+            IsRemoval = header.IsRemoval,
+            IsReplacement = header.IsReplacement
+        };
+    }
 }

src/Stott.Security.Optimizely/Features/Header/HeaderDto.cs

@@ -17,4 +17,9 @@ public sealed class HeaderDto
     /// When true, the middleware will call context.Response.Headers.Remove(Key).
     /// </summary>
     public bool IsRemoval { get; set; }
+
+    /// <summary>
+    /// Gets or sets a value indicating whether this header should replace existing values.
+    /// </summary>
+    public bool IsReplacement { get; set; }
 }

src/Stott.Security.Optimizely/Features/Middleware/SecurityHeaderMiddleware.cs

@@ -35,18 +35,22 @@ public async Task Invoke(
             var headers = await securityHeaderService.GetSecurityHeadersAsync(routeData, context.Request);
             foreach (var header in headers)
             {
-                if (string.IsNullOrWhiteSpace(header.Key))
+                if (header is null)
                 {
                     continue;
                 }
 
                 if (header.IsRemoval)
                 {
-                    context.Response.Headers.Remove(header.Key);
+                    HandleRemoval(context, header);
                 }
-                else if (!string.IsNullOrWhiteSpace(header.Value))
+                else if (header.IsReplacement)
                 {
-                    context.Response.Headers.Append(header.Key, header.Value);
+                    HandleReplacement(context, header);
+                }
+                else 
+                {
+                    HandleAppend(context, header);
                 }
             }
         }
@@ -57,4 +61,41 @@ public async Task Invoke(
 
         await _next(context);
     }
+
+    private static void HandleAppend(HttpContext context, HeaderDto header)
+    {
+        if (string.IsNullOrWhiteSpace(header?.Key) || string.IsNullOrWhiteSpace(header?.Value))
+        {
+            return;
+        }
+
+        context.Response.Headers.Append(header.Key, header.Value);
+    }
+
+    private static void HandleReplacement(HttpContext context, HeaderDto header)
+    {
+        if (string.IsNullOrWhiteSpace(header?.Key) || string.IsNullOrWhiteSpace(header?.Value))
+        { 
+            return;
+        }
+
+        if (context.Response.Headers.ContainsKey(header.Key))
+        {
+            context.Response.Headers[header.Key] = header.Value;
+        }
+        else
+        {   
+            context.Response.Headers.Append(header.Key, header.Value);
+        }
+    }
+
+    private static void HandleRemoval(HttpContext context, HeaderDto header)
+    {
+        if (string.IsNullOrWhiteSpace(header?.Key))
+        {
+            return;
+        }
+
+        context.Response.Headers.Remove(header.Key);
+    }
 }