Move rejection handler to DriverRoute

1 files changed, 15 insertions, 2 deletions

diff --git a/src/main/scala/xyz/driver/core/rest/DriverRoute.scala b/src/main/scala/xyz/driver/core/rest/DriverRoute.scala
index 1fe5e3f..5e629be 100644
--- a/src/main/scala/xyz/driver/core/rest/DriverRoute.scala
+++ b/src/main/scala/xyz/driver/core/rest/DriverRoute.scala
@@ -6,7 +6,7 @@ import akka.http.scaladsl.model._
 import akka.http.scaladsl.model.StatusCodes
 import akka.http.scaladsl.model.headers._
 import akka.http.scaladsl.server.Directives._
-import akka.http.scaladsl.server.{Directive0, ExceptionHandler, RequestContext, Route}
+import akka.http.scaladsl.server._
 import com.typesafe.config.Config
 import com.typesafe.scalalogging.Logger
 import org.slf4j.MDC
@@ -22,7 +22,7 @@ trait DriverRoute {
   def route: Route
 
   def routeWithDefaults: Route = {
-    (defaultResponseHeaders & handleExceptions(ExceptionHandler(exceptionHandler))) {
+    (defaultResponseHeaders & handleRejections(rejectionHandler) & handleExceptions(ExceptionHandler(exceptionHandler))) {
       route ~ defaultOptionsRoute
     }
   }
@@ -84,6 +84,19 @@ trait DriverRoute {
     }
   }
 
+  protected def rejectionHandler: RejectionHandler =
+    RejectionHandler
+      .newBuilder()
+      .handle {
+        case rejection =>
+          respondWithAllCorsHeaders {
+            RejectionHandler
+              .default(scala.collection.immutable.Seq(rejection))
+              .getOrElse(complete("OK"))
+          }
+      }
+      .result()
+
   /**
     * Override me for custom exception handling
     *