Visit parents of Lambda classes to catch potential inherited random generators (#23)

Author: mbfreder

src/main/java/software/amazon/lambda/snapstart/ByteCodeIntrospector.java

@@ -8,6 +8,7 @@
 import edu.umd.cs.findbugs.ba.XMethod;
 import edu.umd.cs.findbugs.classfile.CheckedAnalysisException;
 import edu.umd.cs.findbugs.classfile.ClassDescriptor;
+import edu.umd.cs.findbugs.classfile.Global;
 import edu.umd.cs.findbugs.classfile.FieldDescriptor;
 import edu.umd.cs.findbugs.classfile.Global;
 import org.apache.bcel.generic.Type;
@@ -40,6 +41,7 @@ public class ByteCodeIntrospector {
         put("java.lang.System", setOf("currentTimeMillis", "nanoTime"));
     }};
 
+    private LambdaHandlerParentsDatabase lambdaHandlerParentsDatabase;
     private LambdaHandlerFieldsDatabase database;
 
     private static Set<String> setOf(String ... strings) {
@@ -109,6 +111,11 @@ boolean implementsLambdaInterface(XClass xClass) {
         return false;
     }
 
+    boolean isLambdaHandlerParentClass(XClass xClass) {
+        lambdaHandlerParentsDatabase = Global.getAnalysisCache().getDatabase(LambdaHandlerParentsDatabase.class);
+        return lambdaHandlerParentsDatabase.getParentClasses().contains(xClass.toString());
+    }
+
     /**
      * This returns true only when the class directly implements
      * <a href="https://docs.oracle.com/javase/8/docs/api/java/util/function/Function.html">Java Functional Interface</a>.

src/main/java/software/amazon/lambda/snapstart/CacheLambdaHandlerParentClasses.java

@@ -0,0 +1,52 @@
+package software.amazon.lambda.snapstart;
+
+import edu.umd.cs.findbugs.BugReporter;
+import edu.umd.cs.findbugs.Detector;
+import edu.umd.cs.findbugs.ba.ClassContext;
+import edu.umd.cs.findbugs.ba.XClass;
+import org.apache.bcel.classfile.JavaClass;
+import edu.umd.cs.findbugs.classfile.Global;
+
+import java.util.Objects;
+
+
+public class CacheLambdaHandlerParentClasses implements Detector {
+
+    private final ByteCodeIntrospector introspector;
+    private ClassContext classContext;
+    private XClass xClass;
+    private final LambdaHandlerParentsDatabase database;
+
+    public CacheLambdaHandlerParentClasses(BugReporter bugReporter) {
+        this.introspector = new ByteCodeIntrospector();
+        database = new LambdaHandlerParentsDatabase();
+        Global.getAnalysisCache().eagerlyPutDatabase(LambdaHandlerParentsDatabase.class, database);
+    }
+
+    @Override
+    public void visitClassContext(ClassContext classContext) {
+        this.classContext = classContext;
+        this.xClass = classContext.getXClass();
+        if (introspector.isLambdaHandler(xClass)) {
+            JavaClass[] parentClasses = null;
+            try {
+                parentClasses = classContext.getJavaClass().getSuperClasses();
+            } catch (ClassNotFoundException e) {
+                // Do nothing
+            }
+
+            if (Objects.nonNull(parentClasses)) {
+                for (JavaClass parentClass : parentClasses) {
+                    if (!parentClass.getClassName().equals("java.lang.Object")) {
+                        database.addLambdaParentClass(parentClass.getClassName().replace(".", "/"));
+                    }
+                }
+            }
+        }
+    }
+
+    @Override
+    public void report() {
+        // this is a non-reporting detector
+    }
+}

src/main/java/software/amazon/lambda/snapstart/LambdaHandlerInitedWithRandomValue.java

@@ -24,6 +24,7 @@ public class LambdaHandlerInitedWithRandomValue extends OpcodeStackDetector {
 
     private final BugReporter bugReporter;
     private boolean isLambdaHandlerClass;
+    private boolean isLambdaHandlerParentClass;
     private boolean implementsFunctionalInterface;
     private boolean isLambdaHandlerField;
     private boolean isCracResource;
@@ -36,6 +37,7 @@ public class LambdaHandlerInitedWithRandomValue extends OpcodeStackDetector {
     public LambdaHandlerInitedWithRandomValue(BugReporter bugReporter) {
         this.bugReporter = bugReporter;
         this.isLambdaHandlerClass = false;
+        this.isLambdaHandlerParentClass = false;
         this.implementsFunctionalInterface = false;
         this.isLambdaHandlerField = false;
         this.isCracResource = false;
@@ -52,6 +54,7 @@ public void visit(JavaClass obj) {
         inCracBeforeCheckpoint = false;
         XClass xClass = getXClass();
         isLambdaHandlerClass = introspector.isLambdaHandler(xClass);
+        isLambdaHandlerParentClass = introspector.isLambdaHandlerParentClass(xClass);
         implementsFunctionalInterface = introspector.implementsFunctionalInterface(xClass);
         isLambdaHandlerField = introspector.isLambdaHandlerField(xClass);
         isCracResource = introspector.isCracResource(xClass);
@@ -60,7 +63,7 @@ public void visit(JavaClass obj) {
     @Override
     public boolean shouldVisitCode(Code code) {
         boolean shouldVisit = false;
-        if (isLambdaHandlerClass || implementsFunctionalInterface || isLambdaHandlerField) {
+        if (isLambdaHandlerClass || implementsFunctionalInterface || isLambdaHandlerField || isLambdaHandlerParentClass) {
             inStaticInitializer = getMethodName().equals(Const.STATIC_INITIALIZER_NAME);
             inInitializer = getMethodName().equals(Const.CONSTRUCTOR_NAME);
             database = Global.getAnalysisCache().getDatabase(ReturnValueRandomnessPropertyDatabase.class);

src/main/java/software/amazon/lambda/snapstart/LambdaHandlerParentsDatabase.java

@@ -0,0 +1,17 @@
+package software.amazon.lambda.snapstart;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class LambdaHandlerParentsDatabase {
+
+    private final List<String> parentClasses = new ArrayList<>();
+
+    public List<String> getParentClasses() {
+        return this.parentClasses;
+    }
+
+    public void addLambdaParentClass(String parentClass) {
+        this.parentClasses.add(parentClass);
+    }
+}

src/main/resources/findbugs.xml

@@ -20,10 +20,16 @@
                 <Earlier class="software.amazon.lambda.snapstart.CacheLambdaHandlerFields"/>
                 <Later class="software.amazon.lambda.snapstart.LambdaHandlerInitedWithRandomValue"/>
             </SplitPass>
+            <SplitPass>
+                <Earlier class="software.amazon.lambda.snapstart.CacheLambdaHandlerParentClasses"/>
+                <Later class="software.amazon.lambda.snapstart.LambdaHandlerInitedWithRandomValue"/>
+            </SplitPass>
         </OrderingConstraints>
 
         <Detector class="software.amazon.lambda.snapstart.BuildRandomReturningMethodsDatabase"
                   speed="fast" reports="" disabled="false" hidden="true"/>
+        <Detector class="software.amazon.lambda.snapstart.CacheLambdaHandlerParentClasses"
+                  speed="fast" reports="" disabled="false" hidden="true"/>
         <Detector class="software.amazon.lambda.snapstart.CacheLambdaHandlerFields"
                   speed="fast" reports="" disabled="false" hidden="true"/>
         <Detector class="software.amazon.lambda.snapstart.LambdaHandlerInitedWithRandomValue"

src/main/resources/messages.xml

@@ -18,6 +18,12 @@
       Detector that stores all the fields of the Lambda Handler class to be visited later.
     </Details>
   </Detector>
+  
+  <Detector class="software.amazon.lambda.snapstart.CacheLambdaHandlerParentClasses">
+    <Details>
+      Detector that stores parent classes of Lambda Handler classes to be visited later.
+    </Details>
+  </Detector>
 
   <Detector class="software.amazon.lambda.snapstart.LambdaHandlerInitedWithRandomValue">
     <Details>

src/test/java/software/amazon/lambda/snapstart/LambdaHandlerInitedWithRandomValueTest.java

@@ -152,6 +152,13 @@ public void testLambdaWithToString() {
         assertThat(bugCollection, containsExactly(1, snapStartBugMatcher().inClass("LambdaWithToString").atField("random").atLine(11).build()));
     }
 
+    @Test
+    public void testLambdaWithParentClass() {
+        BugCollection bugCollection = findBugsInClasses("LambdaWithParentClass", "ParentHandler", "SuperParentHandler");
+        assertThat(bugCollection, containsExactly(1, snapStartBugMatcher().inClass("ParentHandler").atField("parentId").atLine(6).build()));
+        assertThat(bugCollection, containsExactly(1, snapStartBugMatcher().inClass("SuperParentHandler").atField("superParentId").atLine(6).build()));
+    }
+  
     @Test
     public void testClassImplementingFunctionalInterface() {
         BugCollection bugCollection = findBugsInClasses("ImplementsFunctionalInterface");

src/test/java/software/amazon/lambda/snapstart/lambdaexamples/LambdaWithParentClass.java

@@ -0,0 +1,17 @@
+package software.amazon.lambda.snapstart.lambdaexamples;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+public class LambdaWithParentClass extends ParentHandler implements RequestHandler<String, String> {
+
+    @Override
+    public String handleRequest(String s, Context context) {
+        Logger logger = Logger.getLogger("LambdaWithParentClass");
+        logger.log(Level.INFO, superParentId.toString());
+        return parentId.toString(); 
+    }
+}

src/test/java/software/amazon/lambda/snapstart/lambdaexamples/ParentHandler.java

@@ -0,0 +1,9 @@
+package software.amazon.lambda.snapstart.lambdaexamples;
+
+import java.util.UUID;
+
+public abstract class ParentHandler extends SuperParentHandler {
+    protected final UUID parentId = UUID.randomUUID();
+
+    protected ParentHandler() {}
+}

src/test/java/software/amazon/lambda/snapstart/lambdaexamples/SuperParentHandler.java

@@ -0,0 +1,9 @@
+package software.amazon.lambda.snapstart.lambdaexamples;
+
+import java.util.UUID;
+
+public class SuperParentHandler {
+    protected final UUID superParentId = UUID.randomUUID();
+
+    protected SuperParentHandler() {}
+}