fix 'runBefore' in server mode (#180)

root issue of joernio/joern#4999

Author: mpollmeier

README.md

@@ -451,6 +451,9 @@ curl http://localhost:8080/query-sync -X POST -d '{"query": "val bar = foo + 1"}
 
 curl http://localhost:8080/query-sync -X POST -d '{"query":"println(\"OMG remote code execution!!1!\")"}'
 # {"success":true,"stdout":"",...}%
+
+# combine with `jq` to directly get the output, as if you had a local console:
+curl --silent http://localhost:8080/query-sync -X POST -d '{"query": "val baz = 43"}' | jq --raw-output .stdout
 ```
 
 The same for windows and powershell:

build.sbt

@@ -5,6 +5,7 @@ ThisBuild / organization := "com.michaelpollmeier"
 ThisBuild / scalaVersion := "3.4.3"
 
 lazy val ScalaTestVersion = "3.2.18"
+lazy val Slf4jVersion = "2.0.16"
 
 lazy val shadedLibs = project.in(file("shaded-libs"))
   .settings(
@@ -26,7 +27,7 @@ lazy val core = project.in(file("core"))
     executableScriptName := "srp",
     libraryDependencies ++= Seq(
       "org.scala-lang" %% "scala3-compiler" % scalaVersion.value,
-      "org.slf4j"       % "slf4j-simple"    % "2.0.16" % Optional,
+      "org.slf4j"       % "slf4j-simple"    % Slf4jVersion % Optional,
     ),
     assemblyJarName := "srp.jar", // TODO remove the '.jar' suffix - when doing so, it doesn't work any longer
   )
@@ -43,6 +44,7 @@ lazy val server = project.in(file("server"))
     fork := true, // important: otherwise we run into classloader issues
     libraryDependencies ++= Seq(
       "com.lihaoyi"   %% "cask"         % "0.8.3",
+      "org.slf4j"      % "slf4j-simple" % Slf4jVersion % Optional,
       "com.lihaoyi"   %% "requests"     % "0.8.2" % Test,
       "org.scalatest" %% "scalatest"    % ScalaTestVersion % "it",
     )

core/src/main/scala/replpp/Config.scala

@@ -153,7 +153,7 @@ object Config {
 
     def runBefore[C](using builder: OParserBuilder[C])(action: Action[String, C]) = {
       builder.opt[String]("runBefore")
-        .valueName("'val foo = 42'")
+        .valueName("\"import Int.MaxValue\"")
         .unbounded()
         .optional()
         .action(action)

server/src/it/scala/replpp/server/EmbeddedReplTests.scala

@@ -13,19 +13,19 @@ import scala.concurrent.duration.Duration
 class EmbeddedReplTests extends AnyWordSpec with Matchers {
 
   "execute commands synchronously" in {
-    val repl = new EmbeddedRepl(defaultCompilerArgs)
+    val repl = new EmbeddedRepl(defaultCompilerArgs, runBeforeCode = Seq("import Short.MaxValue"))
 
-    repl.query("val x = 0").output.trim shouldBe "val x: Int = 0"
-    repl.query("x + 1").output.trim     shouldBe "val res0: Int = 1"
+    repl.query("val x = MaxValue").output.trim shouldBe "val x: Short = 32767"
+    repl.query("x + 1").output.trim     shouldBe "val res0: Int = 32768"
 
     repl.shutdown()
   }
 
   "execute a command asynchronously" in {
-    val repl = new EmbeddedRepl(defaultCompilerArgs)
-    val (uuid, futureResult) = repl.queryAsync("val x = 0")
+    val repl = new EmbeddedRepl(defaultCompilerArgs, runBeforeCode = Seq("import Short.MaxValue"))
+    val (uuid, futureResult) = repl.queryAsync("val x = MaxValue")
     val result = Await.result(futureResult, Duration.Inf)
-    result.trim shouldBe "val x: Int = 0"
+    result.trim shouldBe "val x: Short = 32767"
     repl.shutdown()
   }
 

server/src/main/scala/replpp/server/EmbeddedRepl.scala

@@ -12,17 +12,20 @@ import java.util.concurrent.Executors
 import scala.concurrent.duration.Duration
 import scala.concurrent.{Await, ExecutionContext, ExecutionContextExecutorService, Future}
 
-class EmbeddedRepl(compilerArgs: Array[String]) {
+class EmbeddedRepl(compilerArgs: Array[String], runBeforeCode: Seq[String]) {
   private val logger: Logger = LoggerFactory.getLogger(getClass)
 
   /** repl and compiler output ends up in this replOutputStream */
   private val replOutputStream = new ByteArrayOutputStream()
 
-  private val replDriver: ReplDriver = {
-    new ReplDriver(compilerArgs, new PrintStream(replOutputStream), classLoader = None)
-  }
+  private val replDriver = ReplDriver(compilerArgs, new PrintStream(replOutputStream), classLoader = None)
 
-  private var state: State = replDriver.initialState
+  private var state: State = {
+    if (runBeforeCode.nonEmpty)
+      replDriver.execute(runBeforeCode)
+    else 
+      replDriver.initialState
+  }
 
   private val singleThreadedJobExecutor: ExecutionContextExecutorService =
     ExecutionContext.fromExecutorService(Executors.newSingleThreadExecutor())

server/src/main/scala/replpp/server/ReplServer.scala

@@ -23,7 +23,7 @@ object ReplServer {
 
     val baseConfig = precompilePredefFiles(serverConfig.baseConfig)
     val compilerArgs = replpp.compilerArgs(baseConfig)
-    val embeddedRepl = new EmbeddedRepl(compilerArgs)
+    val embeddedRepl = new EmbeddedRepl(compilerArgs, baseConfig.runBefore)
     Runtime.getRuntime.addShutdownHook(new Thread(() => {
       logger.info("Shutting down server...")
       embeddedRepl.shutdown()

server/src/test/scala/replpp/server/ReplServerTests.scala

@@ -123,6 +123,29 @@ class ReplServerTests extends AnyWordSpec with Matchers {
           getResultResponse("stdout").str shouldBe "val bar: Int = 42\n"
         }
 
+        "use runBefore code" in Fixture(runBeforeCode = Seq("import Int.MaxValue")) { url =>
+          val wsMsgPromise = scala.concurrent.Promise[String]()
+          val connectedPromise = scala.concurrent.Promise[String]()
+          cask.util.WsClient.connect(s"$url/connect") {
+            case cask.Ws.Text(msg) if msg == "connected" =>
+              connectedPromise.success(msg)
+            case cask.Ws.Text(msg) =>
+              wsMsgPromise.success(msg)
+          }
+          Await.result(connectedPromise.future, DefaultPromiseAwaitTimeout)
+          val postQueryResponse = postQueryAsync(url, "val bar = MaxValue")
+          val queryUUID = postQueryResponse("uuid").str
+          queryUUID.length should not be 0
+
+          val queryResultWSMessage = Await.result(wsMsgPromise.future, DefaultPromiseAwaitTimeout)
+          queryResultWSMessage.length should not be 0
+
+          val getResultResponse = getResponse(url, queryUUID)
+          getResultResponse.obj.keySet should contain("success")
+          getResultResponse("uuid").str shouldBe queryResultWSMessage
+          getResultResponse("stdout").str shouldBe "val bar: Int = 2147483647\n"
+        }
+
         "disallow fetching the result of a completed query with an invalid auth header" in Fixture() { url =>
           val wsMsgPromise = scala.concurrent.Promise[String]()
           val connectedPromise = scala.concurrent.Promise[String]()
@@ -384,7 +407,7 @@ object Fixture {
     )
   }
 
-  def apply[T](predefCode: String = "")(urlToResult: String => T): T = {
+  def apply[T](predefCode: String = "", runBeforeCode: Seq[String] = Seq.empty)(urlToResult: String => T): T = {
     val additionalClasspathEntryMaybe: Option[Path] =
       if (predefCode.trim.isEmpty) None
       else {
@@ -395,7 +418,7 @@ object Fixture {
         Files.delete(predefFile)
         predefClassfiles.toOption
       }
-    val embeddedRepl = new EmbeddedRepl(compilerArgs(additionalClasspathEntryMaybe))
+    val embeddedRepl = new EmbeddedRepl(compilerArgs(additionalClasspathEntryMaybe), runBeforeCode)
 
     val host = "localhost"
     val port = 8081