Skip to content

Commit

Permalink
Create transformer to fold basic operations with known inputs and out…
Browse files Browse the repository at this point in the history
…puts

This implementation only works with linear code. To allow folding of code obscured with opaque predicates, we'd want to create and run a transformer to solve those first.
  • Loading branch information
Col-E committed Jan 26, 2025
1 parent 4960ef3 commit ba7801b
Show file tree
Hide file tree
Showing 3 changed files with 557 additions and 4 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,322 @@
package software.coley.recaf.services.deobfuscation.transform.generic;

import jakarta.annotation.Nonnull;
import jakarta.annotation.Nullable;
import jakarta.enterprise.context.Dependent;
import jakarta.inject.Inject;
import org.objectweb.asm.Type;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.FieldInsnNode;
import org.objectweb.asm.tree.InsnList;
import org.objectweb.asm.tree.InsnNode;
import org.objectweb.asm.tree.LdcInsnNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.analysis.Frame;
import software.coley.recaf.info.JvmClassInfo;
import software.coley.recaf.services.inheritance.InheritanceGraph;
import software.coley.recaf.services.inheritance.InheritanceGraphService;
import software.coley.recaf.services.transform.JvmClassTransformer;
import software.coley.recaf.services.transform.JvmTransformerContext;
import software.coley.recaf.services.transform.TransformationException;
import software.coley.recaf.services.workspace.WorkspaceManager;
import software.coley.recaf.util.AsmInsnUtil;
import software.coley.recaf.util.analysis.ReAnalyzer;
import software.coley.recaf.util.analysis.ReInterpreter;
import software.coley.recaf.util.analysis.lookup.InvokeStaticLookup;
import software.coley.recaf.util.analysis.value.DoubleValue;
import software.coley.recaf.util.analysis.value.FloatValue;
import software.coley.recaf.util.analysis.value.IntValue;
import software.coley.recaf.util.analysis.value.LongValue;
import software.coley.recaf.util.analysis.value.ReValue;
import software.coley.recaf.util.analysis.value.StringValue;
import software.coley.recaf.workspace.model.Workspace;
import software.coley.recaf.workspace.model.bundle.JvmClassBundle;
import software.coley.recaf.workspace.model.resource.WorkspaceResource;

import java.util.ArrayList;
import java.util.List;

import static org.objectweb.asm.Opcodes.*;

/**
* A transformer that folds <i>basic</i> linear constant usages.
*
* @author Matt Coley
*/
@Dependent
public class LinearOpaqueConstantFoldingTransformer implements JvmClassTransformer {
private final InheritanceGraphService graphService;
private final WorkspaceManager workspaceManager;
private InheritanceGraph inheritanceGraph;

@Inject
public LinearOpaqueConstantFoldingTransformer(@Nonnull WorkspaceManager workspaceManager, @Nonnull InheritanceGraphService graphService) {
this.workspaceManager = workspaceManager;
this.graphService = graphService;
}

@Override
public void setup(@Nonnull JvmTransformerContext context, @Nonnull Workspace workspace) {
inheritanceGraph = workspace == workspaceManager.getCurrent() ?
graphService.getCurrentWorkspaceInheritanceGraph() :
graphService.newInheritanceGraph(workspace);
}

@Override
public void transform(@Nonnull JvmTransformerContext context, @Nonnull Workspace workspace,
@Nonnull WorkspaceResource resource, @Nonnull JvmClassBundle bundle,
@Nonnull JvmClassInfo initialClassState) throws TransformationException {
boolean dirty = false;
String className = initialClassState.getName();
ClassNode node = context.getNode(bundle, initialClassState);
for (MethodNode method : node.methods) {
InsnList instructions = method.instructions;
if (instructions == null)
continue;
ReInterpreter interpreter = new ReInterpreter(inheritanceGraph);
// TODO: A fleshed out implementation for each to facilitate
// - interpreter.setInvokeStaticLookup(...);
// - interpreter.setInvokeVirtualLookup(...);
// - interpreter.setGetStaticLookup(...);
ReAnalyzer analyzer = new ReAnalyzer(interpreter);
try {
boolean localDirty = false;
Frame<ReValue>[] frames = analyzer.analyze(className, method);
for (int i = 1; i < method.instructions.size() - 1; i++) {
// We must know the contents of the next frame, and it must have 1 or more values on the stack.
Frame<ReValue> nextFrame = frames[i + 1];
if (nextFrame == null || nextFrame.getStackSize() == 0)
continue;

// The next frame's stack top is the result of the operation.
ReValue nextFrameStackTop = nextFrame.getStack(nextFrame.getStackSize() - 1);
if (!nextFrameStackTop.hasKnownValue())
continue;

// Handle folding for specific instructions.
AbstractInsnNode instruction = method.instructions.get(i);
int opcode = instruction.getOpcode();
switch (opcode) {
case IADD:
case FADD:
case ISUB:
case FSUB:
case IMUL:
case FMUL:
case IDIV:
case FDIV:
case IREM:
case FREM:
case ISHL:
case ISHR:
case IUSHR:
case IAND:
case IXOR:
case IOR:
case DREM:
case DDIV:
case DMUL:
case DSUB:
case DADD:
case LUSHR:
case LSHR:
case LSHL:
case LREM:
case LDIV:
case LMUL:
case LSUB:
case LADD:
case LAND:
case LOR:
case LXOR:
case FCMPL:
case FCMPG:
case LCMP:
case DCMPL:
case DCMPG: {
// Get instruction of the top stack's contributing instruction.
AbstractInsnNode argument2 = method.instructions.get(i - 1);
while (argument2 != null && argument2.getOpcode() == NOP)
argument2 = argument2.getPrevious();
if (argument2 == null)
continue;

// Get instruction of the 2nd-to-top stack's contributing instruction.
AbstractInsnNode argument1 = argument2.getPrevious();
while (argument1 != null && argument1.getOpcode() == NOP)
argument1 = argument1.getPrevious();
if (argument1 == null)
continue;

// Both argument instructions must be value producers.
if (!isValueProducer(argument1) || !isValueProducer(argument2))
continue;

// We must have a viable replacement to offer.
AbstractInsnNode replacement = toInsn(nextFrameStackTop);
if (replacement == null)
continue;

// Replace the arguments and operation instructions with the replacement const value.
method.instructions.set(instruction, replacement);
method.instructions.set(argument2, new InsnNode(NOP));
method.instructions.set(argument1, new InsnNode(NOP));
localDirty = true;
break;
}
case INEG:
case FNEG:
case DNEG:
case LNEG:
case I2L:
case I2F:
case I2D:
case F2I:
case F2L:
case F2D:
case I2B:
case I2C:
case I2S:
case D2I:
case D2L:
case D2F:
case L2I:
case L2F:
case L2D: {
// Get instruction of the top stack's contributing instruction.
// It must also be a value producing instruction.
AbstractInsnNode argument = method.instructions.get(i - 1);
while (argument != null && argument.getOpcode() == NOP)
argument = argument.getPrevious();
if (argument == null || !isValueProducer(argument))
continue;

// We must have a viable replacement to offer.
AbstractInsnNode replacement = toInsn(nextFrameStackTop);
if (replacement == null)
continue;

// Replace the argument and operation instructions with the replacement const value.
method.instructions.set(instruction, replacement);
method.instructions.set(argument, new InsnNode(NOP));
localDirty = true;
break;
}
case INVOKESPECIAL:
case INVOKEINTERFACE:
case INVOKEVIRTUAL:
case INVOKESTATIC: {
// We'll have some loops further below that will set this flag to indicate to skip this
// instruction after the loop completes.
boolean skip = false;

// Get the contributing instructions.
MethodInsnNode min = (MethodInsnNode) instruction;
Type methodType = Type.getMethodType(min.desc);
List<AbstractInsnNode> argumentInstructions = new ArrayList<>(methodType.getArgumentCount());
int start = opcode == INVOKESTATIC ? 1 : 0; // non-static methods start at zero to include the method instance host.
for (int arg = methodType.getArgumentCount(); arg >= start; arg--) {
// Get the contributing instruction for this argument (or method instance host for non-static methods when arg == 0)
AbstractInsnNode lastArgumentInsn = argumentInstructions.isEmpty() ? null : argumentInstructions.getLast();
AbstractInsnNode argument = lastArgumentInsn == null ?
method.instructions.get(i - 1) : lastArgumentInsn.getPrevious();

// Argument must be a value producing instruction.
while (argument != null && argument.getOpcode() == NOP)
argument = argument.getPrevious();
if (argument == null || !isValueProducer(argument)) {
skip = true;
break;
}
argumentInstructions.add(argument);
}
if (skip)
continue;

// We must have a viable replacement to offer.
AbstractInsnNode replacement = toInsn(nextFrameStackTop);
if (replacement == null)
continue;

// Replace the arguments and invoke instructions with the replacement const value.
method.instructions.set(instruction, replacement);
for (AbstractInsnNode argument : argumentInstructions)
method.instructions.set(argument, new InsnNode(NOP));
localDirty = true;
break;
}
}
}

// We replace instructions with NOP in the code above because it reduces the headache of managing
// the proper index of instructions. Now that we are done, we'll prune any NOP instructions.
if (localDirty) {
dirty = true;
for (AbstractInsnNode insn : method.instructions.toArray()) {
if (insn.getOpcode() == NOP)
method.instructions.remove(insn);
}
}
} catch (Throwable t) {
throw new TransformationException("Error encountered when folding constants", t);
}
}
if (dirty)
context.setNode(bundle, initialClassState, node);
}

@Nonnull
@Override
public String name() {
return "Opaque constant folding";
}


/**
* @param insn
* Instruction to check.
*
* @return {@code true} when the instruction will produce a single value.
*/
private static boolean isValueProducer(@Nonnull AbstractInsnNode insn) {
if (AsmInsnUtil.isConstValue(insn))
return true;

// Fields always provide single values.
if (insn instanceof FieldInsnNode)
return true;

// We only support context-less provider methods at the moment.
return insn instanceof MethodInsnNode min &&
min.getOpcode() == INVOKESTATIC &&
min.desc.startsWith("()") &&
!min.desc.endsWith(")V");
}

/**
* @param value
* Value to convert.
*
* @return Instruction representing the value,
* or {@code null} if we don't/can't provide a mapping for the value content.
*/
@Nullable
@SuppressWarnings("OptionalGetWithoutIsPresent")
private static AbstractInsnNode toInsn(@Nonnull ReValue value) {
// Skip if value is not known.
if (!value.hasKnownValue())
return null;

// Map known value types to constant value instructions.
return switch (value) {
case IntValue intValue -> AsmInsnUtil.intToInsn(intValue.value().getAsInt());
case FloatValue floatValue -> AsmInsnUtil.floatToInsn((float) floatValue.value().getAsDouble());
case DoubleValue doubleValue -> AsmInsnUtil.doubleToInsn(doubleValue.value().getAsDouble());
case LongValue longValue -> AsmInsnUtil.longToInsn(longValue.value().getAsLong());
case StringValue stringValue -> new LdcInsnNode(stringValue.getText().get());
default -> null;
};
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,7 @@ else if (field.hasPrivateModifier())
}
}
} catch (Throwable t) {
throw new TransformationException("Analysis failure", t);
throw new TransformationException("Error encountered when computing static constants", t);
}
}
}
Expand Down
Loading

0 comments on commit ba7801b

Please sign in to comment.