From d20f96b902ac4324d2c43c78e3169ed5652b87b9 Mon Sep 17 00:00:00 2001
From: Nabil Abdel-Hafeez <7283535+987Nabil@users.noreply.github.com>
Date: Thu, 19 Sep 2024 00:20:09 +0200
Subject: [PATCH] Allow applying aspects only to handlers with `Request` input
 (#3141)

---
 .../src/test/scala/zio/http/RouteSpec.scala   | 21 +++++++++++++++++++
 .../zio/http/HandlerVersionSpecific.scala     |  2 +-
 .../http/RouteCompanionVersionSpecific.scala  | 12 +++++++++++
 .../zio/http/HandlerVersionSpecific.scala     |  2 +-
 .../http/RouteCompanionVersionSpecific.scala  | 10 +++++++++
 .../src/main/scala/zio/http/Handler.scala     |  4 ++--
 .../main/scala/zio/http/HandlerAspect.scala   | 20 ++++++++++++++++++
 .../src/main/scala/zio/http/Middleware.scala  |  3 +++
 .../zio/http/RequestHandlerInput.scala.scala  | 17 +++++++++++++--
 .../src/main/scala/zio/http/Route.scala       | 18 ++++++++++++++--
 10 files changed, 101 insertions(+), 8 deletions(-)
 create mode 100644 zio-http/shared/src/main/scala-2/zio/http/RouteCompanionVersionSpecific.scala
 create mode 100644 zio-http/shared/src/main/scala-3/zio/http/RouteCompanionVersionSpecific.scala

diff --git a/zio-http/jvm/src/test/scala/zio/http/RouteSpec.scala b/zio-http/jvm/src/test/scala/zio/http/RouteSpec.scala
index a0fc0f8b0ae..009b4ee5ff8 100644
--- a/zio-http/jvm/src/test/scala/zio/http/RouteSpec.scala
+++ b/zio-http/jvm/src/test/scala/zio/http/RouteSpec.scala
@@ -16,6 +16,8 @@
 
 package zio.http
 
+import scala.annotation.nowarn
+
 import zio._
 import zio.test._
 
@@ -209,5 +211,24 @@ object RouteSpec extends ZIOHttpSpec {
         } yield assertTrue(bodyString == expected)
       },
     ),
+    test("Apply context middleware to route with path parameter") {
+      case class WebSession(id: Int)
+
+      def maybeWebSession: HandlerAspect[Any, Option[WebSession]] =
+        HandlerAspect.interceptIncomingHandler(
+          Handler.fromFunctionZIO[Request] { req =>
+            ZIO.succeed((req, None))
+          },
+        )
+
+      val route = (Method.GET / "base" / string("1") -> handler((_: String, _: Request) => {
+        withContext((_: Option[WebSession]) => {
+          ZIO.logInfo("Hello").as(Response.ok)
+        })
+      })) @@ maybeWebSession
+      route(Request.get(url"/base/1")).map { response =>
+        assertTrue(extractStatus(response) == Status.Ok)
+      }
+    },
   )
 }
diff --git a/zio-http/shared/src/main/scala-2/zio/http/HandlerVersionSpecific.scala b/zio-http/shared/src/main/scala-2/zio/http/HandlerVersionSpecific.scala
index 151331cb7d0..1b391308b84 100644
--- a/zio-http/shared/src/main/scala-2/zio/http/HandlerVersionSpecific.scala
+++ b/zio-http/shared/src/main/scala-2/zio/http/HandlerVersionSpecific.scala
@@ -5,7 +5,7 @@ import zio._
 trait HandlerVersionSpecific {
   private[http] class ApplyContextAspect[-Env, +Err, -In, +Out, Env0](private val self: Handler[Env, Err, In, Out]) {
     def apply[Env1, Ctx: Tag, In1 <: In](aspect: HandlerAspect[Env1, Ctx])(implicit
-      in: Handler.IsRequest[In1],
+      ev0: Request <:< In,
       ev: Env0 with Ctx <:< Env,
       out: Out <:< Response,
       err: Err <:< Response,
diff --git a/zio-http/shared/src/main/scala-2/zio/http/RouteCompanionVersionSpecific.scala b/zio-http/shared/src/main/scala-2/zio/http/RouteCompanionVersionSpecific.scala
new file mode 100644
index 00000000000..171355a7313
--- /dev/null
+++ b/zio-http/shared/src/main/scala-2/zio/http/RouteCompanionVersionSpecific.scala
@@ -0,0 +1,12 @@
+package zio.http
+
+import zio.Tag
+
+trait RouteCompanionVersionSpecific {
+  private[http] class ApplyContextAspect[-Env, +Err, Env0](private val self: Route[Env, Err]) {
+    def apply[Env1, Env2 <: Env, Ctx: Tag](aspect: HandlerAspect[Env1, Ctx])(implicit
+      ev: Env0 with Ctx <:< Env,
+    ): Route[Env0 with Env1, Err] = self.transform(_.@@[Env0](aspect))
+  }
+
+}
diff --git a/zio-http/shared/src/main/scala-3/zio/http/HandlerVersionSpecific.scala b/zio-http/shared/src/main/scala-3/zio/http/HandlerVersionSpecific.scala
index 9ff1a6730ac..8352fe3b48d 100644
--- a/zio-http/shared/src/main/scala-3/zio/http/HandlerVersionSpecific.scala
+++ b/zio-http/shared/src/main/scala-3/zio/http/HandlerVersionSpecific.scala
@@ -5,7 +5,7 @@ import zio.*
 trait HandlerVersionSpecific {
   private[http] class ApplyContextAspect[-Env, +Err, -In, +Out, Env0](self: Handler[Env, Err, In, Out]) {
     transparent inline def apply[Env1, Ctx, In1 <: In](aspect: HandlerAspect[Env1, Ctx])(implicit
-      in: Handler.IsRequest[In1],
+      ev0: Request <:< In,
       ev: Env0 with Ctx <:< Env,
       out: Out <:< Response,
       err: Err <:< Response,
diff --git a/zio-http/shared/src/main/scala-3/zio/http/RouteCompanionVersionSpecific.scala b/zio-http/shared/src/main/scala-3/zio/http/RouteCompanionVersionSpecific.scala
new file mode 100644
index 00000000000..3122762c10d
--- /dev/null
+++ b/zio-http/shared/src/main/scala-3/zio/http/RouteCompanionVersionSpecific.scala
@@ -0,0 +1,10 @@
+package zio.http
+
+trait RoutesCompanionVersionSpecific {
+  private[http] class ApplyContextAspect[-Env, +Err, Env0](private val self: Route[Env, Err]) {
+    transparent inline def apply[Env1, Env2 <: Env, Ctx](aspect: HandlerAspect[Env1, Ctx])(implicit
+      ev: Env0 with Ctx <:< Env,
+    ): Route[Env0 with Env1, Err] = self.transform(_.@@[Env0](aspect))
+  }
+
+}
diff --git a/zio-http/shared/src/main/scala/zio/http/Handler.scala b/zio-http/shared/src/main/scala/zio/http/Handler.scala
index 2bbdae0034e..d09a7aa0af9 100644
--- a/zio-http/shared/src/main/scala/zio/http/Handler.scala
+++ b/zio-http/shared/src/main/scala/zio/http/Handler.scala
@@ -35,7 +35,7 @@ import zio.http.template._
 sealed trait Handler[-R, +Err, -In, +Out] { self =>
 
   def @@[Env1 <: R, In1 <: In](aspect: HandlerAspect[Env1, Unit])(implicit
-    in: Handler.IsRequest[In1],
+    ev: Request <:< In,
     out: Out <:< Response,
     err: Err <:< Response,
   ): Handler[Env1, Response, Request, Response] = {
@@ -46,7 +46,7 @@ sealed trait Handler[-R, +Err, -In, +Out] { self =>
   }
 
   def @@[Env0, Ctx <: R, In1 <: In](aspect: HandlerAspect[Env0, Ctx])(implicit
-    in: Handler.IsRequest[In1],
+    ev: Request <:< In,
     out: Out <:< Response,
     err: Err <:< Response,
     trace: Trace,
diff --git a/zio-http/shared/src/main/scala/zio/http/HandlerAspect.scala b/zio-http/shared/src/main/scala/zio/http/HandlerAspect.scala
index 7f28352559d..647f2f5e6cb 100644
--- a/zio-http/shared/src/main/scala/zio/http/HandlerAspect.scala
+++ b/zio-http/shared/src/main/scala/zio/http/HandlerAspect.scala
@@ -103,6 +103,26 @@ final case class HandlerAspect[-Env, +CtxOut](
       }
     }
 
+  /**
+   * Applies middleware to the specified handler, which may ignore the context
+   * produced by this middleware.
+   */
+  override def apply[Env1 <: Env, Err](
+    routes: Route[Env1, Err],
+  ): Route[Env1, Err] =
+    routes.transform[Env1] { handler =>
+      if (self == HandlerAspect.identity) handler
+      else {
+        for {
+          tuple <- protocol.incomingHandler
+          (state, (request, ctxOut)) = tuple
+          either   <- Handler.fromZIO(handler(request)).either
+          response <- Handler.fromZIO(protocol.outgoingHandler((state, either.merge)))
+          response <- if (either.isLeft) Handler.fail(response) else Handler.succeed(response)
+        } yield response
+      }
+    }
+
   /**
    * Applies middleware to the specified handler, which must process the context
    * produced by this middleware.
diff --git a/zio-http/shared/src/main/scala/zio/http/Middleware.scala b/zio-http/shared/src/main/scala/zio/http/Middleware.scala
index 9898968c238..f47809ad2a9 100644
--- a/zio-http/shared/src/main/scala/zio/http/Middleware.scala
+++ b/zio-http/shared/src/main/scala/zio/http/Middleware.scala
@@ -30,6 +30,9 @@ trait Middleware[-UpperEnv] { self =>
 
   def apply[Env1 <: UpperEnv, Err](routes: Routes[Env1, Err]): Routes[Env1, Err]
 
+  def apply[Env1 <: UpperEnv, Err](route: Route[Env1, Err]): Route[Env1, Err] =
+    (route.toRoutes @@ self).routes.head
+
   def @@[UpperEnv1 <: UpperEnv](
     that: Middleware[UpperEnv1],
   ): Middleware[UpperEnv1] =
diff --git a/zio-http/shared/src/main/scala/zio/http/RequestHandlerInput.scala.scala b/zio-http/shared/src/main/scala/zio/http/RequestHandlerInput.scala.scala
index a9db787330a..37ab5c5643a 100644
--- a/zio-http/shared/src/main/scala/zio/http/RequestHandlerInput.scala.scala
+++ b/zio-http/shared/src/main/scala/zio/http/RequestHandlerInput.scala.scala
@@ -22,11 +22,24 @@ import zio.Zippable
 
 @implicitNotFound("""
 Your request handler is required to accept both parameters ${A}, as well as the incoming [[zio.http.Request]].
-This is true even if you wish to ignore some parameters or the request itself. Try to add missing parameters 
-until you no longer receive this error message. If all else fails, you can construct a handler manually using 
+This is true even if you wish to ignore some parameters or the request itself. Try to add missing parameters
+until you no longer receive this error message. If all else fails, you can construct a handler manually using
 the constructors in the companion object of [[zio.http.Handler]] using the precise types.""")
 final class RequestHandlerInput[A, I](val zippable: Zippable.Out[A, Request, I])
 object RequestHandlerInput {
   implicit def apply[A, I](implicit zippable: Zippable.Out[A, Request, I]): RequestHandlerInput[A, I] =
     new RequestHandlerInput(zippable)
 }
+
+@implicitNotFound("""
+Your request handler is required to accept both parameters ${A}, as well as the incoming [[zio.http.Request]].
+This is true even if you wish to ignore some parameters or the request itself. Try to add missing parameters
+until you no longer receive this error message. If all else fails, you can construct a handler manually using
+the constructors in the companion object of [[zio.http.Handler]] using the precise types.""")
+final class RequestHandlerInputWithContext[A, I, Ctx](val zippable: Zippable.Out[A, (Ctx, Request), I])
+object RequestHandlerInputWithContext {
+  implicit def apply[A, I, Ctx](implicit
+    zippable: Zippable.Out[A, (Ctx, Request), I],
+  ): RequestHandlerInputWithContext[A, I, Ctx] =
+    new RequestHandlerInputWithContext(zippable)
+}
diff --git a/zio-http/shared/src/main/scala/zio/http/Route.scala b/zio-http/shared/src/main/scala/zio/http/Route.scala
index 9db2d28d13b..fbaa7ea03d5 100644
--- a/zio-http/shared/src/main/scala/zio/http/Route.scala
+++ b/zio-http/shared/src/main/scala/zio/http/Route.scala
@@ -32,7 +32,21 @@ import zio.http.codec.PathCodec
  * Individual routes can be aggregated using [[zio.http.Routes]].
  */
 sealed trait Route[-Env, +Err] { self =>
-  import Route.{Augmented, Handled, Provided, Unhandled}
+  import Route._
+
+  def @@[Env1 <: Env](aspect: Middleware[Env1]): Route[Env1, Err] =
+    aspect(self)
+
+  def @@[Env0](aspect: HandlerAspect[Env0, Unit]): Route[Env with Env0, Err] =
+    aspect(self)
+
+  def @@[Env0, Ctx <: Env](
+    aspect: HandlerAspect[Env0, Ctx],
+  )(implicit tag: Tag[Ctx]): Route[Env0, Err] =
+    self.transform(_ @@ aspect)
+
+  def @@[Env0]: ApplyContextAspect[Env, Err, Env0] =
+    new ApplyContextAspect[Env, Err, Env0](self)
 
   /**
    * Applies the route to the specified request. The route must be defined for
@@ -332,7 +346,7 @@ sealed trait Route[-Env, +Err] { self =>
   ): Route[Env1, Err] =
     Route.Augmented(self, f)
 }
-object Route                   {
+object Route extends RouteCompanionVersionSpecific {
 
   def handledIgnoreParams[Env](
     routePattern: RoutePattern[_],