API-1595 Move rejectionHandler to DriverApp to avoid 404 rejection issues (#167)v1.9.1

2 files changed, 11 insertions, 12 deletions

diff --git a/src/main/scala/xyz/driver/core/app/DriverApp.scala b/src/main/scala/xyz/driver/core/app/DriverApp.scala
index bf18c3d..56ee898 100644
--- a/src/main/scala/xyz/driver/core/app/DriverApp.scala
+++ b/src/main/scala/xyz/driver/core/app/DriverApp.scala
@@ -127,7 +127,7 @@ class DriverApp(
     val combinedRoute =
       Route.seal(modules.map(_.route).foldLeft(basicRoutes.routeWithDefaults)(_ ~ _) ~ defaultOptionsRoute)
 
-    (extractHost & extractClientIP & trace(tracer)) {
+    (extractHost & extractClientIP & trace(tracer) & handleRejections(authenticationRejectionHandler)) {
       case (origin, ip) =>
         ctx =>
           val trackingId = extractTrackingId(ctx.request)
@@ -153,6 +153,15 @@ class DriverApp(
     }
   }
 
+  protected def authenticationRejectionHandler: RejectionHandler =
+    RejectionHandler
+      .newBuilder()
+      .handle {
+        case AuthenticationFailedRejection(_, challenge) =>
+          complete(HttpResponse(StatusCodes.Unauthorized, entity = challenge.realm))
+      }
+      .result()
+
   protected def bindHttp(modules: Seq[Module]): Unit = {
     val _ = http.bindAndHandle(route2HandlerFlow(appRoute), interface, port)(materializer)
   }
diff --git a/src/main/scala/xyz/driver/core/rest/DriverRoute.scala b/src/main/scala/xyz/driver/core/rest/DriverRoute.scala
index 538d918..c23c44d 100644
--- a/src/main/scala/xyz/driver/core/rest/DriverRoute.scala
+++ b/src/main/scala/xyz/driver/core/rest/DriverRoute.scala
@@ -19,8 +19,7 @@ trait DriverRoute {
   def route: Route
 
   def routeWithDefaults: Route = {
-    (defaultResponseHeaders & handleExceptions(ExceptionHandler(exceptionHandler)) & handleRejections(
-      authenticationRejectionHandler)) {
+    (defaultResponseHeaders & handleExceptions(ExceptionHandler(exceptionHandler))) {
       route
     }
   }
@@ -98,13 +97,4 @@ trait DriverRoute {
   protected def errorResponse[T <: Exception](statusCode: StatusCode, message: String, exception: T): Route = {
     complete(HttpResponse(statusCode, entity = message))
   }
-
-  protected def authenticationRejectionHandler: RejectionHandler =
-    RejectionHandler
-      .newBuilder()
-      .handle {
-        case AuthenticationFailedRejection(_, challenge) =>
-          complete(HttpResponse(StatusCodes.Unauthorized, entity = challenge.realm))
-      }
-      .result()
 }