* Made the Mask annotation repeatable

gmu-swe · Nov 29, 2023 · 5469584 · 5469584
1 parent 5c3aae3
commit 5469584
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 58 deletions.
diff --git a/Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/agent/MaskRegistryPatchingCV.java b/Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/agent/MaskRegistryPatchingCV.java
@@ -1,27 +1,23 @@
 package edu.columbia.cs.psl.phosphor.agent;
 
 import edu.columbia.cs.psl.phosphor.Configuration;
-import edu.columbia.cs.psl.phosphor.Phosphor;
 import edu.columbia.cs.psl.phosphor.instrumenter.MethodRecordImpl;
 import edu.columbia.cs.psl.phosphor.instrumenter.TaintMethodRecord;
 import edu.columbia.cs.psl.phosphor.mask.MaskRegistry;
 import edu.columbia.cs.psl.phosphor.runtime.mask.Mask;
 import edu.columbia.cs.psl.phosphor.runtime.mask.SunUnsafeMasker;
 import edu.columbia.cs.psl.phosphor.struct.harmony.util.HashMap;
 import edu.columbia.cs.psl.phosphor.struct.harmony.util.Map;
-import org.objectweb.asm.*;
-import org.objectweb.asm.tree.AnnotationNode;
-import org.objectweb.asm.tree.ClassNode;
-import org.objectweb.asm.tree.MethodNode;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
 
-import java.io.IOException;
-import java.io.InputStream;
-import java.util.List;
+import java.lang.reflect.Method;
 
 public class MaskRegistryPatchingCV extends ClassVisitor {
     private static final String MASK_REGISTRY_INTERNAL_NAME = Type.getInternalName(MaskRegistry.class);
     private static final String TARGET_METHOD_NAME = "initialize";
-    private static final String MONITOR_DESC = Type.getDescriptor(Mask.class);
     private static final Class<?>[] MASKERS = new Class<?>[] {SunUnsafeMasker.class};
 
     public MaskRegistryPatchingCV(ClassVisitor cv) {
@@ -31,59 +27,30 @@ public MaskRegistryPatchingCV(ClassVisitor cv) {
     public static Map<String, MaskRegistry.MaskInfo> readMasks() {
         Map<String, MaskRegistry.MaskInfo> map = new HashMap<>();
         for (Class<?> clazz : MASKERS) {
-            ClassNode cn = getClassNode(clazz);
-            for (MethodNode mn : cn.methods) {
-                readMasks(cn.name, mn, map);
-            }
-        }
-        return map;
-    }
-
-    private static ClassNode getClassNode(Class<?> clazz) {
-        try {
-            ClassNode cn = new ClassNode();
-            String name = clazz.getName().replace('.', '/') + ".class";
-            ClassLoader classLoader = clazz.getClassLoader();
-            try (InputStream in = classLoader == null || Phosphor.RUNTIME_INST
-                    ? ClassLoader.getSystemResourceAsStream(name)
-                    : classLoader.getResourceAsStream(name)) {
-                if (in == null) {
-                    throw new IllegalStateException("Failed to read class: " + clazz);
-                }
-                new ClassReader(in).accept(cn, ClassReader.SKIP_CODE);
-            }
-            return cn;
-        } catch (IOException e) {
-            throw new IllegalStateException("Failed to read class: " + clazz, e);
-        }
-    }
-
-    private static void readMasks(String className, MethodNode mn, Map<String, MaskRegistry.MaskInfo> map) {
-        if ((mn.access & Opcodes.ACC_STATIC) != 0 && mn.invisibleAnnotations != null) {
-            MethodRecordImpl record = new MethodRecordImpl(Opcodes.INVOKESTATIC, className, mn.name, mn.desc, false);
-            for (AnnotationNode an : mn.invisibleAnnotations) {
-                if (MONITOR_DESC.equals(an.desc)) {
-                    map.put(createKey(mn, an), new MaskRegistry.MaskInfo(record));
+            for (Method method : clazz.getDeclaredMethods()) {
+                for (Mask mask : method.getAnnotationsByType(Mask.class)) {
+                    MethodRecordImpl record = new MethodRecordImpl(
+                            Opcodes.INVOKEVIRTUAL,
+                            Type.getInternalName(clazz),
+                            method.getName(),
+                            Type.getMethodDescriptor(method),
+                            false);
+                    String key = createKey(record, mask);
+                    map.put(key, new MaskRegistry.MaskInfo(record));
                 }
             }
         }
+        return map;
     }
 
-    private static String createKey(MethodNode mn, AnnotationNode an) {
-        List<Object> values = an.values;
-        String className = ((Type) values.get(1)).getInternalName();
-        String methodName = mn.name;
-        boolean isStatic = false;
-        if (values.size() == 4) {
-            isStatic = (boolean) values.get(3);
-        }
-        String descriptor = isStatic ? mn.desc : removeFirstParameter(mn.desc);
-        return MaskRegistry.getKey(className, methodName, descriptor);
+    private static String createKey(MethodRecordImpl record, Mask mask) {
+        String className = Type.getInternalName(mask.owner());
+        String descriptor = mask.isStatic() ? record.getDescriptor() : removeFirstParameter(record.getDescriptor());
+        return MaskRegistry.getKey(className, record.getName(), descriptor);
     }
 
     private static String removeFirstParameter(String descriptor) {
-        int index = descriptor.indexOf(';');
-        return '(' + descriptor.substring(index + 1);
+        return '(' + descriptor.substring(descriptor.indexOf(';') + 1);
     }
 
     @Override

diff --git a/Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/runtime/mask/Mask.java b/Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/runtime/mask/Mask.java
@@ -1,9 +1,6 @@
 package edu.columbia.cs.psl.phosphor.runtime.mask;
 
-import java.lang.annotation.ElementType;
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-import java.lang.annotation.Target;
+import java.lang.annotation.*;
 
 /**
  * Indicates that calls to the annotated method should be added during instrumentation to
@@ -19,6 +16,7 @@
  */
 @Target(ElementType.METHOD)
 @Retention(RetentionPolicy.CLASS)
+@Repeatable(Masks.class)
 public @interface Mask {
     Class<?> owner();
 

diff --git a/Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/runtime/mask/Masks.java b/Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/runtime/mask/Masks.java
@@ -0,0 +1,12 @@
+package edu.columbia.cs.psl.phosphor.runtime.mask;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Target(ElementType.METHOD)
+@Retention(RetentionPolicy.CLASS)
+public @interface Masks {
+    Mask[] value();
+}