Replace pseudofunction mechanism with JS sandboxing

Tinkoff · Apr 4, 2023 · cc22bd3 · cc22bd3
1 parent bc541c0
commit cc22bd3
Show file tree

Hide file tree

Showing 14 changed files with 85 additions and 184 deletions.
diff --git a/backend/mockingbird-api/src/main/scala/ru/tinkoff/tcb/mockingbird/Mockingbird.scala b/backend/mockingbird-api/src/main/scala/ru/tinkoff/tcb/mockingbird/Mockingbird.scala
@@ -21,7 +21,6 @@ import sttp.client3.armeria.zio.ArmeriaZioBackend
 import tofu.logging.Logging
 import tofu.logging.impl.ZUniversalLogging
 import zio.managed.*
-
 import ru.tinkoff.tcb.mockingbird.api.AdminApiHandler
 import ru.tinkoff.tcb.mockingbird.api.AdminHttp
 import ru.tinkoff.tcb.mockingbird.api.MetricsHttp
@@ -51,6 +50,8 @@ import ru.tinkoff.tcb.mockingbird.stream.EphemeralCleaner
 import ru.tinkoff.tcb.mockingbird.stream.EventSpawner
 import ru.tinkoff.tcb.mockingbird.stream.SDFetcher
 import ru.tinkoff.tcb.utils.metrics.makeRegistry
+import ru.tinkoff.tcb.utils.resource.readStr
+import ru.tinkoff.tcb.utils.sandboxing.GraalJsSandbox
 
 object Mockingbird extends scala.App {
   type FL = WLD & ServerConfig & PublicHttp & EventSpawner & ResourceManager & EphemeralCleaner & GrpcRequestHandler
@@ -142,6 +143,7 @@ object Mockingbird extends scala.App {
             scopedBackend <- ArmeriaZioBackend.scopedUsingClient(webClient)
           } yield scopedBackend
         },
+        (ZLayer.service[ServerConfig].project(_.sandbox) ++ ZLayer.fromZIO(ZIO.attempt(readStr("prelude.js")).map(Option(_)))) >>> GraalJsSandbox.live,
         mongoLayer,
         aesEncoder,
         collection(_.stub) >>> HttpStubDAOImpl.live,
@@ -186,10 +188,12 @@ object Mockingbird extends scala.App {
           .exec(bytes)
           .provideSome[RequestContext](
             Tracing.live,
+            MockingbirdConfiguration.server,
             MockingbirdConfiguration.mongo,
             mongoLayer,
             collection(_.state) >>> PersistentStateDAOImpl.live,
             collection(_.grpcStub) >>> GrpcStubDAOImpl.live,
+            (ZLayer.service[ServerConfig].project(_.sandbox) ++ ZLayer.fromZIO(ZIO.attempt(readStr("prelude.js")).map(Option(_)))) >>> GraalJsSandbox.live,
             GrpcStubResolverImpl.live,
             GrpcRequestHandlerImpl.live
           )

diff --git a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/api/PublicApiHandler.scala b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/api/PublicApiHandler.scala
@@ -41,6 +41,7 @@ import ru.tinkoff.tcb.mockingbird.scenario.CallbackEngine
 import ru.tinkoff.tcb.mockingbird.scenario.ScenarioEngine
 import ru.tinkoff.tcb.utils.circe.optics.JsonOptic
 import ru.tinkoff.tcb.utils.regex.*
+import ru.tinkoff.tcb.utils.sandboxing.GraalJsSandbox
 import ru.tinkoff.tcb.utils.transformation.json.*
 import ru.tinkoff.tcb.utils.transformation.string.*
 import ru.tinkoff.tcb.utils.transformation.xml.*
@@ -54,6 +55,7 @@ final class PublicApiHandler(
     stateDAO: PersistentStateDAO[Task],
     resolver: StubResolver,
     engine: CallbackEngine,
+    implicit val jsSandbox: GraalJsSandbox,
     private val httpBackend: SttpBackend[Task, ?],
     proxyConfig: ProxyConfig
 ) {
@@ -304,8 +306,9 @@ object PublicApiHandler {
       ssd        <- ZIO.service[PersistentStateDAO[Task]]
       resolver   <- ZIO.service[StubResolver]
       engine     <- ZIO.service[ScenarioEngine]
+      jsSandbox  <- ZIO.service[GraalJsSandbox]
       sttpClient <- ZIO.service[SttpBackend[Task, Any]]
       proxyCfg   <- ZIO.service[ProxyConfig]
-    } yield new PublicApiHandler(hsd, ssd, resolver, engine, sttpClient, proxyCfg)
+    } yield new PublicApiHandler(hsd, ssd, resolver, engine, jsSandbox, sttpClient, proxyCfg)
   }
 }
diff --git a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/config/Configuration.scala b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/config/Configuration.scala
@@ -8,7 +8,9 @@ import net.ceedubs.ficus.Ficus.*
 import net.ceedubs.ficus.readers.ArbitraryTypeReader.*
 import net.ceedubs.ficus.readers.EnumerationReader.*
 
-case class ServerConfig(interface: String, port: Int, allowedOrigins: Seq[String], healthCheckRoute: Option[String])
+case class JsSandboxConfig(allowedClasses: Set[String] = Set())
+
+case class ServerConfig(interface: String, port: Int, allowedOrigins: Seq[String], healthCheckRoute: Option[String], sandbox: JsSandboxConfig)
 
 case class SecurityConfig(secret: String)
 

diff --git a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/grpc/GrpcRequestHandler.scala b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/grpc/GrpcRequestHandler.scala
@@ -18,13 +18,15 @@ import ru.tinkoff.tcb.mockingbird.model.FillResponse
 import ru.tinkoff.tcb.mockingbird.model.GProxyResponse
 import ru.tinkoff.tcb.mockingbird.model.PersistentState
 import ru.tinkoff.tcb.mockingbird.model.Scope
+import ru.tinkoff.tcb.utils.sandboxing.GraalJsSandbox
 import ru.tinkoff.tcb.utils.transformation.json.*
 
 trait GrpcRequestHandler {
   def exec(bytes: Array[Byte]): RIO[WLD & RequestContext, Array[Byte]]
 }
 
-class GrpcRequestHandlerImpl(stubResolver: GrpcStubResolver) extends GrpcRequestHandler {
+class GrpcRequestHandlerImpl(stubResolver: GrpcStubResolver, implicit val jsSandbox: GraalJsSandbox)
+    extends GrpcRequestHandler {
   override def exec(bytes: Array[Byte]): RIO[WLD & RequestContext, Array[Byte]] =
     for {
       context <- ZIO.service[RequestContext]
@@ -86,7 +88,8 @@ class GrpcRequestHandlerImpl(stubResolver: GrpcStubResolver) extends GrpcRequest
 }
 
 object GrpcRequestHandlerImpl {
-  val live: URLayer[GrpcStubResolver, GrpcRequestHandler] = ZLayer.fromFunction(new GrpcRequestHandlerImpl(_))
+  val live: URLayer[GrpcStubResolver & GraalJsSandbox, GrpcRequestHandlerImpl] =
+    ZLayer.fromFunction(new GrpcRequestHandlerImpl(_, _))
 }
 
 object GrpcRequestHandler {

diff --git a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/scenario/ScenarioEngine.scala b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/scenario/ScenarioEngine.scala
@@ -42,6 +42,7 @@ import ru.tinkoff.tcb.mockingbird.model.XMLCallbackRequest
 import ru.tinkoff.tcb.mockingbird.model.XmlOutput
 import ru.tinkoff.tcb.mockingbird.stream.SDFetcher
 import ru.tinkoff.tcb.utils.id.SID
+import ru.tinkoff.tcb.utils.sandboxing.GraalJsSandbox
 import ru.tinkoff.tcb.utils.transformation.json.*
 import ru.tinkoff.tcb.utils.transformation.string.*
 import ru.tinkoff.tcb.utils.transformation.xml.*
@@ -62,6 +63,7 @@ final class ScenarioEngine(
     stateDAO: PersistentStateDAO[Task],
     resolver: ScenarioResolver,
     fetcher: SDFetcher,
+    implicit val jsSandbox: GraalJsSandbox,
     private val httpBackend: SttpBackend[Task, ?]
 ) extends CallbackEngine {
   private val log = MDCLogging.`for`[WLD](this)
@@ -222,7 +224,8 @@ object ScenarioEngine {
       psd        <- ZIO.service[PersistentStateDAO[Task]]
       resolver   <- ZIO.service[ScenarioResolver]
       fetcher    <- ZIO.service[SDFetcher]
+      jsSandbox  <- ZIO.service[GraalJsSandbox]
       sttpClient <- ZIO.service[SttpBackend[Task, Any]]
-    } yield new ScenarioEngine(sd, psd, resolver, fetcher, sttpClient)
+    } yield new ScenarioEngine(sd, psd, resolver, fetcher, jsSandbox, sttpClient)
   }
 }
diff --git a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/sandboxing/ClassAccessRule.scala b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/sandboxing/ClassAccessRule.scala
diff --git a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/sandboxing/GraalJsSandbox.scala b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/sandboxing/GraalJsSandbox.scala
@@ -7,22 +7,21 @@ import scala.util.Using
 
 import org.graalvm.polyglot.*
 
-import ru.tinkoff.tcb.utils.instances.predicate.or.*
+import ru.tinkoff.tcb.mockingbird.config.JsSandboxConfig
 
 class GraalJsSandbox(
-    classAccessRules: List[ClassAccessRule] = GraalJsSandbox.DefaultAccess,
+    jsSandboxConfig: JsSandboxConfig,
     prelude: Option[String] = None
 ) {
-  private val accessRule = classAccessRules.asInstanceOf[List[String => Boolean]].combineAll
-
-  private val preludeSource = prelude.map(Source.create("js", _))
+  private val allowedClasses = GraalJsSandbox.DefaultAccess ++ jsSandboxConfig.allowedClasses
+  private val preludeSource  = prelude.map(Source.create("js", _))
 
   def eval[T: ClassTag](code: String, environment: Map[String, Any] = Map.empty): Try[T] =
     Using(
       Context
         .newBuilder("js")
         .allowHostAccess(HostAccess.ALL)
-        .allowHostClassLookup((t: String) => accessRule(t))
+        .allowHostClassLookup((t: String) => allowedClasses(t))
         .option("engine.WarnInterpreterOnly", "false")
         .build()
     ) { context =>
@@ -36,24 +35,31 @@ class GraalJsSandbox(
 }
 
 object GraalJsSandbox {
-  val DefaultAccess: List[ClassAccessRule] = List(
-    ClassAccessRule.Exact("java.lang.Byte"),
-    ClassAccessRule.Exact("java.lang.Boolean"),
-    ClassAccessRule.Exact("java.lang.Double"),
-    ClassAccessRule.Exact("java.lang.Float"),
-    ClassAccessRule.Exact("java.lang.Integer"),
-    ClassAccessRule.Exact("java.lang.Long"),
-    ClassAccessRule.Exact("java.lang.Math"),
-    ClassAccessRule.Exact("java.lang.Short"),
-    ClassAccessRule.Exact("java.lang.String"),
-    ClassAccessRule.Exact("java.math.BigDecimal"),
-    ClassAccessRule.Exact("java.math.BigInteger"),
-    ClassAccessRule.Exact("java.time.LocalDate"),
-    ClassAccessRule.Exact("java.time.LocalDateTime"),
-    ClassAccessRule.Exact("java.time.format.DateTimeFormatter"),
-    ClassAccessRule.Exact("java.util.List"),
-    ClassAccessRule.Exact("java.util.Map"),
-    ClassAccessRule.Exact("java.util.Random"),
-    ClassAccessRule.Exact("java.util.Set")
+  val live: URLayer[Option[String] & JsSandboxConfig, GraalJsSandbox] = ZLayer {
+    for {
+      sandboxConfig <- ZIO.service[JsSandboxConfig]
+      prelude       <- ZIO.service[Option[String]]
+    } yield new GraalJsSandbox(sandboxConfig, prelude)
+  }
+
+  val DefaultAccess: Set[String] = Set(
+    "java.lang.Byte",
+    "java.lang.Boolean",
+    "java.lang.Double",
+    "java.lang.Float",
+    "java.lang.Integer",
+    "java.lang.Long",
+    "java.lang.Math",
+    "java.lang.Short",
+    "java.lang.String",
+    "java.math.BigDecimal",
+    "java.math.BigInteger",
+    "java.time.LocalDate",
+    "java.time.LocalDateTime",
+    "java.time.format.DateTimeFormatter",
+    "java.util.List",
+    "java.util.Map",
+    "java.util.Random",
+    "java.util.Set"
   )
 }
diff --git a/...mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/json/js_eval/package.scala b/...mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/json/js_eval/package.scala
@@ -1,6 +1,7 @@
 package ru.tinkoff.tcb.utils.transformation.json
 
-import java.lang.Boolean as JBoolean
+import java.lang as jl
+import java.math as jm
 import scala.jdk.CollectionConverters.*
 
 import io.circe.Json
@@ -11,7 +12,7 @@ package object js_eval {
   val circe2js: Json.Folder[AnyRef] = new Json.Folder[AnyRef] {
     override def onNull: AnyRef = null
 
-    override def onBoolean(value: Boolean): AnyRef = JBoolean.valueOf(value)
+    override def onBoolean(value: Boolean): AnyRef = jl.Boolean.valueOf(value)
 
     override def onNumber(value: JsonNumber): AnyRef = value.toBigDecimal.map(_.bigDecimal).orNull
 
@@ -25,4 +26,12 @@ package object js_eval {
         .toMap
         .asJava
   }
+
+  val fold2Json: PartialFunction[AnyRef, Json] = {
+    case b: jl.Boolean     => Json.fromBoolean(b)
+    case s: String         => Json.fromString(s)
+    case bd: jm.BigDecimal => Json.fromBigDecimal(bd)
+    case i: jl.Integer     => Json.fromInt(i.intValue())
+    case l: jl.Long        => Json.fromLong(l.longValue())
+  }
 }
diff --git a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/json/package.scala b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/json/package.scala
@@ -1,7 +1,5 @@
 package ru.tinkoff.tcb.utils.transformation
 
-import java.lang as jl
-import java.math as jm
 import scala.util.Failure
 import scala.util.Success
 import scala.util.control.TailCalls
@@ -10,13 +8,13 @@ import scala.util.control.TailCalls.TailRec
 import io.circe.Json
 import io.circe.JsonNumber as JNumber
 import kantan.xpath.*
-import mouse.boolean.*
 
 import ru.tinkoff.tcb.utils.circe.*
 import ru.tinkoff.tcb.utils.circe.optics.JsonOptic
 import ru.tinkoff.tcb.utils.json.json2StringFolder
 import ru.tinkoff.tcb.utils.regex.OneOrMore
 import ru.tinkoff.tcb.utils.sandboxing.GraalJsSandbox
+import ru.tinkoff.tcb.utils.transformation.json.js_eval.fold2Json
 import ru.tinkoff.tcb.utils.transformation.xml.nodeTemplater
 
 package object json {
@@ -91,39 +89,13 @@ package object json {
         }.result
       }
 
-    def eval: Json =
-      transformValues { case js @ JsonString(str) =>
-        str
-          .foldTemplate(
-            Json.fromString,
-            Json.fromInt,
-            Json.fromLong
-          )
-          .orElse {
-            FunRx
-              .findFirstIn(str)
-              .isDefined
-              .option(
-                FunRx
-                  .replaceSomeIn(str, m => m.matched.foldTemplate(identity, _.toString(), _.toString()))
-              )
-              .map(Json.fromString)
-          }
-          .getOrElse(js)
-      }.result
-
-    def eval2(implicit sandbox: GraalJsSandbox): Json =
-      transformValues {
-        case js @ JsonString(CodeRx(code)) =>
-          (sandbox.eval[AnyRef](code) match {
-            case Success(str: String)       => Option(Json.fromString(str))
-            case Success(bd: jm.BigDecimal) => Option(Json.fromBigDecimal(bd))
-            case Success(i: jl.Integer)     => Option(Json.fromInt(i.intValue()))
-            case Success(l: jl.Long)        => Option(Json.fromLong(l.longValue()))
-            case Success(other)             => throw new Exception(s"${other.getClass.getCanonicalName}: $other")
-            case Failure(exception)         => throw exception
-          }).getOrElse(js)
-        case JsonString(other) => throw new Exception(other)
+    def eval(implicit sandbox: GraalJsSandbox): Json =
+      transformValues { case js @ JsonString(CodeRx(code)) =>
+        (sandbox.eval[AnyRef](code) match {
+          case Success(value) if fold2Json.isDefinedAt(value) => Option(fold2Json(value))
+          case Success(other)     => throw new Exception(s"${other.getClass.getCanonicalName}: $other")
+          case Failure(exception) => throw exception
+        }).getOrElse(js)
       }.result
 
     def patch(values: Json, schema: Map[JsonOptic, String]): Json =

diff --git a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/package.scala b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/package.scala
@@ -1,13 +1,7 @@
 package ru.tinkoff.tcb.utils
 
-import java.time.LocalDate
-import java.time.LocalDateTime
-import java.util.UUID
-import scala.util.Random
 import scala.util.matching.Regex
 
-import ru.tinkoff.tcb.utils.time.*
-
 package object transformation {
   val SubstRx: Regex = """\$\{(.*?)\}""".r
   val FunRx: Regex   = """%\{.*?\}""".r
@@ -24,35 +18,4 @@ package object transformation {
 
   val Today: Regex = """%\{today\((.*?)\)\}""".r
   val Now: Regex   = """%\{now\((.*?)\)\}""".r
-
-  implicit final class TemplateTransformations(private val template: String) extends AnyVal {
-    def foldTemplate[T](foldString: String => T, foldInt: Int => T, foldLong: Long => T): Option[T] =
-      template match {
-        case RandStr(len) =>
-          Some(foldString(Random.alphanumeric.take(len.toInt).mkString))
-        case RandAlphabetStr(alphabet, minLen, maxLen) =>
-          Some(
-            foldString(
-              List.fill(Random.between(minLen.toInt, maxLen.toInt))(alphabet(Random.nextInt(alphabet.length))).mkString
-            )
-          )
-        case RandNumStr(len) =>
-          Some(foldString(Seq.fill(len.toInt)(Random.nextInt(10)).mkString))
-        case RandInt(max) =>
-          Some(foldInt(Random.nextInt(max.toInt)))
-        case RandIntInterval(min, max) =>
-          Some(foldInt(Random.between(min.toInt, max.toInt)))
-        case RandLong(max) =>
-          Some(foldLong(Random.nextLong(max.toLong)))
-        case RandLongInterval(min, max) =>
-          Some(foldLong(Random.between(min.toLong, max.toLong)))
-        case RandUUID() =>
-          Some(foldString(UUID.randomUUID().toString))
-        case Today(Formatter(fmt)) =>
-          Some(foldString(LocalDate.now().format(fmt)))
-        case Now(Formatter(fmt)) =>
-          Some(foldString(LocalDateTime.now().format(fmt)))
-        case _ => None
-      }
-  }
 }
diff --git a/backend/mockingbird/src/test/scala/ru/tinkoff/tcb/utils/sandboxing/GraalJsSandboxSpec.scala b/backend/mockingbird/src/test/scala/ru/tinkoff/tcb/utils/sandboxing/GraalJsSandboxSpec.scala
@@ -5,8 +5,10 @@ import org.scalatest.TryValues
 import org.scalatest.funsuite.AnyFunSuite
 import org.scalatest.matchers.should.Matchers
 
+import ru.tinkoff.tcb.mockingbird.config.JsSandboxConfig
+
 class GraalJsSandboxSpec extends AnyFunSuite with Matchers with TryValues {
-  private val sandbox = new GraalJsSandbox
+  private val sandbox = new GraalJsSandbox(new JsSandboxConfig())
 
   test("Eval simple arithmetics") {
     sandbox.eval[Int]("1 + 2").success.value shouldBe 3