Allow setting CSP per specific request (#5)

LukasForst · Sep 19, 2022 · a3eef16 · a3eef16
1 parent c0b0051
commit a3eef16
Show file tree

Hide file tree

Showing 6 changed files with 67 additions and 34 deletions.
diff --git a/Makefile b/Makefile
@@ -1,6 +1,10 @@
 detekt:
 	./gradlew detekt
 
+detekt-correct:
+	./gradlew detekt --auto-correct || true
+	./gradlew detekt
+
 test:
 	./gradlew test
 

diff --git a/README.md b/README.md
@@ -69,13 +69,21 @@ Plugin that allows setting [Content-Security-Policy](https://developer.mozilla.o
  * Minimal Ktor application using Content Security Policy.
  */
 fun Application.minimalExample() {
+    // this sets Content-Security-Policy
     install(ContentSecurityPolicy) {
-        skipWhen { call ->
-            call.request.path().startsWith("/some-ignored-route")
+        policy { call, body ->
+            when (call.request.path()) {
+                "/specific" -> mapOf("default-src" to "'none'")
+                "/ignored" -> null
+                else -> mapOf("default-src" to "'self'")
+            }
         }
-        policy(
-            "default-src" to "'none'"
-        )
+    }
+    // basic routing
+    routing {
+        get("/specific") { call.respond(HttpStatusCode.OK) }
+        get("/ignored") { call.respond(HttpStatusCode.OK) }
+        get("/") { call.respond(HttpStatusCode.OK) }
     }
 }
 ```
@@ -170,4 +178,4 @@ fun Application.minimalExample() {
         }
     }
 }
-```
+```
diff --git a/ktor-content-security-policy/README.md b/ktor-content-security-policy/README.md
@@ -21,17 +21,25 @@ Minimal usage:
  * Minimal Ktor application using Content Security Policy.
  */
 fun Application.minimalExample() {
+    // this sets Content-Security-Policy
     install(ContentSecurityPolicy) {
-        skipWhen { call ->
-            call.request.path().startsWith("/some-ignored-route")
+        policy { call, body ->
+            when (call.request.path()) {
+                "/specific" -> mapOf("default-src" to "'none'")
+                "/ignored" -> null
+                else -> mapOf("default-src" to "'self'")
+            }
         }
-        policy(
-            "default-src" to "'none'"
-        )
+    }
+    // basic routing
+    routing {
+        get("/specific") { call.respond(HttpStatusCode.OK) }
+        get("/ignored") { call.respond(HttpStatusCode.OK) }
+        get("/") { call.respond(HttpStatusCode.OK) }
     }
 }
 ```
 
 For details see [MinimalExampleApp.kt](src/test/kotlin/dev/forst/ktor/csp/MinimalExampleApp.kt) with this example
 application and [TestMinimalExampleApp.kt](src/test/kotlin/dev/forst/ktor/csp/TestMinimalExampleApp.kt) which verifies
-that this app works as expected.
+that this app works as expected.
diff --git a/ktor-content-security-policy/src/main/kotlin/dev/forst/ktor/csp/ContentSecurityPolicy.kt b/ktor-content-security-policy/src/main/kotlin/dev/forst/ktor/csp/ContentSecurityPolicy.kt
@@ -10,14 +10,17 @@ import io.ktor.server.application.createRouteScopedPlugin
  */
 class ContentSecurityPolicyConfiguration {
     /**
-     * Final value that is injected to the request as a value of "Content-Security-Policy: [cspHeader]".
+     * Selector when not to include Content-Security-Policy header.
      */
-    var cspHeader: String = "default-src 'none'"
+    var skipWhen: (ApplicationCall) -> Boolean = { false }
 
     /**
-     * Selector when not to include Content-Security-Policy header.
+     * Call specific policy, if it returns map, it is used as a csp policy. When returns null,
+     * the CSP Header is not set.
+     *
+     * By default, it uses strict policy "default-src 'none'".
      */
-    var skipWhen: (ApplicationCall) -> Boolean = { false }
+    var policy: (call: ApplicationCall, body: Any) -> Map<String, String?>? = { _, _ -> mapOf("default-src" to "'none'") }
 
     /**
      * Selector when not to include Content-Security-Policy header.
@@ -34,17 +37,18 @@ class ContentSecurityPolicyConfiguration {
     fun policy(vararg policies: Pair<String, String?>) = policy(policies.toMap())
 
     /**
-     * Builder for policies.
+     * Adds policies as map. These are used if there are no other more specific policies set
+     * for the given call.
      */
-    fun policy(policyBuilder: () -> Map<String, String?>) = policy(policyBuilder())
+    fun policy(policies: Map<String, String?>) {
+        policy { _, _ -> policies }
+    }
 
     /**
-     * Adds policies as map.
+     * Adds policies that are specific per request, when [policyBuilder] returns null, no CSP header is set.
      */
-    fun policy(policies: Map<String, String?>) {
-        cspHeader = policies
-            .map { (key, value) -> if (value != null) "$key $value" else key }
-            .joinToString(";")
+    fun policy(policyBuilder: (call: ApplicationCall, body: Any) -> Map<String, String?>?) {
+        this.policy = policyBuilder
     }
 }
 
@@ -55,12 +59,17 @@ val ContentSecurityPolicy = createRouteScopedPlugin(
     name = "ContentSecurityPolicy",
     createConfiguration = ::ContentSecurityPolicyConfiguration
 ) {
-    val headerValue = pluginConfig.cspHeader
+    val policy = pluginConfig.policy
     val skip = pluginConfig.skipWhen
 
-    onCallRespond { call, _ ->
+    onCallRespond { call, body ->
         if (!skip(call)) {
-            call.response.headers.append("Content-Security-Policy", headerValue)
+            val header = policy(call, body)?.toCspHeader() ?: return@onCallRespond
+            call.response.headers.append("Content-Security-Policy", header)
         }
     }
 }
+
+private fun Map<String, String?>.toCspHeader(): String = this
+    .map { (key, value) -> if (value != null) "$key $value" else key }
+    .joinToString(";")
diff --git a/ktor-content-security-policy/src/test/kotlin/dev/forst/ktor/csp/MinimalExampleApp.kt b/ktor-content-security-policy/src/test/kotlin/dev/forst/ktor/csp/MinimalExampleApp.kt
@@ -15,16 +15,18 @@ import io.ktor.server.routing.routing
 fun Application.minimalExample() {
     // this sets Content-Security-Policy
     install(ContentSecurityPolicy) {
-        skipWhen { call ->
-            call.request.path().startsWith("/ignored")
+        policy { call, body ->
+            when (call.request.path()) {
+                "/specific" -> mapOf("default-src" to "'none'")
+                "/ignored" -> null
+                else -> mapOf("default-src" to "'self'")
+            }
         }
-        policy(
-            "default-src" to "'none'"
-        )
     }
     // basic routing
     routing {
-        get("/") { call.respond(HttpStatusCode.OK) }
+        get("/specific") { call.respond(HttpStatusCode.OK) }
         get("/ignored") { call.respond(HttpStatusCode.OK) }
+        get("/") { call.respond(HttpStatusCode.OK) }
     }
 }
diff --git a/ktor-content-security-policy/src/test/kotlin/dev/forst/ktor/csp/TestMinimalExampleApp.kt b/ktor-content-security-policy/src/test/kotlin/dev/forst/ktor/csp/TestMinimalExampleApp.kt
@@ -3,16 +3,18 @@ package dev.forst.ktor.csp
 import io.ktor.client.request.get
 import io.ktor.server.application.Application
 import io.ktor.server.testing.testApplication
+import org.junit.jupiter.api.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertNull
-import org.junit.jupiter.api.Test
 
 class TestMinimalExampleApp {
     @Test
     fun `test minimal example app works as expected`() = testApplication {
         application(Application::minimalExample)
         // this should return csp header
-        val responseWithCsp = client.get("/")
+        var responseWithCsp = client.get("/")
+        assertEquals("default-src 'self'", responseWithCsp.headers["Content-Security-Policy"])
+        responseWithCsp = client.get("/specific")
         assertEquals("default-src 'none'", responseWithCsp.headers["Content-Security-Policy"])
         // this should not
         val skippedResponse = client.get("/ignored")