* Refactored

* Fixed missing mask from SunUnsafeMasker
gmu-swe · Nov 29, 2023 · efd20e6 · efd20e6
1 parent 3919742
commit efd20e6
Show file tree

Hide file tree

Showing 6 changed files with 210 additions and 238 deletions.
diff --git a/Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/Phosphor.java b/Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/Phosphor.java
@@ -42,13 +42,15 @@ public static InstrumentationAdaptor getInstrumentation() {
     }
 
     @InvokedViaInstrumentation(record = TaintMethodRecord.INSTRUMENT_CLASS_BYTES)
-    public static byte[] instrumentClassBytes(byte[] in) {
-        return new PCLoggingTransformer().transform(null, null, null, null, in, false);
+    public static byte[] instrumentClassBytes(byte[] classFileBuffer) {
+        byte[] result = new PCLoggingTransformer().transform(null, null, null, null, classFileBuffer, false);
+        return result == null ? classFileBuffer : result;
     }
 
     @InvokedViaInstrumentation(record = TaintMethodRecord.INSTRUMENT_CLASS_BYTES_ANONYMOUS)
-    public static byte[] instrumentClassBytesAnonymous(byte[] in) {
-        return new PCLoggingTransformer().transform(null, null, null, null, in, true);
+    public static byte[] instrumentClassBytesAnonymous(byte[] classFileBuffer) {
+        byte[] result = new PCLoggingTransformer().transform(null, null, null, null, classFileBuffer, true);
+        return result == null ? classFileBuffer : result;
     }
 
     private static String[] parseOptions(String agentArgs) {

diff --git a/Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/runtime/mask/JdkUnsafeAdapter.java b/Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/runtime/mask/JdkUnsafeAdapter.java
@@ -18,6 +18,11 @@ public long getInvalidFieldOffset() {
         return Unsafe.INVALID_FIELD_OFFSET;
     }
 
+    @Override
+    public Class<?> defineAnonymousClass(Class<?> hostClass, byte[] data, Object[] cpPatches) {
+        return unsafe.defineAnonymousClass(hostClass, data, cpPatches);
+    }
+
     @Override
     public Class<?> defineClass(
             String name, byte[] b, int off, int len, ClassLoader loader, ProtectionDomain protectionDomain) {

diff --git a/Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/runtime/mask/SunUnsafeAdapter.java b/Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/runtime/mask/SunUnsafeAdapter.java
@@ -25,6 +25,11 @@ private static Class<?> defineClass(
         return null;
     }
 
+    @Override
+    public Class<?> defineAnonymousClass(Class<?> hostClass, byte[] data, Object[] cpPatches) {
+        return unsafe.defineAnonymousClass(hostClass, data, cpPatches);
+    }
+
     @Override
     public Class<?> defineClass(
             String name, byte[] b, int off, int len, ClassLoader loader, ProtectionDomain protectionDomain) {