fix: avoid self-loop for exception handlers (#2147)

skylot · Apr 11, 2024 · 6182332 · 6182332
1 parent 37b5709
commit 6182332
Show file tree

Hide file tree

Showing 5 changed files with 78 additions and 0 deletions.
diff --git a/jadx-core/src/main/java/jadx/core/dex/attributes/AFlag.java b/jadx-core/src/main/java/jadx/core/dex/attributes/AFlag.java
@@ -106,4 +106,5 @@ public enum AFlag {
 	DONT_UNLOAD_CLASS, // don't unload class after code generation (only for tests and debug!)
 
 	RESOLVE_JAVA_JSR,
+	COMPUTE_POST_DOM,
 }
diff --git a/jadx-core/src/main/java/jadx/core/dex/instructions/InsnDecoder.java b/jadx-core/src/main/java/jadx/core/dex/instructions/InsnDecoder.java
@@ -521,6 +521,7 @@ private SwitchInsn makeSwitch(InsnData insn, boolean packed) {
 		if (payload != null) {
 			swInsn.attachSwitchData(new SwitchData((ISwitchPayload) payload), insn.getTarget());
 		}
+		method.add(AFlag.COMPUTE_POST_DOM);
 		return swInsn;
 	}
 

diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/blocks/BlockExceptionHandler.java b/jadx-core/src/main/java/jadx/core/dex/visitors/blocks/BlockExceptionHandler.java
@@ -331,6 +331,17 @@ private static boolean wrapBlocksWithTryCatch(MethodNode mth, TryCatchBlockAttr
 			return false;
 		}
 		BlockNode bottom = searchBottomBlock(mth, blocks);
+		BlockNode splitReturn;
+		if (bottom != null && bottom.isReturnBlock()) {
+			if (Consts.DEBUG_EXC_HANDLERS) {
+				LOG.debug("TryCatch #{} bottom block ({}) is return, split", tryCatchBlock.id(), bottom);
+			}
+			splitReturn = bottom;
+			bottom = BlockSplitter.blockSplitTop(mth, bottom);
+			bottom.add(AFlag.SYNTHETIC);
+		} else {
+			splitReturn = null;
+		}
 		if (Consts.DEBUG_EXC_HANDLERS) {
 			LOG.debug("TryCatch #{} split: top {}, bottom: {}", tryCatchBlock.id(), top, bottom);
 		}
@@ -349,6 +360,18 @@ private static boolean wrapBlocksWithTryCatch(MethodNode mth, TryCatchBlockAttr
 			bottomSplitterBlock.add(AFlag.EXC_BOTTOM_SPLITTER);
 			bottomSplitterBlock.add(AFlag.SYNTHETIC);
 			BlockSplitter.connect(bottom, bottomSplitterBlock);
+			if (splitReturn != null) {
+				// redirect handler to return block instead synthetic split block to avoid self-loop
+				BlockSet bottomPreds = BlockSet.from(mth, bottom.getPredecessors());
+				for (ExceptionHandler handler : tryCatchBlock.getHandlers()) {
+					if (bottomPreds.intersects(handler.getBlocks())) {
+						BlockNode lastBlock = bottomPreds.intersect(handler.getBlocks()).getOne();
+						if (lastBlock != null) {
+							BlockSplitter.replaceConnection(lastBlock, bottom, splitReturn);
+						}
+					}
+				}
+			}
 		}
 
 		if (Consts.DEBUG_EXC_HANDLERS) {

diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/blocks/PostDominatorTree.java b/jadx-core/src/main/java/jadx/core/dex/visitors/blocks/PostDominatorTree.java
@@ -4,6 +4,7 @@
 import java.util.BitSet;
 import java.util.List;
 
+import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.nodes.BlockNode;
 import jadx.core.dex.nodes.MethodNode;
 import jadx.core.utils.BlockUtils;
@@ -12,6 +13,9 @@
 public class PostDominatorTree {
 
 	public static void compute(MethodNode mth) {
+		if (!mth.contains(AFlag.COMPUTE_POST_DOM)) {
+			return;
+		}
 		try {
 			int mthBlocksCount = mth.getBasicBlocks().size();
 			List<BlockNode> sorted = new ArrayList<>(mthBlocksCount);
@@ -57,6 +61,9 @@ public static void compute(MethodNode mth) {
 				}
 				mth.addInfoComment("Infinite loop detected, blocks: " + blocksDelta + ", insns: " + insnsCount);
 			}
+		} catch (Throwable e) {
+			// show error as a warning because this info not always used
+			mth.addWarnComment("Failed to build post-dominance tree", e);
 		} finally {
 			// revert block ids change
 			mth.updateBlockIds(mth.getBasicBlocks());

diff --git a/jadx-core/src/main/java/jadx/core/utils/blocks/BlockSet.java b/jadx-core/src/main/java/jadx/core/utils/blocks/BlockSet.java
@@ -2,16 +2,25 @@
 
 import java.util.ArrayList;
 import java.util.BitSet;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.function.Consumer;
 
+import org.jetbrains.annotations.Nullable;
+
 import jadx.core.dex.nodes.BlockNode;
 import jadx.core.dex.nodes.MethodNode;
 import jadx.core.utils.EmptyBitSet;
 
 public class BlockSet {
 
+	public static BlockSet from(MethodNode mth, Collection<BlockNode> blocks) {
+		BlockSet newBS = new BlockSet(mth);
+		newBS.set(blocks);
+		return newBS;
+	}
+
 	private final MethodNode mth;
 	private final BitSet bs;
 
@@ -28,13 +37,50 @@ public void set(BlockNode block) {
 		bs.set(block.getId());
 	}
 
+	public void set(Collection<BlockNode> blocks) {
+		blocks.forEach(this::set);
+	}
+
 	public boolean checkAndSet(BlockNode block) {
 		int id = block.getId();
 		boolean state = bs.get(id);
 		bs.set(id);
 		return state;
 	}
 
+	public boolean intersects(List<BlockNode> blocks) {
+		for (BlockNode block : blocks) {
+			if (get(block)) {
+				return true;
+			}
+		}
+		return false;
+	}
+
+	public BlockSet intersect(List<BlockNode> blocks) {
+		BlockSet input = from(mth, blocks);
+		BlockSet result = new BlockSet(mth);
+		BitSet resultBS = result.bs;
+		resultBS.or(this.bs);
+		resultBS.and(input.bs);
+		return result;
+	}
+
+	public boolean isEmpty() {
+		return bs.cardinality() == 0;
+	}
+
+	public int size() {
+		return bs.cardinality();
+	}
+
+	public @Nullable BlockNode getOne() {
+		if (bs.cardinality() == 1) {
+			return mth.getBasicBlocks().get(bs.nextSetBit(0));
+		}
+		return null;
+	}
+
 	public void forEach(Consumer<? super BlockNode> consumer) {
 		if (bs.isEmpty()) {
 			return;