feat: Visit all fields in Lambda Handler class to support Dependency Injection (#26)

* Resolved merge conflicts by incorporating Functional interface changes

* Visit all fields in Lambda Handler class to support Dependency Injection

Author: mbfreder

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

@@ -8,6 +8,10 @@
 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.FieldDescriptor;
+import edu.umd.cs.findbugs.classfile.Global;
+import org.apache.bcel.generic.Type;
+
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -36,6 +40,8 @@ public class ByteCodeIntrospector {
         put("java.lang.System", setOf("currentTimeMillis", "nanoTime"));
     }};
 
+    private LambdaHandlerFieldsDatabase database;
+
     private static Set<String> setOf(String ... strings) {
         Set<String> set = new HashSet<>();
         Collections.addAll(set, strings);
@@ -95,6 +101,23 @@ boolean implementsFunctionalInterface(XClass xClass) {
                 // ignore
             }
         }
+
+        return false;
+    }
+
+    /** 
+     * This returns true only if this class is used as a field in a Lambda handler class
+     */
+    boolean isLambdaHandlerField(XClass xClass) {
+        database = Global.getAnalysisCache().getDatabase(LambdaHandlerFieldsDatabase.class);
+        for (FieldDescriptor fieldDescriptor : database.getKeys()) {
+            
+
+            String fieldType = Type.getReturnType(fieldDescriptor.getSignature()).toString().replace(".", "/");
+            if (fieldType.equals(xClass.toString())) {
+                return true;
+            }
+        }
         return false;
     }
 

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

@@ -0,0 +1,47 @@
+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 edu.umd.cs.findbugs.classfile.DescriptorFactory;
+import edu.umd.cs.findbugs.classfile.FieldDescriptor;
+import edu.umd.cs.findbugs.classfile.Global;
+import org.apache.bcel.classfile.Field;
+
+/**
+ * This detector stores fields with the Lambda Handler and Crac resources to be used later
+ * for visiting the classes passed in through dependency injection
+ */
+public class CacheLambdaHandlerFields implements Detector {
+
+    private final ByteCodeIntrospector introspector;
+    private ClassContext classContext;
+    private XClass xClass;
+    private final LambdaHandlerFieldsDatabase database;
+
+
+    public CacheLambdaHandlerFields(BugReporter reporter) {
+        introspector = new ByteCodeIntrospector();
+        database = new LambdaHandlerFieldsDatabase();
+        Global.getAnalysisCache().eagerlyPutDatabase(LambdaHandlerFieldsDatabase.class, database);
+    }
+
+    @Override
+    public void visitClassContext(ClassContext classContext) {
+        this.classContext = classContext;
+        this.xClass = classContext.getXClass();
+        if (introspector.isLambdaHandler(xClass)) {
+            Field[] fields = classContext.getJavaClass().getFields();
+            for (Field field : fields) {
+                FieldDescriptor fieldDescriptor = DescriptorFactory.instance().getFieldDescriptor(xClass.toString().replace(".", "/"), field);
+                database.setProperty(fieldDescriptor, true);
+            }
+        }
+    }
+
+    @Override
+    public void report() {
+        // this is a non-reporting detector
+    }
+}

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

@@ -0,0 +1,22 @@
+package software.amazon.lambda.snapstart;
+
+import edu.umd.cs.findbugs.ba.interproc.FieldPropertyDatabase;
+import edu.umd.cs.findbugs.ba.interproc.PropertyDatabaseFormatException;
+
+public class LambdaHandlerFieldsDatabase extends FieldPropertyDatabase<Boolean> {
+
+    public LambdaHandlerFieldsDatabase() {
+        super();
+    }
+
+    @Override
+    protected Boolean decodeProperty(String s) throws PropertyDatabaseFormatException {
+        return null;
+    }
+
+    @Override
+    protected String encodeProperty(Boolean aBoolean) {
+        return null;
+    }
+
+}

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

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

src/main/resources/findbugs.xml

@@ -16,10 +16,16 @@
                 <Earlier class="software.amazon.lambda.snapstart.BuildRandomReturningMethodsDatabase"/>
                 <Later class="software.amazon.lambda.snapstart.LambdaHandlerInitedWithRandomValue"/>
             </SplitPass>
+            <SplitPass>
+                <Earlier class="software.amazon.lambda.snapstart.CacheLambdaHandlerFields"/>
+                <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.CacheLambdaHandlerFields"
+                  speed="fast" reports="" disabled="false" hidden="true"/>
         <Detector class="software.amazon.lambda.snapstart.LambdaHandlerInitedWithRandomValue"
                   reports="AWS_LAMBDA_SNAP_START_BUG" />
 

src/main/resources/messages.xml

@@ -13,6 +13,12 @@
     </Details>
   </Detector>
 
+  <Detector class="software.amazon.lambda.snapstart.CacheLambdaHandlerFields">
+    <Details>
+      Detector that stores all the fields of the Lambda Handler class to be visited later.
+    </Details>
+  </Detector>
+
   <Detector class="software.amazon.lambda.snapstart.LambdaHandlerInitedWithRandomValue">
     <Details>
       Main detector to find out SnapStart bugs in Lambda handler classes.

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

@@ -158,6 +158,12 @@ public void testClassImplementingFunctionalInterface() {
         assertThat(bugCollection, containsExactly(1, snapStartBugMatcher().inClass("ImplementsFunctionalInterface").atField("random").atLine(8).build()));
     }
 
+    @Test
+    public void testLambdaWithDependencyInjection() {
+        BugCollection bugCollection = findBugsInClasses("DependencyInjection", "MyDependency");
+        assertThat(bugCollection, containsExactly(1, snapStartBugMatcher().inClass("MyDependency").atField("random").atLine(6).build()));
+    }
+
     // TODO fix all tests using this and remove this method eventually!
     private static void toBeFixed(Executable e) {
         assertThrows(AssertionError.class, e);

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

@@ -0,0 +1,16 @@
+package software.amazon.lambda.snapstart.lambdaexamples;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+
+public class DependencyInjection implements RequestHandler<String, String> {
+    private final MyDependency myDependency;
+
+    public DependencyInjection(MyDependency myDependency) {
+        this.myDependency = myDependency;
+    }
+    @Override
+    public String handleRequest(String s, Context context) {
+        return myDependency.getUUID().toString();
+    }
+}

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

@@ -0,0 +1,11 @@
+package software.amazon.lambda.snapstart.lambdaexamples;
+
+import java.util.UUID;
+
+public class MyDependency {
+    private final UUID random = UUID.randomUUID();
+
+    public UUID getUUID() {
+        return random;
+    }
+}