Skip to content

Commit

Permalink
Refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
@zhangt2333
zhangt2333 committed May 28, 2024
1 parent 1548d49 commit 7037253
Show file tree
Hide file tree
Showing 10 changed files with 175 additions and 250 deletions.
3 changes: 2 additions & 1 deletion src/main/java/pascal/taie/analysis/pta/PointerAnalysis.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
import pascal.taie.analysis.pta.plugin.EntryPointHandler;
import pascal.taie.analysis.pta.plugin.Plugin;
import pascal.taie.analysis.pta.plugin.ReferenceHandler;
import pascal.taie.analysis.pta.plugin.ResultProcessor;
import pascal.taie.analysis.pta.plugin.ThreadHandler;
import pascal.taie.analysis.pta.plugin.exception.ExceptionAnalysis;
import pascal.taie.analysis.pta.plugin.invokedynamic.InvokeDynamicAnalysis;
Expand Down Expand Up @@ -153,7 +154,7 @@ private static void setPlugin(Solver solver, AnalysisOptions options) {
if (options.getString("taint-config") != null) {
plugin.addPlugin(new TaintAnalysis());
}
// plugin.addPlugin(new ResultProcessor());
plugin.addPlugin(new ResultProcessor());
// add plugins specified in options
// noinspection unchecked
addPlugins(plugin, (List<String>) options.get("plugins"));
Expand Down
16 changes: 7 additions & 9 deletions src/main/java/pascal/taie/analysis/pta/core/cs/element/AbstractPointer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,16 +67,9 @@ public void setPointsToSet(PointsToSet pointsToSet) {
}

@Override
public void rmFromPointsToIf(Predicate<CSObj> predicate) {
public void removeObjsIf(Predicate<CSObj> filter) {
if (pointsToSet != null) {
pointsToSet.removeIf(predicate);
}
}

@Override
public void rmFromOutEdgesIf(Predicate<PointerFlowEdge> predicate) {
if (! outEdges.isEmpty()) {
outEdges.removeIf(predicate);
pointsToSet.removeIf(filter);
}
}

Expand Down Expand Up @@ -122,6 +115,11 @@ public PointerFlowEdge addEdge(PointerFlowEdge edge) {
return null;
}

@Override
public void removeEdgesIf(Predicate<PointerFlowEdge> filter) {
outEdges.removeIf(filter);
}

@Override
public Set<PointerFlowEdge> getOutEdges() {
return Collections.unmodifiableSet(new ArraySet<>(outEdges, true));
Expand Down
14 changes: 10 additions & 4 deletions src/main/java/pascal/taie/analysis/pta/core/cs/element/Pointer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,15 +52,16 @@ public interface Pointer extends Indexable {
@Nullable
PointsToSet getPointsToSet();

void rmFromPointsToIf(Predicate<CSObj> predicate);

void rmFromOutEdgesIf(Predicate<PointerFlowEdge> predicate);

/**
* Sets the associated points-to set of this pointer.
*/
void setPointsToSet(PointsToSet pointsToSet);

/**
* Removes objects pointed to by this pointer if they satisfy the filter.
*/
void removeObjsIf(Predicate<CSObj> filter);

/**
* Adds filter to filter out objects pointed to by this pointer.
*/
Expand Down Expand Up @@ -101,6 +102,11 @@ public interface Pointer extends Indexable {
*/
PointerFlowEdge addEdge(PointerFlowEdge edge);

/**
* Removes out edges of this pointer if they satisfy the filter.
*/
void removeEdgesIf(Predicate<PointerFlowEdge> filter);

/**
* @return out edges of this pointer in pointer flow graph.
*/
Expand Down
2 changes: 1 addition & 1 deletion src/main/java/pascal/taie/analysis/pta/plugin/CompositePlugin.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ private void addPlugin(Plugin plugin, List<Plugin> plugins,
}
}

public void clearAllPlugins() {
public void clearPlugins() {
allPlugins.clear();
onNewPointsToSetPlugins.clear();
onNewCallEdgePlugins.clear();
Expand Down
78 changes: 26 additions & 52 deletions src/main/java/pascal/taie/analysis/pta/plugin/taint/SinkHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,22 +25,22 @@
import pascal.taie.analysis.graph.callgraph.CallKind;
import pascal.taie.analysis.graph.callgraph.Edge;
import pascal.taie.analysis.pta.PointerAnalysisResult;
import pascal.taie.analysis.pta.core.cs.element.*;
import pascal.taie.analysis.pta.core.heap.HeapModel;
import pascal.taie.analysis.pta.core.cs.element.ArrayIndex;
import pascal.taie.analysis.pta.core.cs.element.CSObj;
import pascal.taie.analysis.pta.core.cs.element.InstanceField;
import pascal.taie.analysis.pta.core.cs.element.Pointer;
import pascal.taie.analysis.pta.core.heap.Obj;
import pascal.taie.analysis.pta.plugin.util.InvokeUtils;
import pascal.taie.ir.exp.Var;
import pascal.taie.ir.stmt.Invoke;
import pascal.taie.language.classes.JField;
import pascal.taie.language.classes.JMethod;
import pascal.taie.util.Canonicalizer;
import pascal.taie.util.collection.*;
import pascal.taie.util.collection.MultiMap;
import pascal.taie.util.collection.MultiMapCollector;
import pascal.taie.util.collection.Sets;

import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
* Handles sinks in taint analysis.
Expand All @@ -49,17 +49,9 @@ class SinkHandler extends Handler {

private final List<Sink> sinks;

private final CSManager csManager;

private final HeapModel heapModel;

private final PointerToSetManager ptsManager = new PointerToSetManager();

SinkHandler(HandlerContext context) {
super(context);
sinks = context.config().sinks();
csManager = context.solver().getCSManager();
heapModel = context.solver().getHeapModel();
}

Set<TaintFlow> collectTaintFlows() {
Expand Down Expand Up @@ -100,9 +92,25 @@ private Set<TaintFlow> collectTaintFlows(
SinkPoint sinkPoint = new SinkPoint(sinkCall, indexRef);
// obtain objects to check for different IndexRef.Kind
Set<Obj> objs = switch (indexRef.kind()) {
case VAR -> ptsManager.getPointsToSet(arg);
case ARRAY -> ptsManager.getPointsToSet(arg, (Var) null);
case FIELD -> ptsManager.getPointsToSet(arg, indexRef.field());
case VAR -> csManager.getCSVarsOf(arg)
.stream()
.flatMap(Pointer::objects)
.map(CSObj::getObject)
.collect(Collectors.toUnmodifiableSet());
case ARRAY -> csManager.getCSVarsOf(arg)
.stream()
.flatMap(Pointer::objects)
.map(csManager::getArrayIndex)
.flatMap(ArrayIndex::objects)
.map(CSObj::getObject)
.collect(Collectors.toUnmodifiableSet());
case FIELD -> csManager.getCSVarsOf(arg)
.stream()
.flatMap(Pointer::objects)
.map(o -> csManager.getInstanceField(o, indexRef.field()))
.flatMap(InstanceField::objects)
.map(CSObj::getObject)
.collect(Collectors.toUnmodifiableSet());
};
return objs.stream()
.filter(manager::isTaint)
Expand All @@ -111,38 +119,4 @@ private Set<TaintFlow> collectTaintFlows(
.collect(Collectors.toSet());
}

class PointerToSetManager {
// todo: don't repeat urself!

private static final Canonicalizer<Set<Obj>> canonicalizer = new Canonicalizer<>();

private Set<Obj> removeContexts(Stream<CSObj> objects) {
Set<Obj> set = new HybridBitSet<>(heapModel, true);
objects.map(CSObj::getObject).forEach(set::add);
return canonicalizer.get(Collections.unmodifiableSet(set));
}

Set<Obj> getPointsToSet(Var var) {
return removeContexts(csManager.getCSVarsOf(var)
.stream()
.flatMap(Pointer::objects));
}

Set<Obj> getPointsToSet(Var base, JField field) {
return removeContexts(csManager.getCSVarsOf(base)
.stream()
.flatMap(Pointer::objects)
.map(o -> csManager.getInstanceField(o, field))
.flatMap(InstanceField::objects));
}

Set<Obj> getPointsToSet(Var base, Var index) {
return removeContexts(csManager.getCSVarsOf(base)
.stream()
.flatMap(Pointer::objects)
.map(csManager::getArrayIndex)
.flatMap(ArrayIndex::objects));
}

}
}
Loading

0 comments on commit 7037253

Please sign in to comment.