From 0ef027120ec5857ce4c2a86c2aad5a7d925cf573 Mon Sep 17 00:00:00 2001
From: Stewart Stewart <stewinsalot@gmail.com>
Date: Sun, 19 Mar 2017 16:54:55 -0400
Subject: add directive for sanitizing script tags request entities

---
 src/main/scala/xyz/driver/core/rest.scala | 37 +++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

(limited to 'src/main/scala/xyz/driver/core/rest.scala')

diff --git a/src/main/scala/xyz/driver/core/rest.scala b/src/main/scala/xyz/driver/core/rest.scala
index 9923ef8..17837e6 100644
--- a/src/main/scala/xyz/driver/core/rest.scala
+++ b/src/main/scala/xyz/driver/core/rest.scala
@@ -5,9 +5,12 @@ import akka.http.scaladsl.Http
 import akka.http.scaladsl.model._
 import akka.http.scaladsl.model.headers.{HttpChallenges, RawHeader}
 import akka.http.scaladsl.server.AuthenticationFailedRejection.CredentialsRejected
+import akka.http.scaladsl.server.Directive0
 import akka.http.scaladsl.unmarshalling.Unmarshal
 import akka.http.scaladsl.unmarshalling.Unmarshaller
 import akka.stream.ActorMaterializer
+import akka.stream.scaladsl.Flow
+import akka.util.ByteString
 import com.github.swagger.akka.model._
 import com.github.swagger.akka.{HasActorSystem, SwaggerHttpService}
 import com.typesafe.config.Config
@@ -64,6 +67,40 @@ object rest {
     } toMap
   }
 
+  private def escapeScriptTags(byteString: ByteString): ByteString = {
+    def dirtyIndices(from: Int, descIndices: List[Int]): List[Int] = {
+      val index = byteString.indexOf('/', from)
+      if (index === -1) descIndices.reverse
+      else {
+        val (init, tail) = byteString.splitAt(index)
+        if ((init endsWith "<") && (tail startsWith "/sc")) {
+          dirtyIndices(index + 1, index :: descIndices)
+        } else {
+          dirtyIndices(index + 1, descIndices)
+        }
+      }
+    }
+
+    val firstSlash = byteString.indexOf('/')
+    if (firstSlash === -1) byteString
+    else {
+      val indices = dirtyIndices(firstSlash, Nil) :+ byteString.length
+      val builder = ByteString.newBuilder
+      builder ++= byteString.take(firstSlash)
+      indices.sliding(2).foreach {
+        case Seq(start, end) =>
+          builder += ' '
+          builder ++= byteString.slice(start, end)
+      }
+      builder.result
+    }
+  }
+
+  val sanitizeRequestEntity: Directive0 = {
+    mapRequest(
+      request => request.mapEntity(entity => entity.transformDataBytes(Flow.fromFunction(escapeScriptTags))))
+  }
+
   object AuthProvider {
     val AuthenticationTokenHeader    = ContextHeaders.AuthenticationTokenHeader
     val SetAuthenticationTokenHeader = "set-authorization"
-- 
cgit v1.2.3