From 1e66f645944d6ce657c35a3e990938cf86fd61eb Mon Sep 17 00:00:00 2001
From: Saturn225 <101260782+Saturn225@users.noreply.github.com>
Date: Tue, 1 Oct 2024 20:38:27 +0530
Subject: [PATCH] Update Routes.scala

---
 .../src/main/scala/zio/http/Routes.scala      | 48 ++++++++++++-------
 1 file changed, 32 insertions(+), 16 deletions(-)

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..642044ca0e 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.notImplemented
+          case _                =>
+            chunk.length match {
+              case 0 =>
+                if (allowedMethods.nonEmpty) {
+                  val allowHeader = Header.Allow(NonEmptyChunk.fromIterableOption(allowedMethods).get)
+                  Handler.methodNotAllowed.addHeader(allowHeader)
+                } else {
+                  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] =