From f45cc22e11e1a8e35f4717ad0055197642ff5371 Mon Sep 17 00:00:00 2001
From: Evgeniy Cheban <mister.cheban@gmail.com>
Date: Wed, 20 Nov 2024 06:09:14 +0200
Subject: [PATCH] Allow configuring custom ServerHttpHeadersWriter for Kotlin
 DSL

Closes gh-16009
---
 .../config/web/server/ServerHeadersDsl.kt     | 17 ++++++-
 .../web/server/ServerHeadersDslTests.kt       | 47 ++++++++++++++++++-
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/config/src/main/kotlin/org/springframework/security/config/web/server/ServerHeadersDsl.kt b/config/src/main/kotlin/org/springframework/security/config/web/server/ServerHeadersDsl.kt
index 37bd1f177a9..86289744439 100644
--- a/config/src/main/kotlin/org/springframework/security/config/web/server/ServerHeadersDsl.kt
+++ b/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 writers = mutableListOf<ServerHttpHeadersWriter>()
 
     private var disabled = false
 
@@ -198,6 +200,16 @@ class ServerHeadersDsl {
         this.crossOriginResourcePolicy = ServerCrossOriginResourcePolicyDsl().apply(crossOriginResourcePolicyConfig).get()
     }
 
+    /**
+     * Configures custom headers writer
+     *
+     * @since 6.5
+     * @param writer the [ServerHttpHeadersWriter] to provide custom headers writer
+     */
+    fun writer(writer: ServerHttpHeadersWriter) {
+        this.writers.add(writer)
+    }
+
     /**
      * Disables HTTP response headers.
      */
@@ -244,6 +256,9 @@ class ServerHeadersDsl {
             crossOriginResourcePolicy?.also {
                 headers.crossOriginResourcePolicy(crossOriginResourcePolicy)
             }
+            writers.also {
+                writers.forEach { writer -> headers.writer(writer) }
+            }
             if (disabled) {
                 headers.disable()
             }
diff --git a/config/src/test/kotlin/org/springframework/security/config/web/server/ServerHeadersDslTests.kt b/config/src/test/kotlin/org/springframework/security/config/web/server/ServerHeadersDslTests.kt
index dfa78726cac..78bfd4d57d4 100644
--- a/config/src/test/kotlin/org/springframework/security/config/web/server/ServerHeadersDslTests.kt
+++ b/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,48 @@ 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 {
+                    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()
+                    }
+                }
+            }
+        }
+    }
 }