Skip to content

Commit

Permalink
fix: prevent endless loop in pre header insertion mod (#2300)
Browse files Browse the repository at this point in the history
  • Loading branch information
@skylot
skylot committed Oct 31, 2024
1 parent cfbe5ab commit 7544d1a
Show file tree
Hide file tree
Showing 4 changed files with 184 additions and 47 deletions.
140 changes: 93 additions & 47 deletions jadx-core/src/main/java/jadx/core/dex/visitors/blocks/BlockProcessor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,14 @@
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import jadx.api.plugins.input.data.attributes.IJadxAttrType;
import jadx.api.plugins.input.data.attributes.IJadxAttribute;
import jadx.core.dex.attributes.AFlag;
import jadx.core.dex.attributes.AType;
import jadx.core.dex.attributes.nodes.LoopInfo;
Expand All @@ -36,6 +39,8 @@
public class BlockProcessor extends AbstractVisitor {
private static final Logger LOG = LoggerFactory.getLogger(BlockProcessor.class);

private static final boolean DEBUG_MODS = false;

@Override
public void visit(MethodNode mth) {
if (mth.isNoCode() || mth.getBasicBlocks().isEmpty()) {
Expand All @@ -57,13 +62,27 @@ private static void processBlocksTree(MethodNode mth) {
}
updateCleanSuccessors(mth);

int blocksCount = mth.getBasicBlocks().size();
int modLimit = Math.max(100, blocksCount);
if (DEBUG_MODS) {
mth.addAttr(new DebugModAttr());
}
int i = 0;
while (modifyBlocksTree(mth)) {
computeDominators(mth);
if (i++ > 100) {
throw new JadxRuntimeException("CFG modification limit reached, blocks count: " + mth.getBasicBlocks().size());
if (i++ > modLimit) {
mth.addWarn("CFG modification limit reached, blocks count: " + blocksCount);
break;
}
}
if (DEBUG_MODS && i != 0) {
String stats = "CFG modifications count: " + i
+ ", blocks count: " + blocksCount + '\n'
+ mth.get(DebugModAttr.TYPE).formatStats() + '\n';
mth.addDebugComment(stats);
LOG.debug("Method: {}\n{}", mth, stats);
mth.remove(DebugModAttr.TYPE);
}
checkForUnreachableBlocks(mth);

DominatorTree.computeDominanceFrontier(mth);
Expand Down Expand Up @@ -298,6 +317,9 @@ private static boolean mergeConstReturn(MethodNode mth) {
}
if (changed) {
removeMarkedBlocks(mth);
if (DEBUG_MODS) {
mth.get(DebugModAttr.TYPE).addEvent("Merge const return");
}
}
return changed;
}
Expand Down Expand Up @@ -338,17 +360,20 @@ private static boolean independentBlockTreeMod(MethodNode mth) {

private static boolean simplifyLoopEnd(MethodNode mth, LoopInfo loop) {
BlockNode loopEnd = loop.getEnd();
if (loopEnd.getSuccessors().size() > 1) {
// make loop end a simple path block
BlockNode newLoopEnd = BlockSplitter.startNewBlock(mth, -1);
newLoopEnd.add(AFlag.SYNTHETIC);
newLoopEnd.add(AFlag.LOOP_END);
BlockNode loopStart = loop.getStart();
BlockSplitter.replaceConnection(loopEnd, loopStart, newLoopEnd);
BlockSplitter.connect(newLoopEnd, loopStart);
return true;
if (loopEnd.getSuccessors().size() <= 1) {
return false;
}
return false;
// make loop end a simple path block
BlockNode newLoopEnd = BlockSplitter.startNewBlock(mth, -1);
newLoopEnd.add(AFlag.SYNTHETIC);
newLoopEnd.add(AFlag.LOOP_END);
BlockNode loopStart = loop.getStart();
BlockSplitter.replaceConnection(loopEnd, loopStart, newLoopEnd);
BlockSplitter.connect(newLoopEnd, loopStart);
if (DEBUG_MODS) {
mth.get(DebugModAttr.TYPE).addEvent("Simplify loop end");
}
return true;
}

private static boolean checkLoops(MethodNode mth, BlockNode block) {
Expand All @@ -371,7 +396,6 @@ private static boolean checkLoops(MethodNode mth, BlockNode block) {
if (loopsCount == 1) {
LoopInfo loop = loops.get(0);
return insertBlocksForContinue(mth, loop)
|| insertBlockForPredecessors(mth, loop)
|| insertPreHeader(mth, loop)
|| simplifyLoopEnd(mth, loop);
}
Expand All @@ -398,15 +422,21 @@ private static boolean insertPreHeader(MethodNode mth, LoopInfo loop) {
mth.setEnterBlock(newEnterBlock);
start.remove(AFlag.MTH_ENTER_BLOCK);
BlockSplitter.connect(newEnterBlock, start);
return true;
} else {
// multiple predecessors
BlockNode preHeader = BlockSplitter.startNewBlock(mth, -1);
preHeader.add(AFlag.SYNTHETIC);
BlockNode loopEnd = loop.getEnd();
for (BlockNode pred : new ArrayList<>(preds)) {
if (pred != loopEnd) {
BlockSplitter.replaceConnection(pred, start, preHeader);
}
}
BlockSplitter.connect(preHeader, start);
}
// multiple predecessors
BlockNode preHeader = BlockSplitter.startNewBlock(mth, -1);
preHeader.add(AFlag.SYNTHETIC);
for (BlockNode pred : new ArrayList<>(preds)) {
BlockSplitter.replaceConnection(pred, start, preHeader);
if (DEBUG_MODS) {
mth.get(DebugModAttr.TYPE).addEvent("Insert loop pre header");
}
BlockSplitter.connect(preHeader, start);
return true;
}

Expand All @@ -426,6 +456,9 @@ private static boolean insertBlocksForBreak(MethodNode mth, LoopInfo loop) {
}
}
}
if (DEBUG_MODS && change) {
mth.get(DebugModAttr.TYPE).addEvent("Insert loop break blocks");
}
return change;
}

Expand All @@ -444,24 +477,10 @@ private static boolean insertBlocksForContinue(MethodNode mth, LoopInfo loop) {
}
}
}
return change;
}

/**
* Insert additional block if loop header has several predecessors (exclude back edges)
*/
private static boolean insertBlockForPredecessors(MethodNode mth, LoopInfo loop) {
BlockNode loopHeader = loop.getStart();
List<BlockNode> preds = loopHeader.getPredecessors();
if (preds.size() > 2) {
List<BlockNode> blocks = new ArrayList<>(preds);
blocks.removeIf(block -> block.contains(AFlag.LOOP_END));
BlockNode first = blocks.remove(0);
BlockNode preHeader = BlockSplitter.insertBlockBetween(mth, first, loopHeader);
blocks.forEach(block -> BlockSplitter.replaceConnection(block, loopHeader, preHeader));
return true;
if (DEBUG_MODS && change) {
mth.get(DebugModAttr.TYPE).addEvent("Insert loop continue block");
}
return false;
return change;
}

private static boolean splitLoops(MethodNode mth, BlockNode block, List<LoopInfo> loops) {
Expand All @@ -472,17 +491,20 @@ private static boolean splitLoops(MethodNode mth, BlockNode block, List<LoopInfo
break;
}
}
if (oneHeader) {
// several back edges connected to one loop header => make additional block
BlockNode newLoopEnd = BlockSplitter.startNewBlock(mth, block.getStartOffset());
newLoopEnd.add(AFlag.SYNTHETIC);
connect(newLoopEnd, block);
for (LoopInfo la : loops) {
BlockSplitter.replaceConnection(la.getEnd(), block, newLoopEnd);
}
return true;
if (!oneHeader) {
return false;
}
return false;
// several back edges connected to one loop header => make additional block
BlockNode newLoopEnd = BlockSplitter.startNewBlock(mth, block.getStartOffset());
newLoopEnd.add(AFlag.SYNTHETIC);
connect(newLoopEnd, block);
for (LoopInfo la : loops) {
BlockSplitter.replaceConnection(la.getEnd(), block, newLoopEnd);
}
if (DEBUG_MODS) {
mth.get(DebugModAttr.TYPE).addEvent("Split loops");
}
return true;
}

private static boolean splitExitBlocks(MethodNode mth) {
Expand All @@ -496,6 +518,9 @@ private static boolean splitExitBlocks(MethodNode mth) {
}
if (changed) {
updateExitBlockConnections(mth);
if (DEBUG_MODS) {
mth.get(DebugModAttr.TYPE).addEvent("Split exit block");
}
}
return changed;
}
Expand Down Expand Up @@ -691,4 +716,25 @@ private static void clearBlocksState(MethodNode mth) {
block.getDominatesOn().clear();
});
}

private static final class DebugModAttr implements IJadxAttribute {
static final IJadxAttrType<DebugModAttr> TYPE = IJadxAttrType.create("DebugModAttr");

private final Map<String, Integer> statMap = new HashMap<>();

public void addEvent(String name) {
statMap.merge(name, 1, Integer::sum);
}

public String formatStats() {
return statMap.entrySet().stream()
.map(entry -> " " + entry.getKey() + ": " + entry.getValue())
.collect(Collectors.joining("\n"));
}

@Override
public IJadxAttrType<DebugModAttr> getAttrType() {
return TYPE;
}
}
}
18 changes: 18 additions & 0 deletions jadx-core/src/test/java/jadx/tests/integration/loops/TestLoopRestore3.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
package jadx.tests.integration.loops;

import org.junit.jupiter.api.Test;

import jadx.tests.api.SmaliTest;

import static jadx.tests.api.utils.assertj.JadxAssertions.assertThat;

public class TestLoopRestore3 extends SmaliTest {

@Test
public void test() {
disableCompilation();
assertThat(getClassNodeFromSmali())
.code()
.countString(3, "while (");
}
}
64 changes: 64 additions & 0 deletions jadx-core/src/test/smali/loops/TestLoopRestore3.smali
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
.class public Lloops/TestLoopRestore3;
.super Ljava/lang/Object;

.method public final b(Ljava/lang/String;Lb/U53$b;)V
.registers 8

iget-object v0, p0, Lb/X53;->e:Ljava/util/concurrent/atomic/AtomicReference;
:goto_2
invoke-virtual {v0}, Ljava/util/concurrent/atomic/AtomicReference;->get()Ljava/lang/Object;
move-result-object v1
move-object v2, v1
check-cast v2, Ljava/util/List;
move-object v3, v2
check-cast v3, Ljava/lang/Iterable;
instance-of v4, v3, Ljava/util/Collection;
if-eqz v4, :cond_1a

move-object v4, v3
check-cast v4, Ljava/util/Collection;
invoke-interface {v4}, Ljava/util/Collection;->isEmpty()Z
move-result v4
if-eqz v4, :cond_1a
goto :goto_33

:cond_1a
invoke-interface {v3}, Ljava/lang/Iterable;->iterator()Ljava/util/Iterator;
move-result-object v3

:cond_1e
invoke-interface {v3}, Ljava/util/Iterator;->hasNext()Z
move-result v4
if-eqz v4, :cond_33

invoke-interface {v3}, Ljava/util/Iterator;->next()Ljava/lang/Object;
move-result-object v4
check-cast v4, Lb/X53$c;
iget-object v4, v4, Lb/X53$c;->b:Ljava/lang/String;
invoke-static {v4, p1}, Lkotlin/jvm/internal/Intrinsics;->a(Ljava/lang/Object;Ljava/lang/Object;)Z
move-result v4
if-eqz v4, :cond_1e
goto :goto_40

:cond_33
:goto_33
check-cast v2, Ljava/util/Collection;
new-instance v3, Lb/X53$c;
sget-object v4, Lb/Pd2;->a:Lb/Pd2;
invoke-direct {v3, p2, p1, v4}, Lb/X53$c;-><init>(Lb/U53$b;Ljava/lang/String;Ljava/util/List;)V
invoke-static {v2, v3}, Lb/R31;->a0(Ljava/util/Collection;Ljava/lang/Object;)Ljava/util/ArrayList;
move-result-object v2

:cond_40
:goto_40
invoke-virtual {v0, v1, v2}, Ljava/util/concurrent/atomic/AtomicReference;->compareAndSet(Ljava/lang/Object;Ljava/lang/Object;)Z
move-result v3
if-eqz v3, :cond_47
return-void

:cond_47
invoke-virtual {v0}, Ljava/util/concurrent/atomic/AtomicReference;->get()Ljava/lang/Object;
move-result-object v3
if-eq v3, v1, :cond_40
goto :goto_2
.end method
9 changes: 9 additions & 0 deletions ...ns/jadx-input-api/src/main/java/jadx/api/plugins/input/data/attributes/IJadxAttrType.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,4 +18,13 @@ static <A extends IJadxAttribute> IJadxAttrType<A> create() {
return new IJadxAttrType<>() {
};
}

static <A extends IJadxAttribute> IJadxAttrType<A> create(String name) {
return new IJadxAttrType<>() {
@Override
public String toString() {
return name;
}
};
}
}

0 comments on commit 7544d1a

Please sign in to comment.