KTOR-721 Provide way to customize unhandled responses (#3278)

ktorio · Nov 30, 2022 · ee4ca55 · ee4ca55
1 parent 4279241
commit ee4ca55
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 0 deletions.
diff --git a/ktor-server/ktor-server-plugins/ktor-server-status-pages/api/ktor-server-status-pages.api b/ktor-server/ktor-server-plugins/ktor-server-status-pages/api/ktor-server-status-pages.api
@@ -5,6 +5,7 @@ public final class io/ktor/server/plugins/statuspages/StatusPagesConfig {
 	public final fun getStatuses ()Ljava/util/Map;
 	public final fun status ([Lio/ktor/http/HttpStatusCode;Lkotlin/jvm/functions/Function3;)V
 	public final fun statusWithContext ([Lio/ktor/http/HttpStatusCode;Lkotlin/jvm/functions/Function3;)V
+	public final fun unhandled (Lkotlin/jvm/functions/Function2;)V
 }
 
 public final class io/ktor/server/plugins/statuspages/StatusPagesConfig$StatusContext {

diff --git a/.../ktor-server-status-pages/jvmAndNix/src/io/ktor/server/plugins/statuspages/StatusPages.kt b/.../ktor-server-status-pages/jvmAndNix/src/io/ktor/server/plugins/statuspages/StatusPages.kt
@@ -34,6 +34,7 @@ public val StatusPages: ApplicationPlugin<StatusPagesConfig> = createApplication
 
     val exceptions = HashMap(pluginConfig.exceptions)
     val statuses = HashMap(pluginConfig.statuses)
+    val unhandled = pluginConfig.unhandled
 
     fun findHandlerByValue(cause: Throwable): HandlerFunction? {
         val keys = exceptions.keys.filter { cause.instanceOf(it) }
@@ -92,6 +93,11 @@ public val StatusPages: ApplicationPlugin<StatusPagesConfig> = createApplication
             handler(call, cause)
         }
     }
+
+    on(BeforeFallback) { call ->
+        if (call.isHandled) return@on
+        unhandled(call)
+    }
 }
 
 /**
@@ -111,6 +117,8 @@ public class StatusPagesConfig {
         suspend (call: ApplicationCall, content: OutgoingContent, code: HttpStatusCode) -> Unit> =
         mutableMapOf()
 
+    internal var unhandled: suspend (ApplicationCall) -> Unit = {}
+
     /**
      * Register an exception [handler] for the exception type [T] and its children.
      */
@@ -143,6 +151,13 @@ public class StatusPagesConfig {
         }
     }
 
+    /**
+     * Register a [handler] for the unhandled calls.
+     */
+    public fun unhandled(handler: suspend (ApplicationCall) -> Unit) {
+        unhandled = handler
+    }
+
     /**
      * Register a status [handler] for the [status] code.
      */

diff --git a/...-server-status-pages/jvmAndNix/src/io/ktor/server/plugins/statuspages/StatusPagesUtils.kt b/...-server-status-pages/jvmAndNix/src/io/ktor/server/plugins/statuspages/StatusPagesUtils.kt
@@ -4,6 +4,16 @@
 
 package io.ktor.server.plugins.statuspages
 
+import io.ktor.server.application.*
+import io.ktor.util.pipeline.*
 import kotlin.reflect.*
 
 internal expect fun selectNearestParentClass(cause: Throwable, keys: List<KClass<*>>): KClass<*>?
+
+internal object BeforeFallback : Hook<suspend (ApplicationCall) -> Unit> {
+    override fun install(pipeline: ApplicationCallPipeline, handler: suspend (ApplicationCall) -> Unit) {
+        val phase = PipelinePhase("BeforeFallback")
+        pipeline.insertPhaseBefore(ApplicationCallPipeline.Fallback, phase)
+        pipeline.intercept(phase) { handler(context) }
+    }
+}
diff --git a/...-server-status-pages/jvmAndNix/test/io/ktor/server/plugins/statuspages/StatusPagesTest.kt b/...-server-status-pages/jvmAndNix/test/io/ktor/server/plugins/statuspages/StatusPagesTest.kt
@@ -540,4 +540,31 @@ class StatusPagesTest {
         assertEquals(HttpStatusCode.InternalServerError, response.status)
         assertEquals("Handled", response.bodyAsText())
     }
+
+    @Test
+    fun testUnhandled() = testApplication {
+        install(StatusPages) {
+            unhandled { call -> call.respond(HttpStatusCode.InternalServerError, "body") }
+        }
+        routing {
+            get("route") {
+                call.response.headers.append("Custom-Header", "Custom-Value")
+            }
+            get("handled") {
+                call.respond("OK")
+            }
+        }
+
+        val responseHandled = client.get("handled")
+        assertEquals("OK", responseHandled.bodyAsText())
+
+        val responseNoRoute = client.get("/no-route")
+        assertEquals(HttpStatusCode.InternalServerError, responseNoRoute.status)
+        assertEquals("body", responseNoRoute.bodyAsText())
+
+        val responseWithRoute = client.get("/route")
+        assertEquals(HttpStatusCode.InternalServerError, responseWithRoute.status)
+        assertEquals("Custom-Value", responseWithRoute.headers["Custom-Header"])
+        assertEquals("body", responseWithRoute.bodyAsText())
+    }
 }