From 4cdad0e83e0369a82b1e9ed169a5a5d057bcf086 Mon Sep 17 00:00:00 2001
From: Skylot <skylot@gmail.com>
Date: Tue, 17 Mar 2020 19:38:45 +0000
Subject: [PATCH] fix: correct method exit blocks collection (#876)

---
 .../visitors/blocksmaker/BlockProcessor.java  | 33 +++++++++++---
 .../switches/TestSwitchWithThrow.java         | 44 +++++++++++++++++++
 2 files changed, 71 insertions(+), 6 deletions(-)
 create mode 100644 jadx-core/src/test/java/jadx/tests/integration/switches/TestSwitchWithThrow.java

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 4e703458e..64fca3ce4 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
@@ -58,7 +58,7 @@ public class BlockProcessor extends AbstractVisitor {
 			clearBlocksState(mth);
 			computeDominators(mth);
 		}
-		markReturnBlocks(mth);
+		updateExitBlocks(mth);
 
 		int i = 0;
 		while (modifyBlocksTree(mth)) {
@@ -66,7 +66,7 @@ public class BlockProcessor extends AbstractVisitor {
 			clearBlocksState(mth);
 			// recalculate dominators tree
 			computeDominators(mth);
-			markReturnBlocks(mth);
+			updateExitBlocks(mth);
 
 			if (i++ > 100) {
 				mth.addWarn("CFG modification limit reached, blocks count: " + mth.getBasicBlocks().size());
@@ -337,12 +337,33 @@ public class BlockProcessor extends AbstractVisitor {
 		block.setDomFrontier(domFrontier);
 	}
 
-	private static void markReturnBlocks(MethodNode mth) {
+	private static void updateExitBlocks(MethodNode mth) {
 		mth.getExitBlocks().clear();
 		mth.getBasicBlocks().forEach(block -> {
-			if (BlockUtils.checkLastInsnType(block, InsnType.RETURN)) {
-				block.add(AFlag.RETURN);
-				mth.getExitBlocks().add(block);
+			boolean noSuccessors = block.getSuccessors().isEmpty();
+			boolean exitBlock = false;
+			InsnNode lastInsn = BlockUtils.getLastInsn(block);
+			if (lastInsn != null) {
+				InsnType insnType = lastInsn.getType();
+				if (insnType == InsnType.RETURN) {
+					block.add(AFlag.RETURN);
+					exitBlock = true;
+					if (!noSuccessors) {
+						throw new JadxRuntimeException("Found a block after RETURN instruction: " + lastInsn + " in block: " + block);
+					}
+				} else if (insnType == InsnType.THROW) {
+					if (noSuccessors) {
+						exitBlock = true;
+					}
+				}
+			}
+			if (exitBlock) {
+				mth.addExitBlock(block);
+			} else if (noSuccessors
+					&& !mth.isVoidReturn()
+					&& !mth.isConstructor()) {
+				mth.addComment("JADX INFO: Unexpected exit block: " + block
+						+ ". Expect last instruction to be RETURN or THROW, got: " + lastInsn);
 			}
 		});
 	}
diff --git a/jadx-core/src/test/java/jadx/tests/integration/switches/TestSwitchWithThrow.java b/jadx-core/src/test/java/jadx/tests/integration/switches/TestSwitchWithThrow.java
new file mode 100644
index 000000000..d2b65db2a
--- /dev/null
+++ b/jadx-core/src/test/java/jadx/tests/integration/switches/TestSwitchWithThrow.java
@@ -0,0 +1,44 @@
+package jadx.tests.integration.switches;
+
+import org.junit.jupiter.api.Test;
+
+import jadx.tests.api.IntegrationTest;
+
+import static jadx.tests.api.utils.assertj.JadxAssertions.assertThat;
+
+public class TestSwitchWithThrow extends IntegrationTest {
+
+	public static class TestCls {
+		public int test(int i) {
+			if (i != 0) {
+				switch (i % 4) {
+					case 1:
+						throw new IllegalStateException("1");
+					case 2:
+						throw new IllegalStateException("2");
+					default:
+						throw new IllegalStateException("Other");
+				}
+			} else {
+				System.out.println("0");
+				return -1;
+			}
+		}
+
+		public void check() {
+			assertThat(test(0)).isEqualTo(-1);
+			// TODO: implement 'invoke-custom' support
+			// assertThat(catchThrowable(() -> test(1)))
+			// .isInstanceOf(IllegalStateException.class).hasMessageContaining("1");
+			// assertThat(catchThrowable(() -> test(3)))
+			// .isInstanceOf(IllegalStateException.class).hasMessageContaining("Other");
+		}
+	}
+
+	@Test
+	public void test() {
+		assertThat(getClassNode(TestCls.class))
+				.code()
+				.contains("throw new IllegalStateException(\"1\");");
+	}
+}