From a27cb9c34e084cc0bfc1aa288626645dbe7e4f4d Mon Sep 17 00:00:00 2001 From: Skylot Date: Thu, 31 May 2018 14:48:12 +0300 Subject: [PATCH] core: bind blocks for target instructions at early stage --- .../jadx/core/dex/instructions/GotoNode.java | 13 +- .../jadx/core/dex/instructions/IfNode.java | 14 ++ .../core/dex/instructions/SwitchNode.java | 45 +++++- .../core/dex/instructions/TargetInsnNode.java | 15 ++ .../blocksmaker/BlockFinallyExtract.java | 4 +- .../dex/visitors/blocksmaker/BlockFinish.java | 18 --- .../visitors/blocksmaker/BlockProcessor.java | 129 +++++++++--------- .../visitors/blocksmaker/BlockSplitter.java | 35 ++++- .../dex/visitors/regions/RegionMaker.java | 44 ++---- .../main/java/jadx/core/utils/BlockUtils.java | 5 +- 10 files changed, 197 insertions(+), 125 deletions(-) create mode 100644 jadx-core/src/main/java/jadx/core/dex/instructions/TargetInsnNode.java diff --git a/jadx-core/src/main/java/jadx/core/dex/instructions/GotoNode.java b/jadx-core/src/main/java/jadx/core/dex/instructions/GotoNode.java index be6f6316..c4a424dd 100644 --- a/jadx-core/src/main/java/jadx/core/dex/instructions/GotoNode.java +++ b/jadx-core/src/main/java/jadx/core/dex/instructions/GotoNode.java @@ -1,9 +1,9 @@ package jadx.core.dex.instructions; -import jadx.core.dex.nodes.InsnNode; +import jadx.core.dex.nodes.BlockNode; import jadx.core.utils.InsnUtils; -public class GotoNode extends InsnNode { +public class GotoNode extends TargetInsnNode { protected int target; @@ -20,6 +20,15 @@ public class GotoNode extends InsnNode { return target; } + @Override + public boolean replaceTargetBlock(BlockNode origin, BlockNode replace) { + return false; + } + + @Override + public void initBlocks(BlockNode curBlock) { + } + @Override public String toString() { return super.toString() + "-> " + InsnUtils.formatOffset(target); diff --git a/jadx-core/src/main/java/jadx/core/dex/instructions/IfNode.java b/jadx-core/src/main/java/jadx/core/dex/instructions/IfNode.java index 24869652..68f97e32 100644 --- a/jadx-core/src/main/java/jadx/core/dex/instructions/IfNode.java +++ b/jadx-core/src/main/java/jadx/core/dex/instructions/IfNode.java @@ -56,6 +56,7 @@ public class IfNode extends GotoNode { setArg(1, arg2); } + @Override public void initBlocks(BlockNode curBlock) { thenBlock = getBlockByOffset(target, curBlock.getSuccessors()); if (curBlock.getSuccessors().size() == 1) { @@ -65,6 +66,19 @@ public class IfNode extends GotoNode { } } + @Override + public boolean replaceTargetBlock(BlockNode origin, BlockNode replace) { + if (thenBlock == origin) { + thenBlock = replace; + return true; + } + if (elseBlock == origin) { + elseBlock = replace; + return true; + } + return false; + } + public BlockNode getThenBlock() { return thenBlock; } diff --git a/jadx-core/src/main/java/jadx/core/dex/instructions/SwitchNode.java b/jadx-core/src/main/java/jadx/core/dex/instructions/SwitchNode.java index adb858ab..c9273afe 100644 --- a/jadx-core/src/main/java/jadx/core/dex/instructions/SwitchNode.java +++ b/jadx-core/src/main/java/jadx/core/dex/instructions/SwitchNode.java @@ -1,17 +1,24 @@ package jadx.core.dex.instructions; import java.util.Arrays; +import java.util.List; import jadx.core.dex.instructions.args.InsnArg; +import jadx.core.dex.nodes.BlockNode; import jadx.core.dex.nodes.InsnNode; import jadx.core.utils.InsnUtils; -public class SwitchNode extends InsnNode { +import static jadx.core.utils.BlockUtils.getBlockByOffset; + +public class SwitchNode extends TargetInsnNode { private final Object[] keys; private final int[] targets; private final int def; // next instruction + private BlockNode[] targetBlocks; + private BlockNode defTargetBlock; + public SwitchNode(InsnArg arg, Object[] keys, int[] targets, int def) { super(InsnType.SWITCH, 1); this.keys = keys; @@ -36,6 +43,42 @@ public class SwitchNode extends InsnNode { return def; } + public BlockNode[] getTargetBlocks() { + return targetBlocks; + } + + public BlockNode getDefTargetBlock() { + return defTargetBlock; + } + + @Override + public void initBlocks(BlockNode curBlock) { + List successors = curBlock.getSuccessors(); + int len = targets.length; + targetBlocks = new BlockNode[len]; + for (int i = 0; i < len; i++) { + targetBlocks[i] = getBlockByOffset(targets[i], successors); + } + defTargetBlock = getBlockByOffset(def, successors); + } + + @Override + public boolean replaceTargetBlock(BlockNode origin, BlockNode replace) { + int count = 0; + int len = targetBlocks.length; + for (int i = 0; i < len; i++) { + if (targetBlocks[i] == origin) { + targetBlocks[i] = replace; + count++; + } + } + if (defTargetBlock == origin) { + defTargetBlock = replace; + count++; + } + return count > 0; + } + @Override public boolean isSame(InsnNode obj) { if (this == obj) { diff --git a/jadx-core/src/main/java/jadx/core/dex/instructions/TargetInsnNode.java b/jadx-core/src/main/java/jadx/core/dex/instructions/TargetInsnNode.java new file mode 100644 index 00000000..939d32b0 --- /dev/null +++ b/jadx-core/src/main/java/jadx/core/dex/instructions/TargetInsnNode.java @@ -0,0 +1,15 @@ +package jadx.core.dex.instructions; + +import jadx.core.dex.nodes.BlockNode; +import jadx.core.dex.nodes.InsnNode; + +public abstract class TargetInsnNode extends InsnNode { + + public TargetInsnNode(InsnType type, int argsCount) { + super(type, argsCount); + } + + public abstract void initBlocks(BlockNode curBlock); + + public abstract boolean replaceTargetBlock(BlockNode origin, BlockNode replace); +} diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockFinallyExtract.java b/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockFinallyExtract.java index 6902db97..e08d28fb 100644 --- a/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockFinallyExtract.java +++ b/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockFinallyExtract.java @@ -577,8 +577,7 @@ public class BlockFinallyExtract extends AbstractVisitor { BlockNode newPred = BlockSplitter.insertBlockBetween(mth, pred, sOut); for (BlockNode predBlock : new ArrayList<>(sOut.getPredecessors())) { if (predBlock != newPred) { - removeConnection(predBlock, sOut); - connect(predBlock, newPred); + BlockSplitter.replaceConnection(predBlock, sOut, newPred); } } rOut.getPredecessors().clear(); @@ -605,6 +604,7 @@ public class BlockFinallyExtract extends AbstractVisitor { connect(middle, startBlock); addIgnoredEdge(middle, startBlock); connect(middle, rOut); + BlockSplitter.replaceTarget(middle, remBlock, rOut); } // mark blocks for remove diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockFinish.java b/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockFinish.java index 9fb8b7ae..96db02e9 100644 --- a/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockFinish.java +++ b/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockFinish.java @@ -1,17 +1,13 @@ package jadx.core.dex.visitors.blocksmaker; import java.util.HashMap; -import java.util.List; import java.util.Map; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import jadx.core.dex.attributes.AType; -import jadx.core.dex.instructions.IfNode; -import jadx.core.dex.instructions.InsnType; import jadx.core.dex.nodes.BlockNode; -import jadx.core.dex.nodes.InsnNode; import jadx.core.dex.nodes.MethodNode; import jadx.core.dex.trycatch.ExcHandlerAttr; import jadx.core.dex.trycatch.SplitterBlockAttr; @@ -30,26 +26,12 @@ public class BlockFinish extends AbstractVisitor { for (BlockNode block : mth.getBasicBlocks()) { block.updateCleanSuccessors(); - initBlocksInIfNodes(block); fixSplitterBlock(block); } mth.finishBasicBlocks(); } - /** - * Init 'then' and 'else' blocks for 'if' instruction. - */ - private static void initBlocksInIfNodes(BlockNode block) { - List instructions = block.getInstructions(); - if (instructions.size() == 1) { - InsnNode insn = instructions.get(0); - if (insn.getType() == InsnType.IF) { - ((IfNode) insn).initBlocks(block); - } - } - } - /** * For evey exception handler must be only one splitter block, * select correct one and remove others if necessary. diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockProcessor.java b/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockProcessor.java index 31be3f41..4627e321 100644 --- a/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockProcessor.java +++ b/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockProcessor.java @@ -24,7 +24,6 @@ import jadx.core.utils.BlockUtils; import jadx.core.utils.exceptions.JadxRuntimeException; import static jadx.core.dex.visitors.blocksmaker.BlockSplitter.connect; -import static jadx.core.dex.visitors.blocksmaker.BlockSplitter.removeConnection; import static jadx.core.utils.EmptyBitSet.EMPTY; public class BlockProcessor extends AbstractVisitor { @@ -39,7 +38,7 @@ public class BlockProcessor extends AbstractVisitor { } public static void rerun(MethodNode mth) { - removeBlocks(mth); + removeMarkedBlocks(mth); clearBlocksState(mth); processBlocksTree(mth); } @@ -247,65 +246,69 @@ public class BlockProcessor extends AbstractVisitor { if (block.getPredecessors().isEmpty() && block != mth.getEnterBlock()) { throw new JadxRuntimeException("Unreachable block: " + block); } + if (checkLoops(mth, block)) { + return true; + } + } + return splitReturn(mth); + } - // check loops - List loops = block.getAll(AType.LOOP); - if (loops.size() > 1) { - boolean oneHeader = true; - for (LoopInfo loop : loops) { - if (loop.getStart() != block) { - oneHeader = false; - break; - } + private static boolean checkLoops(MethodNode mth, BlockNode block) { + // check loops + List loops = block.getAll(AType.LOOP); + if (loops.size() > 1) { + boolean oneHeader = true; + for (LoopInfo loop : loops) { + if (loop.getStart() != block) { + oneHeader = false; + 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); } - if (oneHeader) { - // several back edges connected to one loop header => make additional block - BlockNode newLoopHeader = BlockSplitter.startNewBlock(mth, block.getStartOffset()); - newLoopHeader.add(AFlag.SYNTHETIC); - connect(newLoopHeader, block); - for (LoopInfo la : loops) { - BlockNode node = la.getEnd(); - removeConnection(node, block); - connect(node, newLoopHeader); + return true; + } + } + if (loops.size() == 1) { + LoopInfo loop = loops.get(0); + // insert additional blocks for possible 'break' insertion + List edges = loop.getExitEdges(); + if (!edges.isEmpty()) { + boolean change = false; + for (Edge edge : edges) { + BlockNode target = edge.getTarget(); + if (!target.contains(AFlag.SYNTHETIC)) { + BlockSplitter.insertBlockBetween(mth, edge.getSource(), target); + change = true; } + } + if (change) { return true; } } - if (loops.size() == 1) { - LoopInfo loop = loops.get(0); - // insert additional blocks for possible 'break' insertion - List edges = loop.getExitEdges(); - if (!edges.isEmpty()) { - boolean change = false; - for (Edge edge : edges) { - BlockNode target = edge.getTarget(); - if (!target.contains(AFlag.SYNTHETIC)) { - BlockSplitter.insertBlockBetween(mth, edge.getSource(), target); - change = true; - } - } - if (change) { - return true; + // insert additional blocks for possible 'continue' insertion + BlockNode loopEnd = loop.getEnd(); + if (loopEnd.getPredecessors().size() > 1) { + boolean change = false; + List nodes = new ArrayList<>(loopEnd.getPredecessors()); + for (BlockNode pred : nodes) { + if (!pred.contains(AFlag.SYNTHETIC)) { + BlockSplitter.insertBlockBetween(mth, pred, loopEnd); + change = true; } } - // insert additional blocks for possible 'continue' insertion - BlockNode loopEnd = loop.getEnd(); - if (loopEnd.getPredecessors().size() > 1) { - boolean change = false; - List nodes = new ArrayList<>(loopEnd.getPredecessors()); - for (BlockNode pred : nodes) { - if (!pred.contains(AFlag.SYNTHETIC)) { - BlockSplitter.insertBlockBetween(mth, pred, loopEnd); - change = true; - } - } - if (change) { - return true; - } + if (change) { + return true; } } } - return splitReturn(mth); + return false; } /** @@ -334,7 +337,7 @@ public class BlockProcessor extends AbstractVisitor { } boolean first = true; for (BlockNode pred : preds) { - BlockNode newRetBlock = BlockSplitter.startNewBlock(mth, exitBlock.getStartOffset()); + BlockNode newRetBlock = BlockSplitter.startNewBlock(mth, -1); newRetBlock.add(AFlag.SYNTHETIC); InsnNode newRetInsn; if (first) { @@ -345,8 +348,7 @@ public class BlockProcessor extends AbstractVisitor { newRetInsn = duplicateReturnInsn(returnInsn); } newRetBlock.getInstructions().add(newRetInsn); - removeConnection(pred, exitBlock); - connect(pred, newRetBlock); + BlockSplitter.replaceConnection(pred, exitBlock, newRetBlock); } cleanExitNodes(mth); return true; @@ -389,35 +391,32 @@ public class BlockProcessor extends AbstractVisitor { return insn; } - private static void removeBlocks(MethodNode mth) { - Iterator it = mth.getBasicBlocks().iterator(); - while (it.hasNext()) { - BlockNode block = it.next(); + private static void removeMarkedBlocks(MethodNode mth) { + mth.getBasicBlocks().removeIf(block -> { if (block.contains(AFlag.REMOVE)) { - if (!block.getPredecessors().isEmpty() - || !block.getSuccessors().isEmpty()) { - LOG.error("Block {} not deleted, method: {}", block, mth); + if (!block.getPredecessors().isEmpty() || !block.getSuccessors().isEmpty()) { + LOG.warn("Block {} not deleted, method: {}", block, mth); } else { CatchAttr catchAttr = block.get(AType.CATCH_BLOCK); if (catchAttr != null) { catchAttr.getTryBlock().removeBlock(mth, block); } - it.remove(); + return true; } } - } + return false; + }); } private static void clearBlocksState(MethodNode mth) { - for (BlockNode block : mth.getBasicBlocks()) { + mth.getBasicBlocks().forEach(block -> { block.remove(AType.LOOP); block.remove(AFlag.LOOP_START); block.remove(AFlag.LOOP_END); - block.setDoms(null); block.setIDom(null); block.setDomFrontier(null); block.getDominatesOn().clear(); - } + }); } } diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockSplitter.java b/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockSplitter.java index 36e532e2..a9b8d6de 100644 --- a/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockSplitter.java +++ b/jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockSplitter.java @@ -11,6 +11,7 @@ import jadx.core.dex.attributes.AType; import jadx.core.dex.attributes.nodes.JumpInfo; import jadx.core.dex.instructions.IfNode; import jadx.core.dex.instructions.InsnType; +import jadx.core.dex.instructions.TargetInsnNode; import jadx.core.dex.nodes.BlockNode; import jadx.core.dex.nodes.InsnNode; import jadx.core.dex.nodes.MethodNode; @@ -18,6 +19,7 @@ import jadx.core.dex.trycatch.CatchAttr; import jadx.core.dex.trycatch.ExceptionHandler; import jadx.core.dex.trycatch.SplitterBlockAttr; import jadx.core.dex.visitors.AbstractVisitor; +import jadx.core.utils.BlockUtils; import jadx.core.utils.exceptions.JadxRuntimeException; public class BlockSplitter extends AbstractVisitor { @@ -42,7 +44,20 @@ public class BlockSplitter extends AbstractVisitor { mth.initBasicBlocks(); splitBasicBlocks(mth); removeInsns(mth); - removeEmptyBlocks(mth); + removeEmptyDetachedBlocks(mth); + initBlocksInTargetNodes(mth); + } + + /** + * Init 'then' and 'else' blocks for 'if' instruction. + */ + private static void initBlocksInTargetNodes(MethodNode mth) { + mth.getBasicBlocks().forEach(block -> { + InsnNode lastInsn = BlockUtils.getLastInsn(block); + if (lastInsn instanceof TargetInsnNode) { + ((TargetInsnNode) lastInsn).initBlocks(block); + } + }); } private static void splitBasicBlocks(MethodNode mth) { @@ -133,15 +148,31 @@ public class BlockSplitter extends AbstractVisitor { to.getPredecessors().remove(from); } + static void replaceConnection(BlockNode source, BlockNode oldDest, BlockNode newDest) { + removeConnection(source, oldDest); + connect(source, newDest); + replaceTarget(source, oldDest, newDest); + } + static BlockNode insertBlockBetween(MethodNode mth, BlockNode source, BlockNode target) { BlockNode newBlock = startNewBlock(mth, target.getStartOffset()); newBlock.add(AFlag.SYNTHETIC); removeConnection(source, target); connect(source, newBlock); connect(newBlock, target); + replaceTarget(source, target, newBlock); + source.updateCleanSuccessors(); + newBlock.updateCleanSuccessors(); return newBlock; } + static void replaceTarget(BlockNode source, BlockNode oldTarget, BlockNode newTarget) { + InsnNode lastInsn = BlockUtils.getLastInsn(source); + if (lastInsn instanceof TargetInsnNode) { + ((TargetInsnNode) lastInsn).replaceTargetBlock(oldTarget, newTarget); + } + } + private static void setupConnections(MethodNode mth, Map blocksMap) { for (BlockNode block : mth.getBasicBlocks()) { for (InsnNode insn : block.getInstructions()) { @@ -222,7 +253,7 @@ public class BlockSplitter extends AbstractVisitor { } } - private void removeEmptyBlocks(MethodNode mth) { + static void removeEmptyDetachedBlocks(MethodNode mth) { mth.getBasicBlocks().removeIf(block -> block.getInstructions().isEmpty() && block.getPredecessors().isEmpty() 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 5c1effde..c7014045 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 @@ -2,8 +2,6 @@ package jadx.core.dex.visitors.regions; import java.util.ArrayList; import java.util.BitSet; -import java.util.Collections; -import java.util.Comparator; import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; @@ -46,13 +44,11 @@ import jadx.core.utils.ErrorsCounter; import jadx.core.utils.InstructionRemover; import jadx.core.utils.RegionUtils; import jadx.core.utils.exceptions.JadxOverflowException; -import jadx.core.utils.exceptions.JadxRuntimeException; import static jadx.core.dex.visitors.regions.IfMakerHelper.confirmMerge; import static jadx.core.dex.visitors.regions.IfMakerHelper.makeIfInfo; import static jadx.core.dex.visitors.regions.IfMakerHelper.mergeNestedIfNodes; import static jadx.core.dex.visitors.regions.IfMakerHelper.searchNestedIf; -import static jadx.core.utils.BlockUtils.getBlockByOffset; import static jadx.core.utils.BlockUtils.getNextBlock; import static jadx.core.utils.BlockUtils.isPathExists; import static jadx.core.utils.BlockUtils.skipSyntheticSuccessor; @@ -691,27 +687,14 @@ public class RegionMaker { int len = insn.getTargets().length; // sort by target - Map> casesMap = new LinkedHashMap<>(len); + Map> blocksMap = new LinkedHashMap<>(len); for (int i = 0; i < len; i++) { Object key = insn.getKeys()[i]; - int targ = insn.getTargets()[i]; - List keys = casesMap.get(targ); - if (keys == null) { - keys = new ArrayList<>(2); - casesMap.put(targ, keys); - } + BlockNode targ = insn.getTargetBlocks()[i]; + List keys = blocksMap.computeIfAbsent(targ, k -> new ArrayList<>(2)); keys.add(key); } - - Map> blocksMap = new LinkedHashMap<>(len); - for (Map.Entry> entry : casesMap.entrySet()) { - BlockNode c = getBlockByOffset(entry.getKey(), block.getSuccessors()); - if (c == null) { - throw new JadxRuntimeException("Switch block not found by offset: " + entry.getKey()); - } - blocksMap.put(c, entry.getValue()); - } - BlockNode defCase = getBlockByOffset(insn.getDefaultCaseOffset(), block.getSuccessors()); + BlockNode defCase = insn.getDefTargetBlock(); if (defCase != null) { blocksMap.remove(defCase); } @@ -860,19 +843,16 @@ public class RegionMaker { Map fallThroughCases) { List list = new ArrayList<>(blocksMap.size()); list.addAll(blocksMap.keySet()); - Collections.sort(list, new Comparator() { - @Override - public int compare(BlockNode a, BlockNode b) { - BlockNode nextA = fallThroughCases.get(a); - if (nextA != null) { - if (b.equals(nextA)) { - return -1; - } - } else if (a.equals(fallThroughCases.get(b))) { - return 1; + list.sort((a, b) -> { + BlockNode nextA = fallThroughCases.get(a); + if (nextA != null) { + if (b.equals(nextA)) { + return -1; } - return 0; + } else if (a.equals(fallThroughCases.get(b))) { + return 1; } + return 0; }); Map> newBlocksMap = new LinkedHashMap<>(blocksMap.size()); diff --git a/jadx-core/src/main/java/jadx/core/utils/BlockUtils.java b/jadx-core/src/main/java/jadx/core/utils/BlockUtils.java index cc62293a..037c537c 100644 --- a/jadx-core/src/main/java/jadx/core/utils/BlockUtils.java +++ b/jadx-core/src/main/java/jadx/core/utils/BlockUtils.java @@ -453,13 +453,12 @@ public class BlockUtils { */ public static List collectBlocksDominatedBy(BlockNode dominator, BlockNode start) { List result = new ArrayList<>(); - HashSet visited = new HashSet(); + Set visited = new HashSet<>(); collectWhileDominates(dominator, start, result, visited); return result; } - private static void collectWhileDominates(BlockNode dominator, BlockNode child, List result, - HashSet visited) { + private static void collectWhileDominates(BlockNode dominator, BlockNode child, List result, Set visited) { if (visited.contains(child)) { return; } -- GitLab