Merge pull request #205 from gmu-swe/annotation-fix

Fixed failure when instrumenting classes with an annotation containing an Enum-type value

Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/PhosphorPatcher.java

@@ -1,6 +1,7 @@
 package edu.columbia.cs.psl.phosphor;
 
 import edu.columbia.cs.psl.phosphor.instrumenter.ConfigurationEmbeddingMV;
+import edu.columbia.cs.psl.phosphor.instrumenter.TaintMethodRecord;
 import edu.columbia.cs.psl.phosphor.runtime.jdk.unsupported.UnsafeProxy;
 import org.objectweb.asm.*;
 import org.objectweb.asm.commons.ModuleHashesAttribute;
@@ -28,6 +29,8 @@ public byte[] patch(String name, byte[] content) throws IOException {
         } else if (name.equals("edu/columbia/cs/psl/phosphor/runtime/RuntimeJDKInternalUnsafePropagator.class")) {
             return transformUnsafePropagator(new ByteArrayInputStream(content),
                     "jdk/internal/misc/Unsafe", patchUnsafeNames);
+        } else if (AsmPatchingCV.isApplicable(name)) {
+            return AsmPatchingCV.patch(content);
         } else {
             return content;
         }
@@ -168,4 +171,43 @@ public void visitLocalVariable(String name, String descriptor, String signature,
             };
         }
     }
+
+    private static class AsmPatchingCV extends ClassVisitor {
+        private static final String ASM_PREFIX = "edu/columbia/cs/psl/phosphor/org/objectweb/asm/";
+        private static final Type OBJECT_TYPE = Type.getType(Object.class);
+
+        public AsmPatchingCV(ClassWriter cw) {
+            super(Configuration.ASM_VERSION, cw);
+        }
+
+        @Override
+        public MethodVisitor visitMethod(
+                int access, String name, String descriptor, String signature, String[] exceptions) {
+            MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
+            return new MethodVisitor(api, mv) {
+
+                @Override
+                public void visitMethodInsn(
+                        int opcode, String owner, String name, String descriptor, boolean isInterface) {
+                    super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
+                    // Ensure that the return value is unwrapped if necessary
+                    if (owner.startsWith("java/") && OBJECT_TYPE.equals(Type.getReturnType(descriptor))) {
+                        TaintMethodRecord.TAINTED_REFERENCE_ARRAY_UNWRAP.delegateVisit(mv);
+                    }
+                }
+            };
+        }
+
+        public static byte[] patch(byte[] classFileBuffer) {
+            ClassReader cr = new ClassReader(classFileBuffer);
+            ClassWriter cw = new ClassWriter(cr, 0);
+            ClassVisitor cv = new AsmPatchingCV(cw);
+            cr.accept(cv, 0);
+            return cw.toByteArray();
+        }
+
+        public static boolean isApplicable(String className) {
+            return className.startsWith(ASM_PREFIX);
+        }
+    }
 }

Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/instrumenter/TaintMethodRecord.java

@@ -124,7 +124,9 @@ public enum TaintMethodRecord implements MethodRecord {
     TAINTED_REFERENCE_ARRAY_GET(INVOKEVIRTUAL, TaggedReferenceArray.class, "get", Object.class, false, int.class, Taint.class, PhosphorStackFrame.class),
     TAINTED_SHORT_ARRAY_GET(INVOKEVIRTUAL, TaggedShortArray.class, "get", short.class, false, int.class, Taint.class, PhosphorStackFrame.class),
     // Methods from TaintSourceWrapper
-    AUTO_TAINT(INVOKEVIRTUAL, TaintSourceWrapper.class, "autoTaint", Object.class, false, Object.class, String.class, String.class, int.class);
+    AUTO_TAINT(INVOKEVIRTUAL, TaintSourceWrapper.class, "autoTaint", Object.class, false, Object.class, String.class, String.class, int.class),
+    // TaggedReferenceArray
+    TAINTED_REFERENCE_ARRAY_UNWRAP(INVOKESTATIC, TaggedReferenceArray.class, "unwrap", Object.class, false, Object.class);
 
     private final int opcode;
     private final String owner;

Phosphor/src/main/java/edu/columbia/cs/psl/phosphor/struct/TaggedReferenceArray.java

@@ -12,8 +12,7 @@
 import java.io.ObjectOutputStream;
 import java.lang.reflect.Array;
 
-import static edu.columbia.cs.psl.phosphor.instrumenter.TaintMethodRecord.TAINTED_REFERENCE_ARRAY_GET;
-import static edu.columbia.cs.psl.phosphor.instrumenter.TaintMethodRecord.TAINTED_REFERENCE_ARRAY_SET;
+import static edu.columbia.cs.psl.phosphor.instrumenter.TaintMethodRecord.*;
 
 public final class TaggedReferenceArray extends TaggedArray {
 
@@ -233,11 +232,15 @@ public static TaggedReferenceArray factory(Object[] array, Taint lengthTaint) {
         return new TaggedReferenceArray(lengthTaint, array);
     }
 
-
     public static Object[] unwrap(TaggedReferenceArray obj) {
         if (obj != null) {
             return obj.val;
         }
         return null;
     }
+
+    @InvokedViaInstrumentation(record = TAINTED_REFERENCE_ARRAY_UNWRAP)
+    public static Object unwrap(Object obj) {
+        return (obj instanceof TaggedReferenceArray) ? ((TaggedReferenceArray) obj).val : obj;
+    }
 }

integration-tests/src/test/java/edu/columbia/cs/psl/test/phosphor/AnnotationInstCase.java

@@ -0,0 +1,37 @@
+package edu.columbia.cs.psl.test.phosphor;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import java.lang.reflect.Method;
+
+public class AnnotationInstCase {
+    @Test
+    public void testEnum() throws ReflectiveOperationException {
+        Method method = Example.class.getDeclaredMethod("x");
+        ExampleAnnotation annotation = method.getAnnotation(ExampleAnnotation.class);
+        Assert.assertNotNull(annotation);
+        Assert.assertEquals(ExampleAnnotation.Color.BLUE, annotation.color());
+    }
+
+    public static class Example {
+        @ExampleAnnotation(color = ExampleAnnotation.Color.BLUE)
+        void x() {}
+    }
+
+    @Target(ElementType.METHOD)
+    @Retention(RetentionPolicy.RUNTIME)
+    public @interface ExampleAnnotation {
+        Color color();
+
+        enum Color {
+            RED,
+            BLUE,
+            GREEN;
+        }
+    }
+}