From e908612f7d52d0b465e875693f443710c98d2fea Mon Sep 17 00:00:00 2001
From: dmitrievanthony <dmitrievanthony@gmail.com>
Date: Wed, 29 Apr 2020 15:53:43 +0300
Subject: [PATCH] Fix auto head response feature (#1835).

---
 ktor-http/api/ktor-http.api                      | 10 +++++++---
 ktor-http/common/src/io/ktor/http/HttpMethod.kt  | 16 +++++++++++++++-
 .../jvm/src/io/ktor/features/AutoHeadResponse.kt |  5 ++++-
 .../jvm/src/io/ktor/routing/RouteSelector.kt     |  2 +-
 .../io/ktor/tests/server/features/HeadTest.kt    |  9 +++++++++
 5 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/ktor-http/api/ktor-http.api b/ktor-http/api/ktor-http.api
index 6b6c946f3be..c3461831b6a 100644
--- a/ktor-http/api/ktor-http.api
+++ b/ktor-http/api/ktor-http.api
@@ -579,13 +579,17 @@ public final class io/ktor/http/HttpMessagePropertiesKt {
 
 public final class io/ktor/http/HttpMethod {
 	public static final field Companion Lio/ktor/http/HttpMethod$Companion;
-	public fun <init> (Ljava/lang/String;)V
+	public fun <init> (Ljava/lang/String;Ljava/util/List;)V
+	public synthetic fun <init> (Ljava/lang/String;Ljava/util/List;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun component1 ()Ljava/lang/String;
-	public final fun copy (Ljava/lang/String;)Lio/ktor/http/HttpMethod;
-	public static synthetic fun copy$default (Lio/ktor/http/HttpMethod;Ljava/lang/String;ILjava/lang/Object;)Lio/ktor/http/HttpMethod;
+	public final fun component2 ()Ljava/util/List;
+	public final fun copy (Ljava/lang/String;Ljava/util/List;)Lio/ktor/http/HttpMethod;
+	public static synthetic fun copy$default (Lio/ktor/http/HttpMethod;Ljava/lang/String;Ljava/util/List;ILjava/lang/Object;)Lio/ktor/http/HttpMethod;
 	public fun equals (Ljava/lang/Object;)Z
+	public final fun getAggregate ()Ljava/util/List;
 	public final fun getValue ()Ljava/lang/String;
 	public fun hashCode ()I
+	public final fun match (Lio/ktor/http/HttpMethod;)Z
 	public fun toString ()Ljava/lang/String;
 }
 
diff --git a/ktor-http/common/src/io/ktor/http/HttpMethod.kt b/ktor-http/common/src/io/ktor/http/HttpMethod.kt
index 8554bcc9a9a..6f9de2bb33a 100644
--- a/ktor-http/common/src/io/ktor/http/HttpMethod.kt
+++ b/ktor-http/common/src/io/ktor/http/HttpMethod.kt
@@ -4,11 +4,25 @@
 
 package io.ktor.http
 
+import io.ktor.util.*
+
 /**
  * Represents an HTTP method (verb)
  * @property value contains method name
  */
-data class HttpMethod(val value: String) {
+data class HttpMethod(val value: String, @InternalAPI val aggregate: List<HttpMethod> = listOf()) {
+    /**
+     * Checks if the specified HTTP [method] matches this instance of the HTTP method. Specified method matches if it's
+     * equal to this HTTP method or at least one of methods this HTTP method aggregates.
+     */
+    fun match(method: HttpMethod): Boolean {
+        if (this == method) {
+            return true
+        }
+
+        return aggregate.contains(method) || method.aggregate.contains(this)
+    }
+
     @Suppress("KDocMissingDocumentation", "PublicApiImplicitType")
     companion object {
         val Get = HttpMethod("GET")
diff --git a/ktor-server/ktor-server-core/jvm/src/io/ktor/features/AutoHeadResponse.kt b/ktor-server/ktor-server-core/jvm/src/io/ktor/features/AutoHeadResponse.kt
index a69f393dcbd..d8da2c9d6d8 100644
--- a/ktor-server/ktor-server-core/jvm/src/io/ktor/features/AutoHeadResponse.kt
+++ b/ktor-server/ktor-server-core/jvm/src/io/ktor/features/AutoHeadResponse.kt
@@ -33,7 +33,10 @@ object AutoHeadResponse : ApplicationFeature<ApplicationCallPipeline, Unit, Unit
 
                 // Pretend the request was with GET method so that all normal routes and interceptors work
                 // but in the end we will drop the content
-                call.mutableOriginConnectionPoint.method = HttpMethod.Get
+                call.mutableOriginConnectionPoint.method = HttpMethod(
+                    "GET_OR_HEAD",
+                    listOf(HttpMethod.Get, HttpMethod.Head)
+                )
             }
         }
     }
diff --git a/ktor-server/ktor-server-core/jvm/src/io/ktor/routing/RouteSelector.kt b/ktor-server/ktor-server-core/jvm/src/io/ktor/routing/RouteSelector.kt
index d8f76da7054..d800b8db5c6 100644
--- a/ktor-server/ktor-server-core/jvm/src/io/ktor/routing/RouteSelector.kt
+++ b/ktor-server/ktor-server-core/jvm/src/io/ktor/routing/RouteSelector.kt
@@ -349,7 +349,7 @@ data class AndRouteSelector(val first: RouteSelector, val second: RouteSelector)
  */
 data class HttpMethodRouteSelector(val method: HttpMethod) : RouteSelector(RouteSelectorEvaluation.qualityParameter) {
     override fun evaluate(context: RoutingResolveContext, segmentIndex: Int): RouteSelectorEvaluation {
-        if (context.call.request.httpMethod == method)
+        if (context.call.request.httpMethod.match(method))
             return RouteSelectorEvaluation.Constant
         return RouteSelectorEvaluation.Failed
     }
diff --git a/ktor-server/ktor-server-tests/jvm/test/io/ktor/tests/server/features/HeadTest.kt b/ktor-server/ktor-server-tests/jvm/test/io/ktor/tests/server/features/HeadTest.kt
index 7c9640c89ab..578da34b880 100644
--- a/ktor-server/ktor-server-tests/jvm/test/io/ktor/tests/server/features/HeadTest.kt
+++ b/ktor-server/ktor-server-tests/jvm/test/io/ktor/tests/server/features/HeadTest.kt
@@ -26,6 +26,10 @@ class HeadTest {
                     call.response.header("M", "1")
                     call.respond("Hello")
                 }
+
+                head("/head") {
+                    call.respond(HttpStatusCode.OK)
+                }
             }
 
             handleRequest(HttpMethod.Get, "/").let { call ->
@@ -39,6 +43,11 @@ class HeadTest {
                 assertNull(call.response.content)
                 assertEquals("1", call.response.headers["M"])
             }
+
+            handleRequest(HttpMethod.Head, "/head").let { call ->
+                assertEquals(HttpStatusCode.OK, call.response.status())
+                assertNull(call.response.content)
+            }
         }
     }