diff --git a/src/Stott.Security.Optimizely/Features/CustomHeaders/Service/CustomHeaderService.cs b/src/Stott.Security.Optimizely/Features/CustomHeaders/Service/CustomHeaderService.cs index 61b7f2fb..6251f328 100644 --- a/src/Stott.Security.Optimizely/Features/CustomHeaders/Service/CustomHeaderService.cs +++ b/src/Stott.Security.Optimizely/Features/CustomHeaders/Service/CustomHeaderService.cs @@ -43,7 +43,8 @@ public async Task> 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(); } diff --git a/src/Stott.Security.Optimizely/Features/Header/HeaderCompilationService.cs b/src/Stott.Security.Optimizely/Features/Header/HeaderCompilationService.cs index 654a749e..050e2d0c 100644 --- a/src/Stott.Security.Optimizely/Features/Header/HeaderCompilationService.cs +++ b/src/Stott.Security.Optimizely/Features/Header/HeaderCompilationService.cs @@ -104,26 +104,49 @@ private async Task> ModifyHeadersForRequest(List 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 + }; + } } diff --git a/src/Stott.Security.Optimizely/Features/Header/HeaderDto.cs b/src/Stott.Security.Optimizely/Features/Header/HeaderDto.cs index e98f2361..e7f3435d 100644 --- a/src/Stott.Security.Optimizely/Features/Header/HeaderDto.cs +++ b/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). /// public bool IsRemoval { get; set; } + + /// + /// Gets or sets a value indicating whether this header should replace existing values. + /// + public bool IsReplacement { get; set; } } \ No newline at end of file diff --git a/src/Stott.Security.Optimizely/Features/Middleware/SecurityHeaderMiddleware.cs b/src/Stott.Security.Optimizely/Features/Middleware/SecurityHeaderMiddleware.cs index aed3a2d2..25edc929 100644 --- a/src/Stott.Security.Optimizely/Features/Middleware/SecurityHeaderMiddleware.cs +++ b/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); + } }