From 681f8a98b5de8a596d336793129d2cf26089c3e0 Mon Sep 17 00:00:00 2001 From: Skylot <118523+skylot@users.noreply.github.com> Date: Sat, 28 Sep 2024 16:23:16 +0100 Subject: [PATCH] fix: improve checks for restore new filled array (#2289) --- .../main/java/jadx/core/codegen/InsnGen.java | 2 +- .../attributes/nodes/CodeFeaturesAttr.java | 5 ++ .../core/dex/instructions/InsnDecoder.java | 1 + .../core/dex/instructions/args/InsnArg.java | 5 ++ .../core/dex/visitors/ReplaceNewArray.java | 52 ++++++++++++++----- .../java/jadx/core/utils/InsnRemover.java | 11 ++++ .../main/java/jadx/core/utils/InsnUtils.java | 30 +++++++++++ .../arrays/TestArrayFillNegative.java | 33 ++++++++++++ .../arrays/TestMultiDimArrayFill.java | 8 +-- 9 files changed, 125 insertions(+), 22 deletions(-) create mode 100644 jadx-core/src/test/java/jadx/tests/integration/arrays/TestArrayFillNegative.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 3daa841c0..c077fa835 100644 --- a/jadx-core/src/main/java/jadx/core/codegen/InsnGen.java +++ b/jadx-core/src/main/java/jadx/core/codegen/InsnGen.java @@ -434,7 +434,7 @@ public class InsnGen { code.add(']'); } int dim = arrayType.getArrayDimension(); - for (; k < dim - 1; k++) { + for (; k < dim; k++) { code.add("[]"); } break; diff --git a/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/CodeFeaturesAttr.java b/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/CodeFeaturesAttr.java index b0fc2b981..cc25bd087 100644 --- a/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/CodeFeaturesAttr.java +++ b/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/CodeFeaturesAttr.java @@ -14,6 +14,11 @@ public class CodeFeaturesAttr implements IJadxAttribute { * Code contains switch instruction */ SWITCH, + + /** + * Code contains new-array instruction + */ + NEW_ARRAY, } public static boolean contains(MethodNode mth, CodeFeature feature) { 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 index 76956585f..248b12a51 100644 --- a/jadx-core/src/main/java/jadx/core/dex/instructions/InsnDecoder.java +++ b/jadx-core/src/main/java/jadx/core/dex/instructions/InsnDecoder.java @@ -540,6 +540,7 @@ public class InsnDecoder { for (int i = 1; i < regsCount; i++) { newArr.addArg(InsnArg.typeImmutableReg(insn, i, ArgType.INT)); } + CodeFeaturesAttr.add(method, CodeFeature.NEW_ARRAY); return newArr; } diff --git a/jadx-core/src/main/java/jadx/core/dex/instructions/args/InsnArg.java b/jadx-core/src/main/java/jadx/core/dex/instructions/args/InsnArg.java index 1ce323d78..4e38809b7 100644 --- a/jadx-core/src/main/java/jadx/core/dex/instructions/args/InsnArg.java +++ b/jadx-core/src/main/java/jadx/core/dex/instructions/args/InsnArg.java @@ -13,6 +13,7 @@ import jadx.core.dex.instructions.InsnType; import jadx.core.dex.nodes.InsnNode; import jadx.core.dex.nodes.MethodNode; import jadx.core.utils.InsnRemover; +import jadx.core.utils.InsnUtils; import jadx.core.utils.exceptions.JadxRuntimeException; /** @@ -301,6 +302,10 @@ public abstract class InsnArg extends Typed { return false; } + public boolean isUseVar(RegisterArg arg) { + return InsnUtils.containsVar(this, arg); + } + protected final T copyCommonParams(T copy) { copy.copyAttributesFrom(this); copy.setParentInsn(parentInsn); diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/ReplaceNewArray.java b/jadx-core/src/main/java/jadx/core/dex/visitors/ReplaceNewArray.java index dc99ceaaa..80ebc34ab 100644 --- a/jadx-core/src/main/java/jadx/core/dex/visitors/ReplaceNewArray.java +++ b/jadx-core/src/main/java/jadx/core/dex/visitors/ReplaceNewArray.java @@ -1,12 +1,17 @@ package jadx.core.dex.visitors; -import java.util.HashSet; +import java.util.ArrayList; +import java.util.Collections; +import java.util.IdentityHashMap; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.SortedMap; import java.util.TreeMap; import jadx.core.dex.attributes.AFlag; +import jadx.core.dex.attributes.nodes.CodeFeaturesAttr; +import jadx.core.dex.attributes.nodes.CodeFeaturesAttr.CodeFeature; import jadx.core.dex.instructions.FilledNewArrayNode; import jadx.core.dex.instructions.IndexInsnNode; import jadx.core.dex.instructions.InsnType; @@ -35,21 +40,20 @@ public class ReplaceNewArray extends AbstractVisitor { @Override public void visit(MethodNode mth) throws JadxException { - if (mth.isNoCode()) { + if (!CodeFeaturesAttr.contains(mth, CodeFeature.NEW_ARRAY)) { return; } + InsnRemover remover = new InsnRemover(mth); int k = 0; while (true) { boolean changed = false; - InsnRemover remover = new InsnRemover(mth); for (BlockNode block : mth.getBasicBlocks()) { - remover.setBlock(block); - List instructions = block.getInstructions(); - int size = instructions.size(); + List insnList = block.getInstructions(); + int size = insnList.size(); for (int i = 0; i < size; i++) { - changed |= processInsn(mth, instructions, i, remover); + changed |= processInsn(mth, insnList, i, remover); } - remover.perform(); + remover.performForBlock(block); } if (changed) { CodeShrinkVisitor.shrinkMethod(mth); @@ -107,12 +111,11 @@ public class ReplaceNewArray extends AbstractVisitor { SortedMap arrPuts = new TreeMap<>(); for (RegisterArg registerArg : useList) { InsnNode parentInsn = registerArg.getParentInsn(); - if (parentInsn == null || parentInsn.getType() != InsnType.APUT) { + if (parentInsn == null + || parentInsn.getType() != InsnType.APUT + || !arrArg.sameRegAndSVar(parentInsn.getArg(0))) { continue; } - if (!arrArg.sameRegAndSVar(parentInsn.getArg(0))) { - return false; - } Object constVal = InsnUtils.getConstValueByArg(mth.root(), parentInsn.getArg(1)); if (!(constVal instanceof LiteralArg)) { return false; @@ -130,8 +133,7 @@ public class ReplaceNewArray extends AbstractVisitor { if (arrPuts.size() < minLen) { return false; } - // expect all puts to be in same block - if (!new HashSet<>(instructions).containsAll(arrPuts.values())) { + if (!verifyPutInsns(arrArg, instructions, arrPuts)) { return false; } @@ -165,6 +167,28 @@ public class ReplaceNewArray extends AbstractVisitor { return true; } + private static boolean verifyPutInsns(RegisterArg arrReg, List insnList, SortedMap arrPuts) { + List puts = new ArrayList<>(arrPuts.values()); + int putsCount = puts.size(); + // expect all puts to be in the same block + if (insnList.size() < putsCount) { + return false; + } + Set insnSet = Collections.newSetFromMap(new IdentityHashMap<>()); + insnSet.addAll(insnList); + if (!insnSet.containsAll(puts)) { + return false; + } + // array arg shouldn't be used in puts insns + for (InsnNode put : puts) { + InsnArg putArg = put.getArg(2); + if (putArg.isUseVar(arrReg)) { + return false; + } + } + return true; + } + private static InsnArg replaceConstInArg(MethodNode mth, InsnArg valueArg) { if (valueArg.isLiteral()) { IFieldInfoRef f = mth.getParentClass().getConstFieldByLiteralArg((LiteralArg) valueArg); diff --git a/jadx-core/src/main/java/jadx/core/utils/InsnRemover.java b/jadx-core/src/main/java/jadx/core/utils/InsnRemover.java index 1e2b1003e..eb737e25f 100644 --- a/jadx-core/src/main/java/jadx/core/utils/InsnRemover.java +++ b/jadx-core/src/main/java/jadx/core/utils/InsnRemover.java @@ -3,6 +3,7 @@ package jadx.core.utils; import java.util.ArrayList; import java.util.Iterator; import java.util.List; +import java.util.Objects; import java.util.stream.Collectors; import org.jetbrains.annotations.Nullable; @@ -75,6 +76,16 @@ public class InsnRemover { toRemove.clear(); } + public void performForBlock(BlockNode block) { + if (toRemove.isEmpty()) { + return; + } + instrList = Objects.requireNonNull(block.getInstructions()); + unbindInsns(mth, toRemove); + removeAll(instrList, toRemove); + toRemove.clear(); + } + public static void unbindInsn(@Nullable MethodNode mth, InsnNode insn) { unbindAllArgs(mth, insn); unbindResult(mth, insn); diff --git a/jadx-core/src/main/java/jadx/core/utils/InsnUtils.java b/jadx-core/src/main/java/jadx/core/utils/InsnUtils.java index 0de6b4ced..e1d73c7f7 100644 --- a/jadx-core/src/main/java/jadx/core/utils/InsnUtils.java +++ b/jadx-core/src/main/java/jadx/core/utils/InsnUtils.java @@ -263,6 +263,36 @@ public class InsnUtils { return false; } + public static boolean containsVar(InsnNode insn, RegisterArg arg) { + if (insn == null) { + return false; + } + RegisterArg result = insn.getResult(); + if (result != null && result.sameRegAndSVar(arg)) { + return true; + } + if (insn.getArgsCount() == 0) { + return false; + } + for (InsnArg insnArg : insn.getArguments()) { + if (containsVar(insnArg, arg)) { + return true; + } + } + return false; + } + + public static boolean containsVar(InsnArg insnArg, RegisterArg arg) { + if (insnArg.isRegister()) { + return ((RegisterArg) insnArg).sameRegAndSVar(arg); + } + if (insnArg.isInsnWrap()) { + InsnNode wrapInsn = ((InsnWrapArg) insnArg).getWrapInsn(); + return containsVar(wrapInsn, arg); + } + return false; + } + public static boolean contains(InsnNode insn, AFlag flag) { return insn != null && insn.contains(flag); } diff --git a/jadx-core/src/test/java/jadx/tests/integration/arrays/TestArrayFillNegative.java b/jadx-core/src/test/java/jadx/tests/integration/arrays/TestArrayFillNegative.java new file mode 100644 index 000000000..715761a60 --- /dev/null +++ b/jadx-core/src/test/java/jadx/tests/integration/arrays/TestArrayFillNegative.java @@ -0,0 +1,33 @@ +package jadx.tests.integration.arrays; + +import org.junit.jupiter.api.Test; + +import jadx.tests.api.IntegrationTest; + +import static jadx.tests.api.utils.assertj.JadxAssertions.assertThat; + +public class TestArrayFillNegative extends IntegrationTest { + + public static class TestCls { + public int[] test() { + int[] arr = new int[3]; + arr[0] = 1; + arr[1] = arr[0] + 1; + arr[2] = arr[1] + 1; + return arr; + } + + public void check() { + assertThat(test()).isEqualTo(new int[] { 1, 2, 3 }); + } + } + + @Test + public void test() { + assertThat(getClassNode(TestCls.class)) + .code() + .doesNotContain("int[] arr = {1, ") + .containsOne("int[] arr = new int[3];") + .containsOne("arr[1] = arr[0] + 1;"); + } +} diff --git a/jadx-core/src/test/java/jadx/tests/integration/arrays/TestMultiDimArrayFill.java b/jadx-core/src/test/java/jadx/tests/integration/arrays/TestMultiDimArrayFill.java index 877f29b86..de388262a 100644 --- a/jadx-core/src/test/java/jadx/tests/integration/arrays/TestMultiDimArrayFill.java +++ b/jadx-core/src/test/java/jadx/tests/integration/arrays/TestMultiDimArrayFill.java @@ -11,13 +11,7 @@ public class TestMultiDimArrayFill extends IntegrationTest { public static Obj test(int a, int b) { return new Obj( - new int[][] { - new int[] { 1 }, - new int[] { 2 }, - { 3 }, - new int[] { 4, 5 }, - new int[0] - }, + new int[][] { { 1 }, { 2 }, { 3 }, { 4, 5 }, new int[0] }, new int[] { a, a, a, a, b }); }