From 453c609b054fb9501ba08a1f4051b6617d5722e9 Mon Sep 17 00:00:00 2001
From: Saturn225 <101260782+Saturn225@users.noreply.github.com>
Date: Tue, 1 Oct 2024 01:48:55 +0530
Subject: [PATCH] chore: cleanup and fix 404 and 405 tests

---
 .../main/scala/zio/http/RoutePattern.scala    |  6 +++
 .../src/main/scala/zio/http/Routes.scala      | 48 ++++++++++++-------
 2 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/zio-http/shared/src/main/scala/zio/http/RoutePattern.scala b/zio-http/shared/src/main/scala/zio/http/RoutePattern.scala
index c42739408b..7bcdff3505 100644
--- a/zio-http/shared/src/main/scala/zio/http/RoutePattern.scala
+++ b/zio-http/shared/src/main/scala/zio/http/RoutePattern.scala
@@ -185,6 +185,12 @@ object RoutePattern                                                       {
       else forMethod ++ wildcardsTree.get(path)
     }
 
+    def getAllMethods(path: Path): Set[Method] = {
+      roots.collect {
+        case (method, subtree) if subtree.get(path).nonEmpty => method
+      }.toSet
+    }
+
     def map[B](f: A => B): Tree[B] =
       Tree(roots.map { case (k, v) =>
         k -> v.map(f)
diff --git a/zio-http/shared/src/main/scala/zio/http/Routes.scala b/zio-http/shared/src/main/scala/zio/http/Routes.scala
index 5847d9524e..05503d5d0d 100644
--- a/zio-http/shared/src/main/scala/zio/http/Routes.scala
+++ b/zio-http/shared/src/main/scala/zio/http/Routes.scala
@@ -248,22 +248,34 @@ final case class Routes[-Env, +Err](routes: Chunk[zio.http.Route[Env, Err]]) { s
     val tree                  = self.tree
     Handler
       .fromFunctionHandler[Request] { req =>
-        val chunk = tree.get(req.method, req.path)
-        chunk.length match {
-          case 0 => Handler.notFound
-          case 1 => chunk(0)
-          case n => // TODO: Support precomputed fallback among all chunk elements
-            var acc = chunk(0)
-            var i   = 1
-            while (i < n) {
-              val h = chunk(i)
-              acc = acc.catchAll { response =>
-                if (response.status == Status.NotFound) h
-                else Handler.fail(response)
-              }
-              i += 1
+        val chunk          = tree.get(req.method, req.path)
+        val allowedMethods = tree.getAllMethods(req.path)
+
+        req.method match {
+          case Method.CUSTOM(_)                              =>
+            Handler.fromZIO(ZIO.succeed(Response.status(Status.NotImplemented)))
+          case _ if chunk.isEmpty && allowedMethods.nonEmpty =>
+            Handler.fromZIO(ZIO.succeed(Response.status(Status.MethodNotAllowed)))
+
+          case _ if chunk.isEmpty && allowedMethods.isEmpty =>
+            Handler.notFound
+          case _                                            =>
+            chunk.length match {
+              case 0 => Handler.notFound 
+              case 1 => chunk(0)
+              case n => // TODO: Support precomputed fallback among all chunk elements
+                var acc = chunk(0)
+                var i   = 1
+                while (i < n) {
+                  val h = chunk(i)
+                  acc = acc.catchAll { response =>
+                    if (response.status == Status.NotFound) h
+                    else Handler.fail(response)
+                  }
+                  i += 1
+                }
+                acc
             }
-            acc
         }
       }
       .merge
@@ -287,6 +299,7 @@ final case class Routes[-Env, +Err](routes: Chunk[zio.http.Route[Env, Err]]) { s
     }
     _tree.asInstanceOf[Routes.Tree[Env]]
   }
+
 }
 
 object Routes extends RoutesCompanionVersionSpecific {
@@ -344,6 +357,9 @@ object Routes extends RoutesCompanionVersionSpecific {
     empty @@ Middleware.serveResources(path, resourcePrefix)
 
   private[http] final case class Tree[-Env](tree: RoutePattern.Tree[RequestHandler[Env, Response]]) { self =>
+
+    def getAllMethods(path: Path): Set[Method] = tree.getAllMethods(path)
+
     final def ++[Env1 <: Env](that: Tree[Env1]): Tree[Env1] =
       Tree(self.tree ++ that.tree)
 
@@ -357,7 +373,7 @@ object Routes extends RoutesCompanionVersionSpecific {
     final def get(method: Method, path: Path): Chunk[RequestHandler[Env, Response]] =
       tree.get(method, path)
   }
-  private[http] object Tree                                                                         {
+  private[http] object Tree {
     val empty: Tree[Any] = Tree(RoutePattern.Tree.empty)
 
     def fromRoutes[Env](routes: Chunk[zio.http.Route[Env, Response]])(implicit trace: Trace): Tree[Env] =