Allow accessing an IServiceProvider when configuring a SecurityHeaderPolicyBuilder

src/NetEscapades.AspNetCore.SecurityHeaders/ServiceCollectionExtensions.cs

@@ -13,7 +13,7 @@ public static class ServiceCollectionExtensions
     /// Creates a builder for configuring security header policies.
     /// </summary>
     /// <param name="services">The <see cref="IServiceCollection" /> to add services to.</param>
-    /// <returns>The <see cref="IServiceCollection"/> so that additional calls can be chained.</returns>
+    /// <returns>The <see cref="SecurityHeaderPolicyBuilder"/> so that header policies can be configured.</returns>
     public static SecurityHeaderPolicyBuilder AddSecurityHeaderPolicies(this IServiceCollection services)
     {
         if (services == null)
@@ -25,4 +25,35 @@ public static SecurityHeaderPolicyBuilder AddSecurityHeaderPolicies(this IServic
         services.AddSingleton(options);
         return new SecurityHeaderPolicyBuilder(options);
     }
+
+    /// <summary>
+    /// Creates a builder for configuring security header policies.
+    /// </summary>
+    /// <param name="services">The <see cref="IServiceCollection" /> to add services to.</param>
+    /// <param name="configure">A configuration method that provides access to an <see cref="IServiceProvider"/>
+    /// to configure your header policies</param>
+    /// <returns>The <see cref="IServiceCollection"/> so that additional calls can be chained.</returns>
+    public static IServiceCollection AddSecurityHeaderPolicies(
+        this IServiceCollection services,
+        Action<SecurityHeaderPolicyBuilder, IServiceProvider> configure)
+    {
+        if (services == null)
+        {
+            throw new ArgumentNullException(nameof(services));
+        }
+
+        if (configure == null)
+        {
+            throw new ArgumentNullException(nameof(configure));
+        }
+
+        var options = new CustomHeaderOptions();
+        services.AddSingleton<CustomHeaderOptions>(provider =>
+        {
+            configure(new SecurityHeaderPolicyBuilder(options), provider);
+            return options;
+        });
+
+        return services;
+    }
 }

test/NetEscapades.AspNetCore.SecurityHeaders.Test/PublicApiTest.PublicApiHasNotChanged.verified.txt

@@ -294,6 +294,7 @@ namespace Microsoft.Extensions.DependencyInjection
     public static class ServiceCollectionExtensions
     {
         public static NetEscapades.AspNetCore.SecurityHeaders.Infrastructure.SecurityHeaderPolicyBuilder AddSecurityHeaderPolicies(this Microsoft.Extensions.DependencyInjection.IServiceCollection services) { }
+        public static Microsoft.Extensions.DependencyInjection.IServiceCollection AddSecurityHeaderPolicies(this Microsoft.Extensions.DependencyInjection.IServiceCollection services, System.Action<NetEscapades.AspNetCore.SecurityHeaders.Infrastructure.SecurityHeaderPolicyBuilder, System.IServiceProvider> configure) { }
     }
 }
 namespace NetEscapades.AspNetCore.SecurityHeaders.Headers.ContentSecurityPolicy

test/NetEscapades.AspNetCore.SecurityHeaders.Test/SecurityHeadersMiddlewareTests.cs

@@ -2355,6 +2355,41 @@ public async Task HttpRequest_CanApplyDifferentPolicyBasedOnResponseContentType(
         response.Headers.AssertHttpRequestDefaultSecurityHeaders();
     }
 
+    [Fact]
+    public async Task HttpRequest_CanUseProviderToConfigurePolicies()
+    {
+        // Arrange
+        var hostBuilder = new WebHostBuilder()
+            .ConfigureServices(s =>
+            {
+                s.AddSingleton<HeaderPolicyCollectionFactory>();
+                s.AddSecurityHeaderPolicies((builder, provider) =>
+                    {
+                        var service = provider.GetRequiredService<HeaderPolicyCollectionFactory>();
+                        builder.SetDefaultPolicy(service.GetPolicy(null));
+                    });
+            })
+            .Configure(app =>
+            {
+                app.UseSecurityHeaders();
+                app.Run(async context =>
+                {
+                    context.Response.ContentType = "text/html";
+                    await context.Response.WriteAsync("Test response");
+                });
+            });
+
+        using (var server = new TestServer(hostBuilder))
+        {
+            // Act
+            // Add header
+            var response = await server.CreateRequest("/")
+                .SendAsync("PUT");
+            response.EnsureSuccessStatusCode();
+            response.Headers.Should().ContainKey("Custom-Header").WhoseValue.Should().Contain("Default");
+        }
+    }
+
     private class HeaderPolicyCollectionFactory
     {
         private readonly HeaderPolicyCollection _default = new HeaderPolicyCollection().AddCustomHeader("Custom-Header", "Default");