From 04db4e857fceeb15196d8f13d63e987ad214be38 Mon Sep 17 00:00:00 2001
From: Zach Smith <zach@driver.xyz>
Date: Thu, 8 Feb 2018 13:02:07 -0800
Subject: Move rejection handler to DriverRoute

---
 src/main/scala/xyz/driver/core/rest/DriverRoute.scala | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

(limited to 'src/main/scala')

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
     *
-- 
cgit v1.2.3