Allow configuring custom ServerHttpHeadersWriter for Kotlin DSL

Closes spring-projectsgh-16009
evgeniycheban · Dec 9, 2024 · ef3f46c · ef3f46c
1 parent 30c9860
commit ef3f46c
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 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
@@ -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,15 @@ class ServerHeadersDsl {
         this.crossOriginResourcePolicy = ServerCrossOriginResourcePolicyDsl().apply(crossOriginResourcePolicyConfig).get()
     }
 
+    /**
+     * Configures custom headers writer
+     *
+     * @param writer the [ServerHttpHeadersWriter] to provide custom headers writer
+     */
+    fun writer(writer: ServerHttpHeadersWriter) {
+        this.writers.add(writer)
+    }
+
     /**
      * Disables HTTP response headers.
      */
@@ -244,6 +255,9 @@ class ServerHeadersDsl {
             crossOriginResourcePolicy?.also {
                 headers.crossOriginResourcePolicy(crossOriginResourcePolicy)
             }
+            writers.also {
+                writers.forEach { writer -> headers.writer(writer) }
+            }
             if (disabled) {
                 headers.disable()
             }

diff --git a/...g/src/test/kotlin/org/springframework/security/config/web/server/ServerHeadersDslTests.kt b/...g/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()
+                    }
+                }
+            }
+        }
+    }
 }