Enable call selection for taint sources and sinks

Author: DIvanov503

base/src/main/java/proguard/analysis/cpa/defaults/SetAbstractState.java

@@ -19,6 +19,8 @@
 package proguard.analysis.cpa.defaults;
 
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 
 /**
@@ -40,7 +42,17 @@ public class SetAbstractState<T>
      */
     public SetAbstractState(T... items)
     {
-        addAll(Arrays.asList(items));
+        this(Arrays.asList(items));
+    }
+
+    /**
+     * Create a set abstract state from a collection.
+     *
+     * @param c a collection of elements
+     */
+    public SetAbstractState(Collection<? extends T> c)
+    {
+        super(c);
     }
 
     // implementations for LatticeAbstractState

base/src/main/java/proguard/analysis/cpa/domain/taint/TaintAbstractState.java

@@ -18,6 +18,8 @@
 
 package proguard.analysis.cpa.domain.taint;
 
+import proguard.classfile.Signature;
+
 /**
  * A {@link TaintSink} specifies a sink for the taint analysis. A sink can be sensitive to
  * the instance, the arguments, or the static fields. If a sink S is sensitive to X, then
@@ -27,16 +29,16 @@
  */
 public abstract class TaintSink
 {
-    public final String fqn;
+    public final Signature signature;
 
     /**
      * Create a taint sink.
      *
-     * @param fqn           the fully qualified name of a sink method
+     * @param signature the signature of a sink method
      */
-    public TaintSink(String fqn)
+    public TaintSink(Signature signature)
     {
-        this.fqn = fqn;
+        this.signature = signature;
     }
 
     @Override

base/src/main/java/proguard/analysis/cpa/domain/taint/TaintSink.java

@@ -21,6 +21,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Collectors;
+import proguard.classfile.Signature;
 
 /**
  * A {@link TaintSource} specifies a method which can taint any (subset) of the following: the instance, the return value, the argument objects, or static fields.
@@ -30,7 +31,7 @@
 public class TaintSource
 {
 
-    public final String       fqn;
+    public final Signature    signature;
     public final boolean      taintsThis;
     public final boolean      taintsReturn;
     public final Set<Integer> taintsArgs;
@@ -39,15 +40,15 @@ public class TaintSource
     /**
      * Create a taint source.
      *
-     * @param fqn           the fully qualified name of a source method
+     * @param signature     the signature a source method
      * @param taintsThis    whether the source taints the calling instance
      * @param taintsReturn  whether the source taints its return
      * @param taintsArgs    a set of tainted arguments
      * @param taintsGlobals a set of tainted global variables
      */
-    public TaintSource(String fqn, boolean taintsThis, boolean taintsReturn, Set<Integer> taintsArgs, Set<String> taintsGlobals)
+    public TaintSource(Signature signature, boolean taintsThis, boolean taintsReturn, Set<Integer> taintsArgs, Set<String> taintsGlobals)
     {
-        this.fqn = fqn;
+        this.signature = signature;
         this.taintsThis = taintsThis;
         this.taintsReturn = taintsReturn;
         this.taintsArgs = taintsArgs;
@@ -68,7 +69,7 @@ public boolean equals(Object obj)
             return false;
         }
         TaintSource other = (TaintSource) obj;
-        return Objects.equals(fqn, other.fqn)
+        return Objects.equals(signature, other.signature)
                && taintsThis == other.taintsThis
                && taintsReturn == other.taintsReturn
                && Objects.equals(taintsArgs, other.taintsArgs)
@@ -78,13 +79,13 @@ public boolean equals(Object obj)
     @Override
     public int hashCode()
     {
-        return Objects.hash(fqn, taintsThis, taintsReturn, taintsArgs, taintsGlobals);
+        return Objects.hash(signature, taintsThis, taintsReturn, taintsArgs, taintsGlobals);
     }
 
     @Override
     public String toString()
     {
-        StringBuilder builder = new StringBuilder("[TaintSource] ").append(fqn);
+        StringBuilder builder = new StringBuilder("[TaintSource] ").append(signature);
         if (taintsThis)
         {
             builder.append(", taints this");

base/src/main/java/proguard/analysis/cpa/domain/taint/TaintSource.java

@@ -22,7 +22,7 @@
 import java.util.Map.Entry;
 import java.util.stream.Collectors;
 import proguard.analysis.cpa.defaults.MapAbstractState;
-import proguard.analysis.cpa.domain.taint.TaintAbstractState;
+import proguard.analysis.cpa.defaults.SetAbstractState;
 import proguard.analysis.cpa.jvm.domain.reference.JvmReferenceAbstractState;
 import proguard.analysis.cpa.jvm.domain.reference.Reference;
 import proguard.analysis.cpa.jvm.state.heap.JvmHeapAbstractState;
@@ -36,7 +36,7 @@
  * @author Dmitry Ivanov
  */
 public class JvmBasicTaintTreeHeapFollowerAbstractState
-    extends JvmTreeHeapFollowerAbstractState<TaintAbstractState>
+    extends JvmTreeHeapFollowerAbstractState<SetAbstractState<JvmTaintSource>>
     implements JvmTaintHeapAbstractState
 {
 
@@ -50,27 +50,27 @@ public class JvmBasicTaintTreeHeapFollowerAbstractState
      * @param objectMapAbstractStateFactory a map abstract state factory used for constructing the mapping from fields to values
      */
     public JvmBasicTaintTreeHeapFollowerAbstractState(JvmReferenceAbstractState principal,
-                                                      TaintAbstractState defaultValue, MapAbstractState<Reference, HeapNode<TaintAbstractState>> referenceToNode,
-                                                      MapAbstractStateFactory<Reference, HeapNode<TaintAbstractState>> heapMapAbstractStateFactory,
-                                                      MapAbstractStateFactory<String, TaintAbstractState> objectMapAbstractStateFactory)
+                                                      SetAbstractState<JvmTaintSource> defaultValue, MapAbstractState<Reference, HeapNode<SetAbstractState<JvmTaintSource>>> referenceToNode,
+                                                      MapAbstractStateFactory<Reference, HeapNode<SetAbstractState<JvmTaintSource>>> heapMapAbstractStateFactory,
+                                                      MapAbstractStateFactory<String, SetAbstractState<JvmTaintSource>> objectMapAbstractStateFactory)
     {
         super(principal, defaultValue, referenceToNode, heapMapAbstractStateFactory, objectMapAbstractStateFactory);
     }
 
     // implementations for JvmTaintHeapAbstractState
 
     @Override
-    public <T> void taintObject(T object, TaintAbstractState value)
+    public <T> void taintObject(T object, SetAbstractState<JvmTaintSource> value)
     {
     }
 
     // implementations for LatticeAbstractState
 
     @Override
-    public JvmBasicTaintTreeHeapFollowerAbstractState join(JvmHeapAbstractState<TaintAbstractState> abstractState)
+    public JvmBasicTaintTreeHeapFollowerAbstractState join(JvmHeapAbstractState<SetAbstractState<JvmTaintSource>> abstractState)
     {
         JvmBasicTaintTreeHeapFollowerAbstractState other = (JvmBasicTaintTreeHeapFollowerAbstractState) abstractState;
-        MapAbstractState<Reference, HeapNode<TaintAbstractState>> newReferenceToNode = referenceToNode.join(other.referenceToNode);
+        MapAbstractState<Reference, HeapNode<SetAbstractState<JvmTaintSource>>> newReferenceToNode = referenceToNode.join(other.referenceToNode);
         if (referenceToNode == newReferenceToNode)
         {
             return this;

base/src/main/java/proguard/analysis/cpa/jvm/domain/taint/JvmBasicTaintTreeHeapFollowerAbstractState.java

@@ -20,14 +20,18 @@
 
 import java.util.HashSet;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import proguard.analysis.cpa.interfaces.CallEdge;
 import proguard.analysis.cpa.jvm.cfa.edges.JvmCallCfaEdge;
 import proguard.analysis.cpa.jvm.cfa.edges.JvmCfaEdge;
 import proguard.analysis.cpa.jvm.witness.JvmMemoryLocation;
 import proguard.analysis.cpa.jvm.witness.JvmStackLocation;
 import proguard.analysis.cpa.jvm.witness.JvmStaticFieldLocation;
+import proguard.analysis.datastructure.callgraph.Call;
+import proguard.classfile.Signature;
 import proguard.classfile.util.ClassUtil;
 
 /**
@@ -41,27 +45,58 @@ public class JvmInvokeTaintSink
     extends JvmTaintSink
 {
 
+    public final Optional<Predicate<Call>> callMatcher;
+
     public final boolean      takesInstance;
     public final Set<Integer> takesArgs;
     public final Set<String>  takesGlobals;
 
     /**
      * Create a taint sink.
      *
-     * @param fqn           the fully qualified name of a sink method
+     * @param signature     the signature of a sink method
+     * @param takesInstance whether the sink is sensitive to the calling instance
+     * @param takesArgs     a set of sensitive arguments
+     * @param takesGlobals  a set of sensitive global variables
+     */
+    public JvmInvokeTaintSink(Signature signature, boolean takesInstance, Set<Integer> takesArgs, Set<String> takesGlobals)
+    {
+        this(signature, Optional.empty(), takesInstance, takesArgs, takesGlobals);
+    }
+
+    /**
+     * Create a taint sink.
+     *
+     * @param signature     the signature of a sink method
+     * @param callMatcher   whether the call matches this taint sink
+     * @param takesInstance whether the sink is sensitive to the calling instance
+     * @param takesArgs     a set of sensitive arguments
+     * @param takesGlobals  a set of sensitive global variables
+     */
+    public JvmInvokeTaintSink(Signature signature, Predicate<Call> callMatcher, boolean takesInstance, Set<Integer> takesArgs, Set<String> takesGlobals)
+    {
+        this(signature, Optional.of(callMatcher), takesInstance, takesArgs, takesGlobals);
+    }
+
+    /**
+     * Create a taint sink.
+     *
+     * @param signature     the signature of a sink method
+     * @param callMatcher   an optional predicate on whether the call matches this taint sink
      * @param takesInstance whether the sink is sensitive to the calling instance
      * @param takesArgs     a set of sensitive arguments
      * @param takesGlobals  a set of sensitive global variables
      */
-    public JvmInvokeTaintSink(String fqn, boolean takesInstance, Set<Integer> takesArgs, Set<String> takesGlobals)
+    private JvmInvokeTaintSink(Signature signature, Optional<Predicate<Call>> callMatcher, boolean takesInstance, Set<Integer> takesArgs, Set<String> takesGlobals)
     {
-        super(fqn);
+        super(signature);
 
         if (!takesInstance && takesArgs.isEmpty() && takesGlobals.isEmpty())
         {
-            throw new RuntimeException(String.format("Tainted sink for method %s must have taint somewhere!", fqn));
+            throw new RuntimeException(String.format("Tainted sink for method %s must have taint somewhere!", signature));
         }
 
+        this.callMatcher = callMatcher;
         this.takesInstance = takesInstance;
         this.takesArgs = takesArgs;
         this.takesGlobals = takesGlobals;
@@ -74,14 +109,15 @@ public JvmInvokeTaintSink(String fqn, boolean takesInstance, Set<Integer> takesA
     public Set<JvmMemoryLocation> getMemoryLocations()
     {
         Set<JvmMemoryLocation> result = new HashSet<>();
+        String fqn = signature.getFqn();
         String descriptor = fqn.substring(fqn.indexOf("("));
         int parameterSize = ClassUtil.internalMethodParameterSize(descriptor);
         if (takesInstance)
         {
             result.add(new JvmStackLocation(parameterSize));
         }
         takesArgs.forEach(i -> result.add(new JvmStackLocation(parameterSize - ClassUtil.internalMethodVariableIndex(descriptor, true, i))));
-        takesGlobals.forEach(fqn -> result.add(new JvmStaticFieldLocation(fqn)));
+        takesGlobals.forEach(n -> result.add(new JvmStaticFieldLocation(n)));
         return result;
     }
 
@@ -91,7 +127,12 @@ public Set<JvmMemoryLocation> getMemoryLocations()
     @Override
     public boolean matchCfaEdge(JvmCfaEdge edge)
     {
-        return edge instanceof JvmCallCfaEdge && fqn.equals(((CallEdge) edge).getCall().getTarget().getFqn());
+        if (!(edge instanceof JvmCallCfaEdge))
+        {
+            return false;
+        }
+        CallEdge callEdge = (CallEdge) edge;
+        return signature.equals(callEdge.getCall().getTarget()) && callMatcher.map(p -> p.test(callEdge.getCall())).orElse(true);
     }
 
     @Override
@@ -107,21 +148,22 @@ public boolean equals(Object o)
         }
         JvmInvokeTaintSink taintSink = (JvmInvokeTaintSink) o;
         return takesInstance == taintSink.takesInstance
-               && Objects.equals(fqn, taintSink.fqn)
+               && Objects.equals(signature, taintSink.signature)
                && Objects.equals(takesArgs, taintSink.takesArgs)
-               && Objects.equals(takesGlobals, taintSink.takesGlobals);
+               && Objects.equals(takesGlobals, taintSink.takesGlobals)
+               && Objects.equals(callMatcher, taintSink.callMatcher);
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(fqn, takesInstance, takesArgs, takesGlobals);
+        return Objects.hash(signature, takesInstance, takesArgs, takesGlobals, callMatcher);
     }
 
     @Override
     public String toString()
     {
-        StringBuilder result = new StringBuilder("[JvmInvokeTaintSink] ").append(fqn);
+        StringBuilder result = new StringBuilder("[JvmInvokeTaintSink] ").append(signature);
         if (takesInstance)
         {
             result.append(", takes instance");
@@ -143,6 +185,7 @@ public String toString()
                                       .collect(Collectors.joining(", ")))
                   .append(")");
         }
+        callMatcher.ifPresent(p -> result.append(" filtered by ").append(p));
 
         return result.toString();
     }