Allow configuring custom ServerHttpHeadersWriter for Kotlin DSL

Closes gh-16009

Author: evgeniycheban

config/src/main/java/org/springframework/security/config/web/server/ServerHttpSecurity.java

@@ -3212,6 +3212,18 @@ public HeaderSpec writer(ServerHttpHeadersWriter serverHttpHeadersWriter) {
 			return this;
 		}
 
+		/**
+		 * Configures custom headers writer
+		 * @param serverHttpHeadersWriterCustomizer the {@link Customizer} to provide more
+		 * options for the {@link ServerHttpHeadersWriterSpec}
+		 * @return the {@link HeaderSpec} to customize
+		 */
+		public HeaderSpec serverHttpHeadersWriter(
+				Customizer<ServerHttpHeadersWriterSpec> serverHttpHeadersWriterCustomizer) {
+			serverHttpHeadersWriterCustomizer.customize(new ServerHttpHeadersWriterSpec());
+			return this;
+		}
+
 		/**
 		 * Configures the Strict Transport Security response headers
 		 * @return the {@link HstsSpec} to configure
@@ -3896,6 +3908,30 @@ public HeaderSpec and() {
 
 		}
 
+		/**
+		 * Configures {@link ServerHttpHeadersWriter}
+		 *
+		 * @author Evgeniy Cheban
+		 * @see #serverHttpHeadersWriter(Customizer)
+		 */
+		public final class ServerHttpHeadersWriterSpec {
+
+			private ServerHttpHeadersWriterSpec() {
+			}
+
+			/**
+			 * Configures custom headers writer
+			 * @param writer the {@link ServerHttpHeadersWriter} to provide custom headers
+			 * writer
+			 * @return the {@link HeaderSpec} to customize
+			 */
+			public ServerHttpHeadersWriterSpec writer(ServerHttpHeadersWriter writer) {
+				HeaderSpec.this.writer(writer);
+				return this;
+			}
+
+		}
+
 	}
 
 	/**

config/src/main/kotlin/org/springframework/security/config/web/server/ServerHeadersDsl.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 package org.springframework.security.config.web.server
 
 import org.springframework.security.web.server.header.CacheControlServerHttpHeadersWriter
+import org.springframework.security.web.server.header.ServerHttpHeadersWriter
 import org.springframework.security.web.server.header.ContentTypeOptionsServerHttpHeadersWriter
 import org.springframework.security.web.server.header.ReferrerPolicyServerHttpHeadersWriter
 import org.springframework.security.web.server.header.StrictTransportSecurityServerHttpHeadersWriter
@@ -43,6 +44,7 @@ class ServerHeadersDsl {
     private var crossOriginOpenerPolicy: ((ServerHttpSecurity.HeaderSpec.CrossOriginOpenerPolicySpec) -> Unit)? = null
     private var crossOriginEmbedderPolicy: ((ServerHttpSecurity.HeaderSpec.CrossOriginEmbedderPolicySpec) -> Unit)? = null
     private var crossOriginResourcePolicy: ((ServerHttpSecurity.HeaderSpec.CrossOriginResourcePolicySpec) -> Unit)? = null
+    private var serverHttpHeadersWriter: ((ServerHttpSecurity.HeaderSpec.ServerHttpHeadersWriterSpec) -> Unit)? = null
 
     private var disabled = false
 
@@ -198,6 +200,15 @@ class ServerHeadersDsl {
         this.crossOriginResourcePolicy = ServerCrossOriginResourcePolicyDsl().apply(crossOriginResourcePolicyConfig).get()
     }
 
+    /**
+     * Allows customizations for the [ServerHttpHeadersWriter] which allows writing headers just before the response is committed
+     *
+     * @param serverHttpHeadersWriterConfig the customization to apply to the header
+     */
+    fun serverHttpHeadersWriter(serverHttpHeadersWriterConfig: ServerHttpHeadersWriterDsl.() -> Unit) {
+        this.serverHttpHeadersWriter = ServerHttpHeadersWriterDsl().apply(serverHttpHeadersWriterConfig).get()
+    }
+
     /**
      * Disables HTTP response headers.
      */
@@ -244,6 +255,9 @@ class ServerHeadersDsl {
             crossOriginResourcePolicy?.also {
                 headers.crossOriginResourcePolicy(crossOriginResourcePolicy)
             }
+            serverHttpHeadersWriter?.also {
+                headers.serverHttpHeadersWriter(serverHttpHeadersWriter)
+            }
             if (disabled) {
                 headers.disable()
             }

config/src/main/kotlin/org/springframework/security/config/web/server/ServerHttpHeadersWriterDsl.kt

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2002-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.security.config.web.server
+
+import org.springframework.security.web.server.header.ServerHttpHeadersWriter
+
+/**
+ * A Kotlin DSL to configure the [ServerHttpSecurity] to use custom [ServerHttpHeadersWriter] using
+ * idiomatic Kotlin code.
+ *
+ * @author Evgeniy Cheban
+ * @property writers the list of custom [ServerHttpHeadersWriter] to be used
+ */
+@ServerSecurityMarker
+class ServerHttpHeadersWriterDsl {
+
+    private var writers = mutableListOf<ServerHttpHeadersWriter>()
+
+    fun writer(writer: ServerHttpHeadersWriter) {
+        writers.add(writer)
+    }
+
+    internal fun get(): (ServerHttpSecurity.HeaderSpec.ServerHttpHeadersWriterSpec) -> Unit {
+        return { writerOptions ->
+            writers.also {
+                writers.forEach { writer -> writerOptions.writer(writer) }
+            }
+        }
+    }
+
+}

config/src/test/kotlin/org/springframework/security/config/web/server/ServerHeadersDslTests.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -37,6 +37,7 @@ import org.springframework.security.web.server.header.XFrameOptionsServerHttpHea
 import org.springframework.security.web.server.header.XXssProtectionServerHttpHeadersWriter
 import org.springframework.test.web.reactive.server.WebTestClient
 import org.springframework.web.reactive.config.EnableWebFlux
+import reactor.core.publisher.Mono
 
 /**
  * Tests for [ServerHeadersDsl]
@@ -198,4 +199,50 @@ class ServerHeadersDslTests {
             }
         }
     }
+
+    @Test
+    fun `request when custom server http headers writer configured then custom http headers added`() {
+        this.spring.register(ServerHttpHeadersWriterCustomConfig::class.java).autowire()
+
+        this.client.get()
+            .uri("/")
+            .exchange()
+            .expectHeader().valueEquals("CUSTOM-HEADER-1", "CUSTOM-VALUE-1")
+            .expectHeader().valueEquals("CUSTOM-HEADER-2", "CUSTOM-VALUE-2")
+    }
+
+    @Configuration
+    @EnableWebFluxSecurity
+    @EnableWebFlux
+    open class ServerHttpHeadersWriterCustomConfig {
+        @Bean
+        open fun springWebFilterChain(http: ServerHttpSecurity): SecurityWebFilterChain {
+            return http {
+                headers {
+                    serverHttpHeadersWriter {
+                        writer { exchange ->
+                            Mono.just(exchange)
+                                .doOnNext {
+                                    it.response.headers.add(
+                                        "CUSTOM-HEADER-1",
+                                        "CUSTOM-VALUE-1"
+                                    )
+                                }
+                                .then()
+                        }
+                        writer { exchange ->
+                            Mono.just(exchange)
+                                .doOnNext {
+                                    it.response.headers.add(
+                                        "CUSTOM-HEADER-2",
+                                        "CUSTOM-VALUE-2"
+                                    )
+                                }
+                                .then()
+                        }
+                    }
+                }
+            }
+        }
+    }
 }