Replace pseudofunction mechanism with JS sandboxing

Author: danslapman

backend/mockingbird-api/src/main/scala/ru/tinkoff/tcb/mockingbird/Mockingbird.scala

@@ -5,7 +5,6 @@ import scalapb.zio_grpc.Server
 import scalapb.zio_grpc.ServerLayer
 import scalapb.zio_grpc.ServiceList
 import scalapb.zio_grpc.server.ZServerCallHandler
-
 import com.mongodb.ConnectionString
 import io.grpc.ServerBuilder
 import io.grpc.Status
@@ -18,7 +17,6 @@ import sttp.client3.asynchttpclient.zio.AsyncHttpClientZioBackend
 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
@@ -48,6 +46,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
@@ -122,6 +122,7 @@ object Mockingbird extends scala.App {
             )
           } yield scopedBackend
         },
+        (ZLayer.service[ServerConfig].project(_.sandbox) ++ ZLayer.fromZIO(ZIO.attempt(readStr("prelude.js")).map(Option(_)))) >>> GraalJsSandbox.live,
         mongoLayer,
         aesEncoder,
         collection(_.stub) >>> HttpStubDAOImpl.live,
@@ -166,10 +167,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
           )

backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/api/PublicApiHandler.scala

@@ -3,7 +3,6 @@ package ru.tinkoff.tcb.mockingbird.api
 import scala.concurrent.duration.FiniteDuration
 import scala.util.control.NonFatal
 import scala.xml.Node
-
 import io.circe.Json
 import io.circe.parser.parse
 import io.circe.syntax.*
@@ -13,7 +12,6 @@ import sttp.client3.*
 import sttp.client3.circe.*
 import sttp.model.Method
 import zio.interop.catz.core.*
-
 import ru.tinkoff.tcb.criteria.*
 import ru.tinkoff.tcb.criteria.Typed.*
 import ru.tinkoff.tcb.logging.MDCLogging
@@ -37,6 +35,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.*
@@ -50,6 +49,7 @@ final class PublicApiHandler(
     stateDAO: PersistentStateDAO[Task],
     resolver: StubResolver,
     engine: CallbackEngine,
+    implicit val jsSandbox: GraalJsSandbox,
     private val httpBackend: SttpBackend[Task, ?],
     proxyConfig: ProxyConfig
 ) {
@@ -243,8 +243,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)
   }
 }

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])
+case class JsSandboxConfig(allowedClasses: Set[String] = Set())
+
+case class ServerConfig(interface: String, port: Int, allowedOrigins: Seq[String], sandbox: JsSandboxConfig)
 
 case class SecurityConfig(secret: String)
 

backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/grpc/GrpcRequestHandler.scala

@@ -3,13 +3,11 @@ package ru.tinkoff.tcb.mockingbird.grpc
 import scalapb.zio_grpc.RequestContext
 import scalapb.zio_grpc.ZManagedChannel
 import scalapb.zio_grpc.client.ClientCalls
-
 import io.circe.Json
 import io.circe.syntax.KeyOps
 import io.grpc.CallOptions
 import io.grpc.ManagedChannelBuilder
 import zio.Duration
-
 import ru.tinkoff.tcb.mockingbird.api.Tracing
 import ru.tinkoff.tcb.mockingbird.api.WLD
 import ru.tinkoff.tcb.mockingbird.error.StubSearchError
@@ -18,13 +16,14 @@ 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 +85,7 @@ 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 {

backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/scenario/ScenarioEngine.scala

@@ -2,7 +2,6 @@ package ru.tinkoff.tcb.mockingbird.scenario
 
 import java.nio.charset.Charset
 import java.util.Base64
-
 import io.circe.Json
 import io.circe.syntax.*
 import kantan.xpath.Node as KNode
@@ -13,7 +12,6 @@ import sttp.client3.*
 import sttp.client3.circe.*
 import sttp.model.Method
 import zio.interop.catz.core.*
-
 import ru.tinkoff.tcb.criteria.*
 import ru.tinkoff.tcb.criteria.Typed.*
 import ru.tinkoff.tcb.logging.MDCLogging
@@ -42,6 +40,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 +61,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)
@@ -217,7 +217,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)
   }
 }

backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/sandboxing/ClassAccessRule.scala

@@ -4,25 +4,22 @@ import scala.reflect.ClassTag
 import scala.reflect.classTag
 import scala.util.Try
 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 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 +33,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"
   )
 }

backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/sandboxing/GraalJsSandbox.scala

@@ -1,6 +1,7 @@
 package ru.tinkoff.tcb.utils.transformation.json
 
-import java.lang.Boolean as JBoolean
+import java.math as jm
+import java.lang as jl
 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())
+  }
 }

backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/json/js_eval/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,14 +8,14 @@ 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.xml.nodeTemplater
+import ru.tinkoff.tcb.utils.transformation.json.js_eval.fold2Json
 
 package object json {
   private val JORx = """\$([\:~])?\{([\p{L}\d\.\[\]\-_]+)\}""".r
@@ -81,39 +79,14 @@ 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 =
+    def eval(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(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)
-        case JsonString(other) => throw new Exception(other)
       }.result
 
     def patch(values: Json, schema: Map[JsonOptic, String]): Json =