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 38ab41961..45ef65907 100644 --- a/jadx-core/src/main/java/jadx/core/codegen/RegionGen.java +++ b/jadx-core/src/main/java/jadx/core/codegen/RegionGen.java @@ -10,6 +10,7 @@ import jadx.core.dex.instructions.SwitchNode; import jadx.core.dex.instructions.args.InsnArg; import jadx.core.dex.instructions.args.RegisterArg; import jadx.core.dex.nodes.BlockNode; +import jadx.core.dex.nodes.FieldNode; import jadx.core.dex.nodes.IBlock; import jadx.core.dex.nodes.IContainer; import jadx.core.dex.nodes.IRegion; @@ -239,7 +240,14 @@ public class RegionGen extends InsnGen { IContainer c = sw.getCases().get(i); for (Object k : keys) { code.startLine("case "); - if (k instanceof IndexInsnNode) { + if (k instanceof FieldNode) { + FieldNode fn = (FieldNode) k; + if (fn.getParentClass().isEnum()) { + code.add(fn.getName()); + } else { + staticField(code, fn.getFieldInfo()); + } + } else if (k instanceof IndexInsnNode) { staticField(code, (FieldInfo) ((IndexInsnNode) k).getIndex()); } else { code.add(TypeGen.literalToString((Integer) k, arg.getType())); 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 248956465..6a3e6231e 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 @@ -4,6 +4,7 @@ import jadx.core.dex.attributes.annotations.AnnotationsList; import jadx.core.dex.attributes.annotations.MethodParameters; import jadx.core.dex.attributes.nodes.DeclareVariablesAttr; import jadx.core.dex.attributes.nodes.EnumClassAttr; +import jadx.core.dex.attributes.nodes.EnumMapAttr; import jadx.core.dex.attributes.nodes.FieldReplaceAttr; import jadx.core.dex.attributes.nodes.ForceReturnAttr; import jadx.core.dex.attributes.nodes.JadxErrorAttr; @@ -40,6 +41,7 @@ public class AType { public static final AType JADX_ERROR = new AType(); public static final AType METHOD_INLINE = new AType(); public static final AType ENUM_CLASS = new AType(); + public static final AType ENUM_MAP = new AType(); public static final AType ANNOTATION_LIST = new AType(); public static final AType ANNOTATION_MTH_PARAMETERS = new AType(); public static final AType PHI_LIST = new AType(); diff --git a/jadx-core/src/main/java/jadx/core/dex/attributes/AttrNode.java b/jadx-core/src/main/java/jadx/core/dex/attributes/AttrNode.java index 28563feae..ff298a827 100644 --- a/jadx-core/src/main/java/jadx/core/dex/attributes/AttrNode.java +++ b/jadx-core/src/main/java/jadx/core/dex/attributes/AttrNode.java @@ -12,25 +12,25 @@ public abstract class AttrNode implements IAttributeNode { @Override public void add(AFlag flag) { - getStorage().add(flag); + initStorage().add(flag); } @Override public void addAttr(IAttribute attr) { - getStorage().add(attr); + initStorage().add(attr); } @Override public void addAttr(AType> type, T obj) { - getStorage().add(type, obj); + initStorage().add(type, obj); } @Override public void copyAttributesFrom(AttrNode attrNode) { - getStorage().addAll(attrNode.storage); + initStorage().addAll(attrNode.storage); } - AttributeStorage getStorage() { + AttributeStorage initStorage() { AttributeStorage store = storage; if (store == EMPTY_ATTR_STORAGE) { store = new AttributeStorage(); diff --git a/jadx-core/src/main/java/jadx/core/dex/attributes/EmptyAttrStorage.java b/jadx-core/src/main/java/jadx/core/dex/attributes/EmptyAttrStorage.java index ac0179977..364db9b78 100644 --- a/jadx-core/src/main/java/jadx/core/dex/attributes/EmptyAttrStorage.java +++ b/jadx-core/src/main/java/jadx/core/dex/attributes/EmptyAttrStorage.java @@ -5,7 +5,7 @@ import jadx.core.dex.attributes.annotations.Annotation; import java.util.Collections; import java.util.List; -public class EmptyAttrStorage extends AttributeStorage { +public final class EmptyAttrStorage extends AttributeStorage { @Override public boolean contains(AFlag flag) { @@ -52,4 +52,9 @@ public class EmptyAttrStorage extends AttributeStorage { public List getAttributeStrings() { return Collections.emptyList(); } + + @Override + public String toString() { + return ""; + } } diff --git a/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/EnumMapAttr.java b/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/EnumMapAttr.java new file mode 100644 index 000000000..c45edb6cd --- /dev/null +++ b/jadx-core/src/main/java/jadx/core/dex/attributes/nodes/EnumMapAttr.java @@ -0,0 +1,27 @@ +package jadx.core.dex.attributes.nodes; + +import jadx.core.dex.attributes.AType; +import jadx.core.dex.attributes.IAttribute; + +import java.util.HashMap; +import java.util.Map; + +public class EnumMapAttr implements IAttribute { + + private Map map = new HashMap(); + + public Map getMap() { + return map; + } + + @Override + public AType getType() { + return AType.ENUM_MAP; + } + + @Override + public String toString() { + return "Enum fields map: " + map; + } + +} diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/EnumVisitor.java b/jadx-core/src/main/java/jadx/core/dex/visitors/EnumVisitor.java index f61b8451a..1d8add9ff 100644 --- a/jadx-core/src/main/java/jadx/core/dex/visitors/EnumVisitor.java +++ b/jadx-core/src/main/java/jadx/core/dex/visitors/EnumVisitor.java @@ -38,13 +38,12 @@ public class EnumVisitor extends AbstractVisitor { // collect enum fields, remove synthetic List enumFields = new ArrayList(); - for (Iterator it = cls.getFields().iterator(); it.hasNext(); ) { - FieldNode f = it.next(); + for (FieldNode f : cls.getFields()) { if (f.getAccessFlags().isEnum()) { enumFields.add(f); - it.remove(); + f.add(AFlag.DONT_GENERATE); } else if (f.getAccessFlags().isSynthetic()) { - it.remove(); + f.add(AFlag.DONT_GENERATE); } } diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/ReSugarCode.java b/jadx-core/src/main/java/jadx/core/dex/visitors/ReSugarCode.java index c7d17af6a..7928e051f 100644 --- a/jadx-core/src/main/java/jadx/core/dex/visitors/ReSugarCode.java +++ b/jadx-core/src/main/java/jadx/core/dex/visitors/ReSugarCode.java @@ -1,12 +1,23 @@ package jadx.core.dex.visitors; +import jadx.core.dex.attributes.AFlag; +import jadx.core.dex.attributes.AType; +import jadx.core.dex.attributes.nodes.EnumMapAttr; +import jadx.core.dex.info.FieldInfo; +import jadx.core.dex.instructions.IndexInsnNode; import jadx.core.dex.instructions.InsnType; +import jadx.core.dex.instructions.InvokeNode; +import jadx.core.dex.instructions.SwitchNode; import jadx.core.dex.instructions.args.InsnArg; +import jadx.core.dex.instructions.args.InsnWrapArg; import jadx.core.dex.instructions.args.LiteralArg; import jadx.core.dex.nodes.BlockNode; +import jadx.core.dex.nodes.ClassNode; +import jadx.core.dex.nodes.FieldNode; import jadx.core.dex.nodes.InsnNode; import jadx.core.dex.nodes.MethodNode; import jadx.core.utils.InstructionRemover; +import jadx.core.utils.exceptions.DecodeException; import jadx.core.utils.exceptions.JadxException; import java.util.List; @@ -43,6 +54,9 @@ public class ReSugarCode extends AbstractVisitor { case NEW_ARRAY: return processNewArray(mth, instructions, i, remover); + case SWITCH: + return processEnumSwitch(mth, (SwitchNode) insn); + default: return null; } @@ -75,4 +89,153 @@ public class ReSugarCode extends AbstractVisitor { } return filledArr; } + + private static InsnNode processEnumSwitch(MethodNode mth, SwitchNode insn) { + InsnArg arg = insn.getArg(0); + if (!arg.isInsnWrap()) { + return null; + } + InsnNode wrapInsn = ((InsnWrapArg) arg).getWrapInsn(); + if (wrapInsn.getType() != InsnType.AGET) { + return null; + } + EnumMapInfo enumMapInfo = checkEnumMapAccess(mth, wrapInsn); + if (enumMapInfo == null) { + return null; + } + FieldNode enumMapField = enumMapInfo.getMapField(); + InsnArg invArg = enumMapInfo.getArg(); + + EnumMapAttr enumMapAttr = getEnumMap(mth, enumMapField); + if (enumMapAttr == null) { + return null; + } + Object[] keys = insn.getKeys(); + for (int i = 0; i < keys.length; i++) { + Object key = keys[i]; + Object newKey = enumMapAttr.getMap().get(key); + if (newKey != null) { + keys[i] = newKey; + } else { + return null; + } + } + enumMapField.getParentClass().add(AFlag.DONT_GENERATE); + insn.replaceArg(arg, invArg); + return null; + } + + private static EnumMapAttr getEnumMap(MethodNode mth, FieldNode field) { + ClassNode syntheticClass = field.getParentClass(); + EnumMapAttr mapAttr = syntheticClass.get(AType.ENUM_MAP); + if (mapAttr != null) { + return mapAttr; + } + MethodNode clsInitMth = syntheticClass.searchMethodByName("()V"); + if (clsInitMth == null || clsInitMth.isNoCode()) { + return null; + } + if (clsInitMth.getBasicBlocks() == null) { + try { + clsInitMth.load(); + } catch (DecodeException e) { + LOG.error("Load failed", e); + return null; + } + if (clsInitMth.getBasicBlocks() == null) { + // TODO: + return null; + } + } + mapAttr = new EnumMapAttr(); + for (BlockNode block : clsInitMth.getBasicBlocks()) { + for (InsnNode insn : block.getInstructions()) { + if (insn.getType() == InsnType.APUT) { + addToEnumMap(mth, field, mapAttr, insn); + } + } + } + if (mapAttr.getMap().isEmpty()) { + return null; + } + syntheticClass.addAttr(mapAttr); + return mapAttr; + } + + private static void addToEnumMap(MethodNode mth, FieldNode field, EnumMapAttr mapAttr, InsnNode insn) { + InsnArg litArg = insn.getArg(2); + if (!litArg.isLiteral()) { + return; + } + EnumMapInfo mapInfo = checkEnumMapAccess(mth, insn); + if (mapInfo == null) { + return; + } + InsnArg enumArg = mapInfo.getArg(); + if (mapInfo.getMapField() != field || !enumArg.isInsnWrap()) { + return; + } + InsnNode sget = ((InsnWrapArg) enumArg).getWrapInsn(); + if (!(sget instanceof IndexInsnNode)) { + return; + } + Object index = ((IndexInsnNode) sget).getIndex(); + if (!(index instanceof FieldInfo)) { + return; + } + FieldNode fieldNode = mth.dex().resolveField((FieldInfo) index); + if (fieldNode == null) { + return; + } + int literal = (int) ((LiteralArg) litArg).getLiteral(); + mapAttr.getMap().put(literal, fieldNode); + } + + public static EnumMapInfo checkEnumMapAccess(MethodNode mth, InsnNode checkInsn) { + InsnArg sgetArg = checkInsn.getArg(0); + InsnArg invArg = checkInsn.getArg(1); + if (!sgetArg.isInsnWrap() || !invArg.isInsnWrap()) { + return null; + } + InsnNode invInsn = ((InsnWrapArg) invArg).getWrapInsn(); + InsnNode sgetInsn = ((InsnWrapArg) sgetArg).getWrapInsn(); + if (invInsn.getType() != InsnType.INVOKE || sgetInsn.getType() != InsnType.SGET) { + return null; + } + InvokeNode inv = (InvokeNode) invInsn; + if (!inv.getCallMth().getShortId().equals("ordinal()I")) { + return null; + } + ClassNode enumCls = mth.dex().resolveClass(inv.getCallMth().getDeclClass()); + if (enumCls == null || !enumCls.isEnum()) { + return null; + } + Object index = ((IndexInsnNode) sgetInsn).getIndex(); + if (!(index instanceof FieldInfo)) { + return null; + } + FieldNode enumMapField = mth.dex().resolveField(((FieldInfo) index)); + if (enumMapField == null || !enumMapField.getAccessFlags().isSynthetic()) { + return null; + } + return new EnumMapInfo(inv.getArg(0), enumMapField); + } + + private static class EnumMapInfo { + private final InsnArg arg; + private final FieldNode mapField; + + public EnumMapInfo(InsnArg arg, FieldNode mapField) { + this.arg = arg; + this.mapField = mapField; + } + + public InsnArg getArg() { + return arg; + } + + public FieldNode getMapField() { + return mapField; + } + } } diff --git a/jadx-core/src/test/java/jadx/tests/internal/enums/TestSwitchOverEnum.java b/jadx-core/src/test/java/jadx/tests/internal/enums/TestSwitchOverEnum.java new file mode 100644 index 000000000..7cc1db24c --- /dev/null +++ b/jadx-core/src/test/java/jadx/tests/internal/enums/TestSwitchOverEnum.java @@ -0,0 +1,39 @@ +package jadx.tests.internal.enums; + +import jadx.api.InternalJadxTest; +import jadx.core.dex.nodes.ClassNode; + +import org.junit.Test; + +import static jadx.tests.utils.JadxMatchers.countString; +import static org.junit.Assert.assertThat; + +public class TestSwitchOverEnum extends InternalJadxTest { + + public enum Count { + ONE, TWO, THREE + } + + public int testEnum(Count c) { + switch (c) { + case ONE: + return 1; + case TWO: + return 2; + case THREE: + return 3; + } + return 0; + } + + @Test + public void test() { + ClassNode cls = getClassNode(TestSwitchOverEnum.class); + String code = cls.getCode().toString(); + System.out.println(code); + + assertThat(code, countString(1, "synthetic")); + assertThat(code, countString(2, "switch (c) {")); + assertThat(code, countString(2, "case ONE:")); + } +}