Fix build.

cmu-pasta · Aug 14, 2024 · 9c8b40c · 9c8b40c
1 parent 07c46bd
commit 9c8b40c
Show file tree

Hide file tree

Showing 6 changed files with 47 additions and 14 deletions.
diff --git a/instrumentation/agent/src/main/kotlin/ApplicationCodeTransformer.kt b/instrumentation/agent/src/main/kotlin/ApplicationCodeTransformer.kt
@@ -8,6 +8,9 @@ import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.ClassWriter
 import org.objectweb.asm.Opcodes
 import org.objectweb.asm.tree.ClassNode
+import org.objectweb.asm.util.CheckClassAdapter
+import org.pastalab.fray.instrumentation.base.Configs.DEBUG_MODE
+import org.pastalab.fray.instrumentation.base.Utils.writeClassFile
 import org.pastalab.fray.instrumentation.base.visitors.*
 
 class ApplicationCodeTransformer : ClassFileTransformer {
@@ -38,8 +41,9 @@ class ApplicationCodeTransformer : ClassFileTransformer {
       // This is likely a JDK class, so skip transformation
       return classfileBuffer
     }
-    File("/tmp/out/origin/${className.replace("/", ".").removePrefix(".")}.class")
-        .writeBytes(classfileBuffer)
+    if (DEBUG_MODE) {
+      writeClassFile(className, classfileBuffer, false)
+    }
     try {
       val classReader = ClassReader(classfileBuffer)
       val cn = ClassNode()
@@ -63,15 +67,17 @@ class ApplicationCodeTransformer : ClassFileTransformer {
       val classWriter = ClassWriter(classReader, ClassWriter.COMPUTE_MAXS)
       if (classVersionInstrumenter.classVersion >= Opcodes.V1_5) {
         cn.accept(classWriter)
-        //        val checkClassAdapter = CheckClassAdapter(classWriter)
-        //        cn.accept(checkClassAdapter)
+        val checkClassAdapter = CheckClassAdapter(classWriter)
+        cn.accept(checkClassAdapter)
       } else {
         cn.accept(classWriter)
       }
       val out = classWriter.toByteArray()
-      File("/tmp/out/app/${className.replace("/",
-          ".").removePrefix(".")}.class")
-          .writeBytes(out)
+      if (DEBUG_MODE) {
+        File("/tmp/out/app/${className.replace("/",
+            ".").removePrefix(".")}.class")
+            .writeBytes(out)
+      }
       return out
     } catch (e: Throwable) {
       println("Failed to instrument: $className")

diff --git a/instrumentation/agent/src/main/kotlin/PreMain.kt b/instrumentation/agent/src/main/kotlin/PreMain.kt
@@ -1,7 +1,9 @@
 package org.pastalab.fray.instrumentation.agent
 
 import java.lang.instrument.Instrumentation
+import org.pastalab.fray.instrumentation.base.Utils
 
 fun premain(arguments: String?, instrumentation: Instrumentation) {
+  Utils.prepareDebugFolder("origin")
   instrumentation.addTransformer(ApplicationCodeTransformer())
 }
diff --git a/instrumentation/base/src/main/kotlin/org/pastalab/fray/instrumentation/base/Configs.kt b/instrumentation/base/src/main/kotlin/org/pastalab/fray/instrumentation/base/Configs.kt
@@ -0,0 +1,5 @@
+package org.pastalab.fray.instrumentation.base
+
+object Configs {
+  val DEBUG_MODE = System.getProperty("fray.debug", "false").toBoolean()
+}
diff --git a/instrumentation/base/src/main/kotlin/org/pastalab/fray/instrumentation/base/Utils.kt b/instrumentation/base/src/main/kotlin/org/pastalab/fray/instrumentation/base/Utils.kt
@@ -1,9 +1,11 @@
 package org.pastalab.fray.instrumentation.base
 
+import java.io.File
 import java.nio.file.Paths
 import kotlin.io.path.ExperimentalPathApi
 import kotlin.io.path.createDirectories
 import kotlin.io.path.deleteRecursively
+import org.pastalab.fray.instrumentation.base.Configs.DEBUG_MODE
 
 object Utils {
   @OptIn(ExperimentalPathApi::class)
@@ -12,4 +14,21 @@ object Utils {
     path.deleteRecursively()
     path.createDirectories()
   }
+
+  fun writeClassFile(className: String, classBytes: ByteArray, instrumented: Boolean) {
+    if (instrumented) {
+      File("/tmp/out/instrumented/${className.replace("/", ".").removePrefix(".")}.class")
+          .writeBytes(classBytes)
+    } else {
+      File("/tmp/out/origin/${className.replace("/", ".").removePrefix(".")}.class")
+          .writeBytes(classBytes)
+    }
+  }
+
+  init {
+    if (DEBUG_MODE) {
+      prepareDebugFolder("instrumented")
+      prepareDebugFolder("origin")
+    }
+  }
 }
diff --git a/instrumentation/jdk/src/main/kotlin/org/pastalab/fray/jdk/JDKInstrumenter.kt b/instrumentation/jdk/src/main/kotlin/org/pastalab/fray/jdk/JDKInstrumenter.kt
@@ -11,11 +11,15 @@ import org.objectweb.asm.commons.ModuleTargetAttribute
 import org.objectweb.asm.tree.ClassNode
 import org.objectweb.asm.tree.ModuleExportNode
 import org.objectweb.asm.util.CheckClassAdapter
+import org.pastalab.fray.instrumentation.base.Configs.DEBUG_MODE
+import org.pastalab.fray.instrumentation.base.Utils.writeClassFile
 import org.pastalab.fray.instrumentation.base.visitors.*
 
 fun instrumentClass(path: String, inputStream: InputStream): ByteArray {
   val byteArray = inputStream.readBytes()
-  File("/tmp/out/origin/${path.replace("/", ".").removePrefix(".")}").writeBytes(byteArray)
+  if (DEBUG_MODE) {
+    writeClassFile(path, byteArray, false)
+  }
   val shouldSkipChecking = path.contains("SystemModules") || path.contains("$")
 
   try {
@@ -53,8 +57,9 @@ fun instrumentClass(path: String, inputStream: InputStream): ByteArray {
       cn.accept(CheckClassAdapter(classWriter))
     }
     val out = classWriter.toByteArray()
-    File("/tmp/out/jdk/${path.replace("/", ".").removePrefix(".")}").writeBytes(out)
-    //        File("/tmp/${path.replace("/", ".").removePrefix(".")}").writeBytes(out)
+    if (DEBUG_MODE) {
+      writeClassFile(path, out, true)
+    }
     return out
   } catch (e: Throwable) {
     println(path)

diff --git a/instrumentation/jdk/src/main/kotlin/org/pastalab/fray/jdk/jlink/JlinkPlugin.kt b/instrumentation/jdk/src/main/kotlin/org/pastalab/fray/jdk/jlink/JlinkPlugin.kt
@@ -8,7 +8,6 @@ import jdk.tools.jlink.plugin.Plugin
 import jdk.tools.jlink.plugin.ResourcePool
 import jdk.tools.jlink.plugin.ResourcePoolBuilder
 import jdk.tools.jlink.plugin.ResourcePoolEntry
-import org.pastalab.fray.instrumentation.base.Utils
 import org.pastalab.fray.jdk.instrumentClass
 import org.pastalab.fray.jdk.instrumentModuleInfo
 
@@ -27,9 +26,6 @@ class JlinkPlugin : Plugin {
 
   override fun transform(input: ResourcePool, output: ResourcePoolBuilder): ResourcePool {
     println("Start fray plugin!")
-
-    Utils.prepareDebugFolder("jdk")
-    Utils.prepareDebugFolder("origin")
     input.transformAndCopy(
         { entry ->
           var resourcePoolEntry = entry