From 2a0dab9a2e030045056207aaea174a80f36ea583 Mon Sep 17 00:00:00 2001
From: Engincan VESKE <43685404+EngincanV@users.noreply.github.com>
Date: Thu, 20 Apr 2023 10:07:27 +0300
Subject: [PATCH] Add custom security header support to security header
 middleware

---
 .../Security/AbpSecurityHeadersMiddleware.cs  | 21 ++++++++++++++-----
 .../Security/AbpSecurityHeadersOptions.cs     |  9 ++++++++
 .../SecurityHeadersTestController_Tests.cs    |  9 ++++++++
 3 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/framework/src/Volo.Abp.AspNetCore/Volo/Abp/AspNetCore/Security/AbpSecurityHeadersMiddleware.cs b/framework/src/Volo.Abp.AspNetCore/Volo/Abp/AspNetCore/Security/AbpSecurityHeadersMiddleware.cs
index 37c18f6b70..29e05b9bdc 100644
--- a/framework/src/Volo.Abp.AspNetCore/Volo/Abp/AspNetCore/Security/AbpSecurityHeadersMiddleware.cs
+++ b/framework/src/Volo.Abp.AspNetCore/Volo/Abp/AspNetCore/Security/AbpSecurityHeadersMiddleware.cs
@@ -20,27 +20,38 @@ public class AbpSecurityHeadersMiddleware : IMiddleware, ITransientDependency
     public async Task InvokeAsync(HttpContext context, RequestDelegate next)
     {
         /*X-Content-Type-Options header tells the browser to not try and “guess” what a mimetype of a resource might be, and to just take what mimetype the server has returned as fact.*/
-        AddHeaderIfNotExists(context, "X-Content-Type-Options", "nosniff");
+        AddHeader(context, "X-Content-Type-Options", "nosniff");
 
         /*X-XSS-Protection is a feature of Internet Explorer, Chrome and Safari that stops pages from loading when they detect reflected cross-site scripting (XSS) attacks*/
-        AddHeaderIfNotExists(context, "X-XSS-Protection", "1; mode=block");
+        AddHeader(context, "X-XSS-Protection", "1; mode=block");
 
         /*The X-Frame-Options HTTP response header can be used to indicate whether or not a browser should be allowed to render a page in a <frame>, <iframe> or <object>. SAMEORIGIN makes it being displayed in a frame on the same origin as the page itself. The spec leaves it up to browser vendors to decide whether this option applies to the top level, the parent, or the whole chain*/
-        AddHeaderIfNotExists(context, "X-Frame-Options", "SAMEORIGIN");
+        AddHeader(context, "X-Frame-Options", "SAMEORIGIN");
 
         if (Options.Value.UseContentSecurityPolicyHeader)
         {
-            AddHeaderIfNotExists(context, "Content-Security-Policy",
+            AddHeader(context, "Content-Security-Policy",
                 Options.Value.ContentSecurityPolicyValue.IsNullOrEmpty()
                     ? "object-src 'none'; form-action 'self'; frame-ancestors 'none'"
                     : Options.Value.ContentSecurityPolicyValue);
         }
+        
+        foreach (var (key, value) in Options.Value.Headers)
+        {
+            AddHeader(context, key, value, true);
+        }
 
         await next.Invoke(context);
     }
 
-    protected virtual void AddHeaderIfNotExists(HttpContext context, string key, string value)
+    protected virtual void AddHeader(HttpContext context, string key, string value, bool overrideIfExists = false)
     {
+        if (overrideIfExists && context.Response.Headers.TryGetValue(key, out _))
+        {
+            context.Response.Headers[key] = value;
+            return;
+        }
+        
         context.Response.Headers.AddIfNotContains(new KeyValuePair<string, StringValues>(key, value));
     }
 }
diff --git a/framework/src/Volo.Abp.AspNetCore/Volo/Abp/AspNetCore/Security/AbpSecurityHeadersOptions.cs b/framework/src/Volo.Abp.AspNetCore/Volo/Abp/AspNetCore/Security/AbpSecurityHeadersOptions.cs
index 198a62723c..93b0dc9236 100644
--- a/framework/src/Volo.Abp.AspNetCore/Volo/Abp/AspNetCore/Security/AbpSecurityHeadersOptions.cs
+++ b/framework/src/Volo.Abp.AspNetCore/Volo/Abp/AspNetCore/Security/AbpSecurityHeadersOptions.cs
@@ -1,3 +1,5 @@
+using System.Collections.Generic;
+
 namespace Volo.Abp.AspNetCore.Security;
 
 public class AbpSecurityHeadersOptions
@@ -5,4 +7,11 @@ public class AbpSecurityHeadersOptions
     public bool UseContentSecurityPolicyHeader { get; set; }
 
     public string ContentSecurityPolicyValue { get; set; }
+
+    public Dictionary<string, string> Headers { get; }
+
+    public AbpSecurityHeadersOptions()
+    {
+        Headers = new Dictionary<string, string>();
+    }
 }
diff --git a/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Headers/SecurityHeadersTestController_Tests.cs b/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Headers/SecurityHeadersTestController_Tests.cs
index fdfda83ac9..7c704905e9 100644
--- a/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Headers/SecurityHeadersTestController_Tests.cs
+++ b/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Headers/SecurityHeadersTestController_Tests.cs
@@ -15,6 +15,7 @@ public class SecurityHeadersTestController_Tests : AspNetCoreMvcTestBase
         services.Configure<AbpSecurityHeadersOptions>(options =>
         {
             options.UseContentSecurityPolicyHeader = true;
+            options.Headers["Referrer-Policy"] = "no-referrer";
         });
 
         base.ConfigureServices(context, services);
@@ -30,4 +31,12 @@ public class SecurityHeadersTestController_Tests : AspNetCoreMvcTestBase
         responseMessage.Headers.ShouldContain(x => x.Key == "X-Content-Type-Options" & x.Value.First().ToString() == "nosniff");
         responseMessage.Headers.ShouldContain(x => x.Key == "Content-Security-Policy" & x.Value.First().ToString() == "object-src 'none'; form-action 'self'; frame-ancestors 'none'");
     }
+
+    [Fact]
+    public async Task SecurityHeaders_Custom_Headers_Should_Be_Added()
+    {
+        var responseMessage = await GetResponseAsync("/SecurityHeadersTest/Get");
+        responseMessage.Headers.ShouldNotBeEmpty();
+        responseMessage.Headers.ShouldContain(x => x.Key == "Referrer-Policy" && x.Value.First().ToString() == "no-referrer");
+    }
 }