Respond with correct cors headers for all rejections

1 files changed, 28 insertions, 10 deletions

diff --git a/src/main/scala/xyz/driver/core/app/DriverApp.scala b/src/main/scala/xyz/driver/core/app/DriverApp.scala
index df80c3d..bc9d634 100644
--- a/src/main/scala/xyz/driver/core/app/DriverApp.scala
+++ b/src/main/scala/xyz/driver/core/app/DriverApp.scala
@@ -223,6 +223,25 @@ class DriverApp(
 }
 
 object DriverApp {
+  def defaultCorsHeaders: Directive0 = {
+    optionalHeaderValueByType[Origin](()) flatMap { originHeader =>
+      respondWithHeaders(
+        List[HttpHeader](
+          allowOrigin(originHeader),
+          `Access-Control-Allow-Headers`(AllowedHeaders: _*),
+          `Access-Control-Expose-Headers`(AllowedHeaders: _*)
+        ))
+    }
+  }
+
+  def corsAllowedMethodHeaders(methods: scala.collection.immutable.Seq[HttpMethod]): Directive0 = {
+    respondWithHeaders(
+      List[HttpHeader](
+        Allow(methods),
+        `Access-Control-Allow-Methods`(methods)
+      ))
+  }
+
   implicit def rejectionHandler: RejectionHandler =
     RejectionHandler
       .newBuilder()
@@ -230,20 +249,19 @@ object DriverApp {
         val methods    = rejections map (_.supported)
         lazy val names = methods map (_.name) mkString ", "
 
-        options { ctx =>
-          optionalHeaderValueByType[Origin](()) { originHeader =>
-            respondWithHeaders(List[HttpHeader](
-              Allow(methods),
-              `Access-Control-Allow-Methods`(methods),
-              allowOrigin(originHeader),
-              `Access-Control-Allow-Headers`(AllowedHeaders: _*),
-              `Access-Control-Expose-Headers`(AllowedHeaders: _*)
-            )) {
+        options {
+          defaultCorsHeaders {
+            corsAllowedMethodHeaders(methods) {
               complete(s"Supported methods: $names.")
             }
-          }(ctx)
+          }
         } ~
           complete(MethodNotAllowed -> s"HTTP method not allowed, supported methods: $names!")
       }
+      .handleAll[Rejection] { rejections =>
+        defaultCorsHeaders {
+          RejectionHandler.default(rejections).getOrElse(reject)
+        }
+      }
       .result()
 }