From 8a4ec47b929f4d26936c0ac3d26afd24fd7df1b3 Mon Sep 17 00:00:00 2001 From: Skylot Date: Mon, 29 Sep 2014 23:44:36 +0400 Subject: [PATCH] core: support break with label for simple cases --- .../main/java/jadx/core/codegen/InsnGen.java | 5 ++ .../main/java/jadx/core/codegen/NameGen.java | 8 +++ .../java/jadx/core/codegen/RegionGen.java | 5 ++ .../java/jadx/core/dex/attributes/AType.java | 4 ++ .../core/dex/attributes/AttributeStorage.java | 4 +- .../core/dex/attributes/nodes/LoopInfo.java | 21 +++++++- .../dex/attributes/nodes/LoopLabelAttr.java | 27 ++++++++++ .../java/jadx/core/dex/nodes/MethodNode.java | 25 +++++++++- .../core/dex/regions/loops/LoopRegion.java | 9 +++- .../core/dex/visitors/BlockMakerVisitor.java | 30 +++++++++++- .../dex/visitors/regions/RegionMaker.java | 34 ++++++++++--- .../integration/loops/TestBreakWithLabel.java | 49 +++++++++++++++++++ 12 files changed, 206 insertions(+), 15 deletions(-) create mode 100644 jadx-core/src/main/java/jadx/core/dex/attributes/nodes/LoopLabelAttr.java create mode 100644 jadx-core/src/test/java/jadx/tests/integration/loops/TestBreakWithLabel.java diff --git a/jadx-core/src/main/java/jadx/core/codegen/InsnGen.java b/jadx-core/src/main/java/jadx/core/codegen/InsnGen.java index 2da48095c..22291c7b0 100644 --- a/jadx-core/src/main/java/jadx/core/codegen/InsnGen.java +++ b/jadx-core/src/main/java/jadx/core/codegen/InsnGen.java @@ -3,6 +3,7 @@ package jadx.core.codegen; import jadx.core.dex.attributes.AFlag; import jadx.core.dex.attributes.AType; import jadx.core.dex.attributes.nodes.FieldReplaceAttr; +import jadx.core.dex.attributes.nodes.LoopLabelAttr; import jadx.core.dex.attributes.nodes.MethodInlineAttr; import jadx.core.dex.info.ClassInfo; import jadx.core.dex.info.FieldInfo; @@ -289,6 +290,10 @@ public class InsnGen { case BREAK: code.add("break"); + LoopLabelAttr labelAttr = insn.get(AType.LOOP_LABEL); + if (labelAttr != null) { + code.add(' ').add(mgen.getNameGen().getLoopLabel(labelAttr)); + } break; case CONTINUE: diff --git a/jadx-core/src/main/java/jadx/core/codegen/NameGen.java b/jadx-core/src/main/java/jadx/core/codegen/NameGen.java index 4a3eb3126..1db788247 100644 --- a/jadx-core/src/main/java/jadx/core/codegen/NameGen.java +++ b/jadx-core/src/main/java/jadx/core/codegen/NameGen.java @@ -2,6 +2,7 @@ package jadx.core.codegen; import jadx.core.Consts; import jadx.core.deobf.NameMapper; +import jadx.core.dex.attributes.nodes.LoopLabelAttr; import jadx.core.dex.info.ClassInfo; import jadx.core.dex.instructions.InvokeNode; import jadx.core.dex.instructions.args.ArgType; @@ -76,6 +77,13 @@ public class NameGen { return name; } + // TODO: avoid name collision with variables names + public String getLoopLabel(LoopLabelAttr attr) { + String name = "loop" + attr.getLoop().getId(); + varNames.add(name); + return name; + } + private String getUniqueVarName(String name) { String r = name; int i = 2; diff --git a/jadx-core/src/main/java/jadx/core/codegen/RegionGen.java b/jadx-core/src/main/java/jadx/core/codegen/RegionGen.java index 6891d7196..15d78f76f 100644 --- a/jadx-core/src/main/java/jadx/core/codegen/RegionGen.java +++ b/jadx-core/src/main/java/jadx/core/codegen/RegionGen.java @@ -4,6 +4,7 @@ import jadx.core.dex.attributes.AFlag; import jadx.core.dex.attributes.AType; import jadx.core.dex.attributes.nodes.DeclareVariablesAttr; import jadx.core.dex.attributes.nodes.ForceReturnAttr; +import jadx.core.dex.attributes.nodes.LoopLabelAttr; import jadx.core.dex.info.FieldInfo; import jadx.core.dex.instructions.IndexInsnNode; import jadx.core.dex.instructions.SwitchNode; @@ -164,6 +165,10 @@ public class RegionGen extends InsnGen { } } } + LoopLabelAttr labelAttr = region.getInfo().getStart().get(AType.LOOP_LABEL); + if (labelAttr != null) { + code.startLine(mgen.getNameGen().getLoopLabel(labelAttr)).add(':'); + } IfCondition condition = region.getCondition(); if (condition == null) { diff --git a/jadx-core/src/main/java/jadx/core/dex/attributes/AType.java b/jadx-core/src/main/java/jadx/core/dex/attributes/AType.java index 6a3e6231e..8f75f0f0f 100644 --- a/jadx-core/src/main/java/jadx/core/dex/attributes/AType.java +++ b/jadx-core/src/main/java/jadx/core/dex/attributes/AType.java @@ -10,6 +10,7 @@ import jadx.core.dex.attributes.nodes.ForceReturnAttr; import jadx.core.dex.attributes.nodes.JadxErrorAttr; import jadx.core.dex.attributes.nodes.JumpInfo; import jadx.core.dex.attributes.nodes.LoopInfo; +import jadx.core.dex.attributes.nodes.LoopLabelAttr; import jadx.core.dex.attributes.nodes.MethodInlineAttr; import jadx.core.dex.attributes.nodes.PhiListAttr; import jadx.core.dex.attributes.nodes.SourceFileAttr; @@ -29,6 +30,8 @@ public class AType { private AType() { } + public static final int FIELDS_COUNT = 18; + public static final AType> JUMP = new AType>(); public static final AType> LOOP = new AType>(); @@ -47,4 +50,5 @@ public class AType { public static final AType PHI_LIST = new AType(); public static final AType SOURCE_FILE = new AType(); public static final AType DECLARE_VARIABLES = new AType(); + public static final AType LOOP_LABEL = new AType(); } diff --git a/jadx-core/src/main/java/jadx/core/dex/attributes/AttributeStorage.java b/jadx-core/src/main/java/jadx/core/dex/attributes/AttributeStorage.java index 0dcc4941d..a34baf820 100644 --- a/jadx-core/src/main/java/jadx/core/dex/attributes/AttributeStorage.java +++ b/jadx-core/src/main/java/jadx/core/dex/attributes/AttributeStorage.java @@ -7,7 +7,7 @@ import jadx.core.utils.Utils; import java.util.ArrayList; import java.util.Collections; import java.util.EnumSet; -import java.util.HashMap; +import java.util.IdentityHashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -24,7 +24,7 @@ public class AttributeStorage { public AttributeStorage() { flags = EnumSet.noneOf(AFlag.class); - attributes = new HashMap, IAttribute>(2); + attributes = new IdentityHashMap, IAttribute>(AType.FIELDS_COUNT); } public void add(AFlag flag) { diff --git a/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/LoopInfo.java b/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/LoopInfo.java index c868855e7..3cdd007c2 100644 --- a/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/LoopInfo.java +++ b/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/LoopInfo.java @@ -17,6 +17,9 @@ public class LoopInfo { private final BlockNode end; private final Set loopBlocks; + private int id; + private LoopInfo parentLoop; + public LoopInfo(BlockNode start, BlockNode end) { this.start = start; this.end = end; @@ -69,8 +72,24 @@ public class LoopInfo { return edges; } + public int getId() { + return id; + } + + public void setId(int id) { + this.id = id; + } + + public LoopInfo getParentLoop() { + return parentLoop; + } + + public void setParentLoop(LoopInfo parentLoop) { + this.parentLoop = parentLoop; + } + @Override public String toString() { - return "LOOP: " + start + "->" + end; + return "LOOP:" + id + ": " + start + "->" + end; } } diff --git a/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/LoopLabelAttr.java b/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/LoopLabelAttr.java new file mode 100644 index 000000000..15381f391 --- /dev/null +++ b/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/LoopLabelAttr.java @@ -0,0 +1,27 @@ +package jadx.core.dex.attributes.nodes; + +import jadx.core.dex.attributes.AType; +import jadx.core.dex.attributes.IAttribute; + +public class LoopLabelAttr implements IAttribute { + + private final LoopInfo loop; + + public LoopLabelAttr(LoopInfo loop) { + this.loop = loop; + } + + public LoopInfo getLoop() { + return loop; + } + + @Override + public AType getType() { + return AType.LOOP_LABEL; + } + + @Override + public String toString() { + return "LOOP_LABEL: " + loop; + } +} diff --git a/jadx-core/src/main/java/jadx/core/dex/nodes/MethodNode.java b/jadx-core/src/main/java/jadx/core/dex/nodes/MethodNode.java index 5cd1b4694..8b56b0288 100644 --- a/jadx-core/src/main/java/jadx/core/dex/nodes/MethodNode.java +++ b/jadx-core/src/main/java/jadx/core/dex/nodes/MethodNode.java @@ -406,10 +406,14 @@ public class MethodNode extends LineAttrNode implements ILoadable { if (loops.isEmpty()) { loops = new ArrayList(5); } + loop.setId(loops.size()); loops.add(loop); } public LoopInfo getLoopForBlock(BlockNode block) { + if (loops.isEmpty()) { + return null; + } for (LoopInfo loop : loops) { if (loop.getLoopBlocks().contains(block)) { return loop; @@ -418,10 +422,27 @@ public class MethodNode extends LineAttrNode implements ILoadable { return null; } + public List getAllLoopsForBlock(BlockNode block) { + if (loops.isEmpty()) { + return Collections.emptyList(); + } + List list = new ArrayList(loops.size()); + for (LoopInfo loop : loops) { + if (loop.getLoopBlocks().contains(block)) { + list.add(loop); + } + } + return list; + } + public int getLoopsCount() { return loops.size(); } + public Iterable getLoops() { + return loops; + } + public ExceptionHandler addExceptionHandler(ExceptionHandler handler) { if (exceptionHandlers.isEmpty()) { exceptionHandlers = new ArrayList(2); @@ -468,7 +489,7 @@ public class MethodNode extends LineAttrNode implements ILoadable { } return false; } - + public boolean isDefaultConstructor() { boolean result = false; @@ -492,7 +513,7 @@ public class MethodNode extends LineAttrNode implements ILoadable { } return result; - } + } public int getRegsCount() { return regsCount; diff --git a/jadx-core/src/main/java/jadx/core/dex/regions/loops/LoopRegion.java b/jadx-core/src/main/java/jadx/core/dex/regions/loops/LoopRegion.java index 2623672c3..855bc01b4 100644 --- a/jadx-core/src/main/java/jadx/core/dex/regions/loops/LoopRegion.java +++ b/jadx-core/src/main/java/jadx/core/dex/regions/loops/LoopRegion.java @@ -1,5 +1,6 @@ package jadx.core.dex.regions.loops; +import jadx.core.dex.attributes.nodes.LoopInfo; import jadx.core.dex.instructions.IfNode; import jadx.core.dex.instructions.args.RegisterArg; import jadx.core.dex.nodes.BlockNode; @@ -15,6 +16,7 @@ import java.util.List; public final class LoopRegion extends AbstractRegion { + private final LoopInfo info; // loop header contains one 'if' insn, equals null for infinite loop private IfCondition condition; private final BlockNode conditionBlock; @@ -25,13 +27,18 @@ public final class LoopRegion extends AbstractRegion { private LoopType type; - public LoopRegion(IRegion parent, BlockNode header, boolean reversed) { + public LoopRegion(IRegion parent, LoopInfo info, BlockNode header, boolean reversed) { super(parent); + this.info = info; this.conditionBlock = header; this.condition = IfCondition.fromIfBlock(header); this.conditionAtEnd = reversed; } + public LoopInfo getInfo() { + return info; + } + public IfCondition getCondition() { return condition; } diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/BlockMakerVisitor.java b/jadx-core/src/main/java/jadx/core/dex/visitors/BlockMakerVisitor.java index 46aa8426b..384f53304 100644 --- a/jadx-core/src/main/java/jadx/core/dex/visitors/BlockMakerVisitor.java +++ b/jadx-core/src/main/java/jadx/core/dex/visitors/BlockMakerVisitor.java @@ -190,6 +190,7 @@ public class BlockMakerVisitor extends AbstractVisitor { } computeDominanceFrontier(mth); registerLoops(mth); + processNestedLoops(mth); } private static BlockNode getBlock(int offset, Map blocksMap) { @@ -357,15 +358,40 @@ public class BlockMakerVisitor extends AbstractVisitor { private static void registerLoops(MethodNode mth) { for (BlockNode block : mth.getBasicBlocks()) { - List loops = block.getAll(AType.LOOP); if (block.contains(AFlag.LOOP_START)) { - for (LoopInfo loop : loops) { + for (LoopInfo loop : block.getAll(AType.LOOP)) { mth.registerLoop(loop); } } } } + private static void processNestedLoops(MethodNode mth) { + if (mth.getLoopsCount() == 0) { + return; + } + for (LoopInfo outLoop : mth.getLoops()) { + for (LoopInfo innerLoop : mth.getLoops()) { + if (outLoop == innerLoop) { + continue; + } + if (outLoop.getLoopBlocks().containsAll(innerLoop.getLoopBlocks())) { + LoopInfo parentLoop = innerLoop.getParentLoop(); + if (parentLoop != null) { + if (parentLoop.getLoopBlocks().containsAll(outLoop.getLoopBlocks())) { + outLoop.setParentLoop(parentLoop); + innerLoop.setParentLoop(outLoop); + } else { + parentLoop.setParentLoop(outLoop); + } + } else { + innerLoop.setParentLoop(outLoop); + } + } + } + } + } + private static boolean modifyBlocksTree(MethodNode mth) { for (BlockNode block : mth.getBasicBlocks()) { if (block.getPredecessors().isEmpty() && block != mth.getEnterBlock()) { diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/regions/RegionMaker.java b/jadx-core/src/main/java/jadx/core/dex/visitors/regions/RegionMaker.java index 3c3858f90..d5f7d97a7 100644 --- a/jadx-core/src/main/java/jadx/core/dex/visitors/regions/RegionMaker.java +++ b/jadx-core/src/main/java/jadx/core/dex/visitors/regions/RegionMaker.java @@ -4,6 +4,7 @@ import jadx.core.Consts; import jadx.core.dex.attributes.AFlag; import jadx.core.dex.attributes.AType; import jadx.core.dex.attributes.nodes.LoopInfo; +import jadx.core.dex.attributes.nodes.LoopLabelAttr; import jadx.core.dex.instructions.IfNode; import jadx.core.dex.instructions.InsnType; import jadx.core.dex.instructions.SwitchNode; @@ -13,12 +14,12 @@ import jadx.core.dex.nodes.Edge; import jadx.core.dex.nodes.IRegion; import jadx.core.dex.nodes.InsnNode; import jadx.core.dex.nodes.MethodNode; -import jadx.core.dex.regions.conditions.IfInfo; -import jadx.core.dex.regions.conditions.IfRegion; -import jadx.core.dex.regions.loops.LoopRegion; import jadx.core.dex.regions.Region; import jadx.core.dex.regions.SwitchRegion; import jadx.core.dex.regions.SynchronizedRegion; +import jadx.core.dex.regions.conditions.IfInfo; +import jadx.core.dex.regions.conditions.IfRegion; +import jadx.core.dex.regions.loops.LoopRegion; import jadx.core.dex.trycatch.ExcHandlerAttr; import jadx.core.dex.trycatch.ExceptionHandler; import jadx.core.dex.trycatch.TryCatchBlock; @@ -184,7 +185,7 @@ public class RegionMaker { if (exitBlocks.size() > 0) { BlockNode loopExit = condInfo.getElseBlock(); if (loopExit != null) { - // add 'break' instruction before path cross between main loop exit and subexit + // add 'break' instruction before path cross between main loop exit and sub-exit for (Edge exitEdge : loop.getExitEdges()) { if (!exitBlocks.contains(exitEdge.getSource())) { continue; @@ -245,7 +246,12 @@ public class RegionMaker { || block.getInstructions().get(0).getType() != InsnType.IF) { continue; } - LoopRegion loopRegion = new LoopRegion(curRegion, block, block == loop.getEnd()); + List loops = block.getAll(AType.LOOP); + if (!loops.isEmpty() && loops.get(0) != loop) { + // skip nested loop condition + continue; + } + LoopRegion loopRegion = new LoopRegion(curRegion, loop, block, block == loop.getEnd()); boolean found; if (block == loop.getStart() || block == loop.getEnd() || BlockUtils.isEmptySimplePath(loop.getStart(), block)) { @@ -266,7 +272,7 @@ public class RegionMaker { } private BlockNode makeEndlessLoop(IRegion curRegion, RegionStack stack, LoopInfo loop, BlockNode loopStart) { - LoopRegion loopRegion = new LoopRegion(curRegion, null, false); + LoopRegion loopRegion = new LoopRegion(curRegion, loop, null, false); curRegion.getSubBlocks().add(loopRegion); loopStart.remove(AType.LOOP); @@ -332,8 +338,22 @@ public class RegionMaker { if (prev != null && isPathExists(loopExit, exit)) { // found cross if (canInsertBreak(exit)) { - prev.getInstructions().add(new InsnNode(InsnType.BREAK, 0)); + InsnNode breakInsn = new InsnNode(InsnType.BREAK, 0); + prev.getInstructions().add(breakInsn); stack.addExit(exit); + // add label to 'break' if needed + List loops = mth.getAllLoopsForBlock(exitEdge.getSource()); + if (loops.size() >= 2) { + // find parent loop + for (LoopInfo loop : loops) { + if (loop.getParentLoop() == null) { + LoopLabelAttr labelAttr = new LoopLabelAttr(loop); + breakInsn.addAttr(labelAttr); + loop.getStart().addAttr(labelAttr); + break; + } + } + } } return; } diff --git a/jadx-core/src/test/java/jadx/tests/integration/loops/TestBreakWithLabel.java b/jadx-core/src/test/java/jadx/tests/integration/loops/TestBreakWithLabel.java new file mode 100644 index 000000000..d44d011cb --- /dev/null +++ b/jadx-core/src/test/java/jadx/tests/integration/loops/TestBreakWithLabel.java @@ -0,0 +1,49 @@ +package jadx.tests.integration.loops; + +import jadx.core.dex.nodes.ClassNode; +import jadx.tests.api.IntegrationTest; + +import java.lang.reflect.Method; + +import org.junit.Test; + +import static jadx.tests.api.utils.JadxMatchers.containsOne; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThat; +import static org.junit.Assert.assertTrue; + +public class TestBreakWithLabel extends IntegrationTest { + + public static class TestCls { + + public boolean test(int[][] arr, int b) { + boolean found = false; + loop0: + for (int i = 0; i < arr.length; i++) { + for (int j = 0; j < arr[i].length; j++) { + if (arr[i][j] == b) { + found = true; + break loop0; + } + } + } + System.out.println("found: " + found); + return found; + } + } + + @Test + public void test() { + ClassNode cls = getClassNode(TestCls.class); + String code = cls.getCode().toString(); + System.out.println(code); + + assertThat(code, containsOne("loop0:")); + assertThat(code, containsOne("break loop0;")); + + Method test = getReflectMethod("test", int[][].class, int.class); + int[][] testArray = {{1, 2}, {3, 4}}; + assertTrue((Boolean) invoke(test, testArray, 3)); + assertFalse((Boolean) invoke(test, testArray, 5)); + } +}