From 32d72195806b807d15cb202b83b9b4258b00f8f3 Mon Sep 17 00:00:00 2001
From: maliming <malimings@gmail.com>
Date: Wed, 17 Feb 2021 09:32:03 +0800
Subject: [PATCH] Add unit test for AbpSecurityHeadersMiddleware.

---
 .../AbpApplicationBuilderExtensions.cs        |  4 +--
 .../Volo.Abp.AspNetCore.csproj                |  1 -
 .../Security/AbpSecurityHeadersMiddleware.cs  | 31 +++++++------------
 .../Authorization/AuthTestController_Tests.cs | 10 ------
 .../Claims/ClaimsMapTestController_Tests.cs   | 10 ------
 .../Headers/SecurityHeadersTestController.cs  | 12 +++++++
 .../SecurityHeadersTestController_Tests.cs    | 19 ++++++++++++
 7 files changed, 45 insertions(+), 42 deletions(-)
 create mode 100644 framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Headers/SecurityHeadersTestController.cs
 create mode 100644 framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Headers/SecurityHeadersTestController_Tests.cs

diff --git a/framework/src/Volo.Abp.AspNetCore/Microsoft/AspNetCore/Builder/AbpApplicationBuilderExtensions.cs b/framework/src/Volo.Abp.AspNetCore/Microsoft/AspNetCore/Builder/AbpApplicationBuilderExtensions.cs
index fab4a4bb80..a409355ef0 100644
--- a/framework/src/Volo.Abp.AspNetCore/Microsoft/AspNetCore/Builder/AbpApplicationBuilderExtensions.cs
+++ b/framework/src/Volo.Abp.AspNetCore/Microsoft/AspNetCore/Builder/AbpApplicationBuilderExtensions.cs
@@ -84,9 +84,9 @@ namespace Microsoft.AspNetCore.Builder
             return app.UseMiddleware<AbpClaimsMapMiddleware>();
         }
 
-        public static void UseAbpSecurityHeaders(this IApplicationBuilder app)
+        public static IApplicationBuilder UseAbpSecurityHeaders(this IApplicationBuilder app)
         {
-            app.UseMiddleware<AbpSecurityHeadersMiddleware>();
+            return app.UseMiddleware<AbpSecurityHeadersMiddleware>();
         }
     }
 }
diff --git a/framework/src/Volo.Abp.AspNetCore/Volo.Abp.AspNetCore.csproj b/framework/src/Volo.Abp.AspNetCore/Volo.Abp.AspNetCore.csproj
index eaa053c91a..b191872339 100644
--- a/framework/src/Volo.Abp.AspNetCore/Volo.Abp.AspNetCore.csproj
+++ b/framework/src/Volo.Abp.AspNetCore/Volo.Abp.AspNetCore.csproj
@@ -25,7 +25,6 @@
     <ProjectReference Include="..\Volo.Abp.Uow\Volo.Abp.Uow.csproj" />
     <ProjectReference Include="..\Volo.Abp.Validation\Volo.Abp.Validation.csproj" />
     <ProjectReference Include="..\Volo.Abp.VirtualFileSystem\Volo.Abp.VirtualFileSystem.csproj" />
-    <PackageReference Include="Microsoft.AspNetCore.Http.Abstractions" Version="2.0.0" />
   </ItemGroup>
 
 </Project>
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 09dddfabe7..28b0f3a48e 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
@@ -1,37 +1,30 @@
+using System.Collections.Generic;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Primitives;
+using Volo.Abp.DependencyInjection;
 
 namespace Volo.Abp.AspNetCore.Security
 {
-    public class AbpSecurityHeadersMiddleware
+    public class AbpSecurityHeadersMiddleware : IMiddleware, ITransientDependency
     {
-        private readonly RequestDelegate _next;
-
-        public AbpSecurityHeadersMiddleware(RequestDelegate next)
-        {
-            _next = next;
-        }
-
-        public async Task Invoke(HttpContext httpContext)
+        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(httpContext, "X-Content-Type-Options", "nosniff");
+            AddHeaderIfNotExists(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(httpContext, "X-XSS-Protection", "1; mode=block");
+            AddHeaderIfNotExists(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(httpContext, "X-Frame-Options", "SAMEORIGIN");
+            AddHeaderIfNotExists(context, "X-Frame-Options", "SAMEORIGIN");
 
-            await _next.Invoke(httpContext);
+            await next.Invoke(context);
         }
 
-        private static void AddHeaderIfNotExists(HttpContext context, string key, string value)
+        protected virtual void AddHeaderIfNotExists(HttpContext context, string key, string value)
         {
-            if (!context.Response.Headers.ContainsKey(key))
-            {
-                context.Response.Headers.Add(key, value);
-            }
+            context.Response.Headers.AddIfNotContains(new KeyValuePair<string, StringValues>(key, value));
         }
     }
-}
\ No newline at end of file
+}
diff --git a/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Authorization/AuthTestController_Tests.cs b/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Authorization/AuthTestController_Tests.cs
index 8caf59d0c0..749f3cb12b 100644
--- a/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Authorization/AuthTestController_Tests.cs
+++ b/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Authorization/AuthTestController_Tests.cs
@@ -2,21 +2,11 @@
 using System.Security.Claims;
 using System.Threading.Tasks;
 using Shouldly;
-using Volo.Abp.AspNetCore.TestBase;
-using Volo.Abp.Autofac;
-using Volo.Abp.MemoryDb;
-using Volo.Abp.Modularity;
 using Volo.Abp.Security.Claims;
 using Xunit;
 
 namespace Volo.Abp.AspNetCore.Mvc.Authorization
 {
-    [DependsOn(
-        typeof(AbpAspNetCoreTestBaseModule),
-        typeof(AbpMemoryDbTestModule),
-        typeof(AbpAspNetCoreMvcModule),
-        typeof(AbpAutofacModule)
-    )]
     public class AuthTestController_Tests : AspNetCoreMvcTestBase
     {
         private readonly FakeUserClaims _fakeRequiredService;
diff --git a/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Claims/ClaimsMapTestController_Tests.cs b/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Claims/ClaimsMapTestController_Tests.cs
index 154956afdf..146e920d20 100644
--- a/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Claims/ClaimsMapTestController_Tests.cs
+++ b/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Claims/ClaimsMapTestController_Tests.cs
@@ -2,20 +2,10 @@
 using System.Threading.Tasks;
 using Shouldly;
 using Volo.Abp.AspNetCore.Mvc.Authorization;
-using Volo.Abp.AspNetCore.TestBase;
-using Volo.Abp.Autofac;
-using Volo.Abp.MemoryDb;
-using Volo.Abp.Modularity;
 using Xunit;
 
 namespace Volo.Abp.AspNetCore.Mvc.Security.Claims
 {
-    [DependsOn(
-        typeof(AbpAspNetCoreTestBaseModule),
-        typeof(AbpMemoryDbTestModule),
-        typeof(AbpAspNetCoreMvcModule),
-        typeof(AbpAutofacModule)
-    )]
     public class ClaimsMapTestController_Tests : AspNetCoreMvcTestBase
     {
         private readonly FakeUserClaims _fakeRequiredService;
diff --git a/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Headers/SecurityHeadersTestController.cs b/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Headers/SecurityHeadersTestController.cs
new file mode 100644
index 0000000000..034d11440b
--- /dev/null
+++ b/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Headers/SecurityHeadersTestController.cs
@@ -0,0 +1,12 @@
+using Microsoft.AspNetCore.Mvc;
+
+namespace Volo.Abp.AspNetCore.Mvc.Security.Headers
+{
+    public class SecurityHeadersTestController : AbpController
+    {
+        public ActionResult Get()
+        {
+            return Content("OK");
+        }
+    }
+}
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
new file mode 100644
index 0000000000..5902d7cbbe
--- /dev/null
+++ b/framework/test/Volo.Abp.AspNetCore.Mvc.Tests/Volo/Abp/AspNetCore/Mvc/Security/Headers/SecurityHeadersTestController_Tests.cs
@@ -0,0 +1,19 @@
+using System.Linq;
+using System.Threading.Tasks;
+using Shouldly;
+using Xunit;
+
+namespace Volo.Abp.AspNetCore.Mvc.Security.Headers
+{
+    public class SecurityHeadersTestController_Tests : AspNetCoreMvcTestBase
+    {
+        [Fact]
+        public async Task SecurityHeaders_Should_Be_Added()
+        {
+            var responseMessage = await GetResponseAsync("/SecurityHeadersTest/Get");
+            responseMessage.Headers.ShouldContain(x => x.Key == "X-Content-Type-Options" & x.Value.First().ToString() == "nosniff");
+            responseMessage.Headers.ShouldContain(x => x.Key == "X-XSS-Protection" & x.Value.First().ToString() == "1; mode=block");
+            responseMessage.Headers.ShouldContain(x => x.Key == "X-Frame-Options" & x.Value.First().ToString() == "SAMEORIGIN");
+        }
+    }
+}