From 1ea7092df3fb4697a96419323faa1228f192a19a Mon Sep 17 00:00:00 2001
From: Hosam Aly <hosamaly6@gmail.com>
Date: Thu, 5 Dec 2019 11:14:15 +0000
Subject: [PATCH] Introduce Local Port Routing Builder (#1392)

* Introduce Local Port Route Builder

See discussion at:
https://github.com/ktorio/ktor/pull/1392#issuecomment-540415051
---
 .../ktor/routing/LocalPortRoutingBuilder.kt   | 54 +++++++++++++++++++
 .../server/routing/RoutingProcessingTest.kt   | 39 ++++++++++++++
 2 files changed, 93 insertions(+)
 create mode 100644 ktor-server/ktor-server-core/jvm/src/io/ktor/routing/LocalPortRoutingBuilder.kt

diff --git a/ktor-server/ktor-server-core/jvm/src/io/ktor/routing/LocalPortRoutingBuilder.kt b/ktor-server/ktor-server-core/jvm/src/io/ktor/routing/LocalPortRoutingBuilder.kt
new file mode 100644
index 00000000000..0bfb28aff87
--- /dev/null
+++ b/ktor-server/ktor-server-core/jvm/src/io/ktor/routing/LocalPortRoutingBuilder.kt
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2014-2019 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package io.ktor.routing
+
+import io.ktor.http.parametersOf
+import io.ktor.util.*
+
+/**
+ * Create a route to match the port on which the call was received.
+ *
+ * The selector checks the [io.ktor.request.ApplicationRequest.local] request port,
+ * _ignoring_ HTTP headers such as "Host" or "X-Forwarded-Host".
+ * This is useful for securing routes under separate ports.
+ *
+ * For multi-tenant applications, you may want to use [io.ktor.routing.port],
+ * which takes HTTP headers into consideration.
+ *
+ * @param port the port to match against
+ *
+ * @throws IllegalArgumentException if the port is outside the range of TCP/UDP ports
+ */
+@KtorExperimentalAPI
+fun Route.localPort(port: Int, build: Route.() -> Unit): Route {
+    require(port in 1..65535) { "Port $port must be a positive number between 1 and 65,535" }
+
+    val selector = LocalPortRouteSelector(port)
+    return createChild(selector).apply(build)
+}
+
+/**
+ * Evaluate a route against the port on which the call was received.
+ *
+ * @param port the port to match against
+ */
+@KtorExperimentalAPI
+data class LocalPortRouteSelector(val port: Int) : RouteSelector(RouteSelectorEvaluation.qualityConstant) {
+    override fun evaluate(context: RoutingResolveContext, segmentIndex: Int) =
+        if (context.call.request.local.port == port) {
+            val parameters = parametersOf(LocalPortParameter, port.toString())
+            RouteSelectorEvaluation(true, RouteSelectorEvaluation.qualityConstant, parameters)
+        } else {
+            RouteSelectorEvaluation.Failed
+        }
+
+    companion object {
+        /**
+         * Parameter name for [RoutingApplicationCall.parameters] for request host
+         */
+        @KtorExperimentalAPI
+        const val LocalPortParameter: String = "\$LocalPort"
+    }
+}
diff --git a/ktor-server/ktor-server-tests/jvm/test/io/ktor/tests/server/routing/RoutingProcessingTest.kt b/ktor-server/ktor-server-tests/jvm/test/io/ktor/tests/server/routing/RoutingProcessingTest.kt
index 228399f009c..c44f8bfad25 100644
--- a/ktor-server/ktor-server-tests/jvm/test/io/ktor/tests/server/routing/RoutingProcessingTest.kt
+++ b/ktor-server/ktor-server-tests/jvm/test/io/ktor/tests/server/routing/RoutingProcessingTest.kt
@@ -447,6 +447,45 @@ class RoutingProcessingTest {
         }
     }
 
+    @Test
+    fun `local port route processing`(): Unit = withTestApplication {
+        application.routing {
+            route("/") {
+                // TestApplicationRequest.local defaults to 80 in the absence of headers
+                // so connections paths to port 80 in tests should work, whereas other ports shouldn't
+                localPort(80) {
+                    get("http") {
+                        call.respond("received")
+                    }
+                }
+                localPort(443) {
+                    get("https") {
+                        fail("shouldn't be received")
+                    }
+                }
+            }
+        }
+
+        // accepts calls to the specified port
+        handleRequest(HttpMethod.Get, "/http").apply {
+            assertEquals("received", response.content)
+        }
+
+        // ignores calls to different ports
+        handleRequest(HttpMethod.Get, "/https").apply {
+            assertNull(response.content)
+        }
+
+        // I tried to write a test to confirm that it ignores the HTTP Host header,
+        // but I couldn't get it to work without adding headers, because
+        // [io.ktor.server.testing.TestApplicationRequest.local] is hard-coded to
+        // extract the value of those headers.
+        // (even though, according to docs, it shouldn't; this should be done by `origin`)
+
+        // I also tried to create a test listening to multiple ports, but I couldn't get it
+        // to work because of the same reason above.
+    }
+
     @Test
     fun `routing with tracing`() = withTestApplication {
         var trace: RoutingResolveTrace? = null