Skip to content

Commit

Permalink
Enhanced method patch detection
Browse files Browse the repository at this point in the history
  • Loading branch information
Su5eD committed Jul 15, 2024
1 parent e8f6c7e commit 9818239
Show file tree
Hide file tree
Showing 7 changed files with 197 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
// Source: https://git.sleeping.town/Nil/NilLoader/src/commit/d66d783a5f7ac72a3688594335b3285fcb975b07/src/main/java/nilloader/api/lib/mini/PatchContext.java
public class InsnComparator {
public static final int IGNORE_VAR_INDEX = 0x001;
public static final int IGNORE_LINE_NUMBERS = 0x010;

public static boolean instructionsEqual(AbstractInsnNode a, AbstractInsnNode b) {
return instructionsEqual(a, b, 0);
Expand Down Expand Up @@ -78,7 +79,7 @@ public static boolean instructionsEqual(AbstractInsnNode a, AbstractInsnNode b,
} else if (a instanceof LineNumberNode) {
LineNumberNode la = (LineNumberNode) a;
LineNumberNode lb = (LineNumberNode) b;
return la.line == lb.line && instructionsEqual(la.start, lb.start);
return (flags & IGNORE_LINE_NUMBERS) != 0 || la.line == lb.line && instructionsEqual(la.start, lb.start);
} else if (a instanceof LookupSwitchInsnNode) {
LookupSwitchInsnNode la = (LookupSwitchInsnNode) a;
LookupSwitchInsnNode lb = (LookupSwitchInsnNode) b;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -75,11 +75,15 @@ public boolean test(InstructionMatcher other, int flags) {
}

public static boolean test(InsnList first, InsnList second) {
return test(first, second, 0);
}

public static boolean test(InsnList first, InsnList second, int flags) {
if (first.size() == second.size()) {
for (int i = 0; i < first.size(); i++) {
AbstractInsnNode insn = first.get(i);
AbstractInsnNode otherInsn = second.get(i);
if (!InsnComparator.instructionsEqual(insn, otherInsn)) {
if (!InsnComparator.instructionsEqual(insn, otherInsn, flags)) {
return false;
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
import org.sinytra.adapter.patch.api.*;
import org.sinytra.adapter.patch.transformer.param.TransformParameters;
import org.sinytra.adapter.patch.util.AdapterUtil;
import org.sinytra.adapter.patch.util.MethodQualifier;
import org.sinytra.adapter.patch.util.OpcodeUtil;

import java.util.*;
Expand All @@ -33,12 +32,8 @@ public Collection<String> getAcceptedAnnotations() {
public Patch.Result apply(ClassNode classNode, MethodNode methodNode, MethodContext methodContext, PatchContext context) {
// Sanity check
boolean isStatic = (methodNode.access & Opcodes.ACC_STATIC) == Opcodes.ACC_STATIC;
MethodQualifier qualifier = methodContext.getTargetMethodQualifier();
if (qualifier == null) {
return Patch.Result.PASS;
}

String owner = Objects.requireNonNullElse(qualifier.internalOwnerName(), this.targetClass);
String owner = Optional.ofNullable(methodContext.findCleanInjectionTarget()).map(t -> t.classNode().name).orElse(this.targetClass);
boolean isInherited = context.environment().inheritanceHandler().isClassInherited(this.targetClass, owner);
Candidates candidates = findCandidates(classNode, methodNode);
if (!candidates.canMove(classNode, isInherited)) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,155 @@
package org.sinytra.adapter.patch.transformer.dynfix;

import com.mojang.datafixers.util.Pair;
import org.jetbrains.annotations.Nullable;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.tree.*;
import org.sinytra.adapter.patch.analysis.InsnComparator;
import org.sinytra.adapter.patch.analysis.InstructionMatcher;
import org.sinytra.adapter.patch.api.MethodContext;
import org.sinytra.adapter.patch.api.MixinConstants;
import org.sinytra.adapter.patch.api.Patch;
import org.sinytra.adapter.patch.transformer.ExtractMixin;
import org.sinytra.adapter.patch.transformer.ModifyInjectionPoint;
import org.sinytra.adapter.patch.transformer.ModifyInjectionTarget;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

public class DynFixMethodComparison implements DynamicFixer<DynFixMethodComparison.Data> {
public record Data(AbstractInsnNode cleanInjectionInsn) {}

@Nullable
@Override
public Data prepare(MethodContext methodContext) {
if (methodContext.methodAnnotation().matchesDesc(MixinConstants.INJECT)) {
MethodContext.TargetPair cleanInjectionTarget = methodContext.findCleanInjectionTarget();
List<AbstractInsnNode> cleanInsns = methodContext.findInjectionTargetInsns(cleanInjectionTarget);
if (cleanInsns.size() == 1) {
return new Data(cleanInsns.getFirst());
}
}
return null;
}

@Override
public Patch.Result apply(ClassNode classNode, MethodNode methodNode, MethodContext methodContext, Data data) {
List<InsnList> cleanLabels = getLabelsInMethod(methodContext.findCleanInjectionTarget().methodNode());
List<InsnList> cleanLabelsOriginal = List.copyOf(cleanLabels);

AbstractInsnNode cleanInjectionInsn = data.cleanInjectionInsn();
List<InsnList> cleanMatchedLabels = cleanLabels.stream()
.filter(insns -> StreamSupport.stream(insns.spliterator(), false)
.anyMatch(insn -> InsnComparator.instructionsEqual(cleanInjectionInsn, insn)))
.toList();
if (cleanMatchedLabels.size() != 1) {
return Patch.Result.PASS;
}
InsnList cleanLabel = cleanMatchedLabels.getFirst();

List<InsnList> dirtyLabels = getLabelsInMethod(methodContext.findDirtyInjectionTarget().methodNode());
List<InsnList> dirtyLabelsOriginal = List.copyOf(dirtyLabels);

Map<InsnList, InsnList> matchedLabels = new LinkedHashMap<>();
for (InsnList cleanInsns : cleanLabelsOriginal) {
for (InsnList dirtyInsns : dirtyLabels) {
if (InstructionMatcher.test(cleanInsns, dirtyInsns, InsnComparator.IGNORE_VAR_INDEX | InsnComparator.IGNORE_LINE_NUMBERS)) {
matchedLabels.put(cleanInsns, dirtyInsns);
cleanLabels.remove(cleanInsns);
dirtyLabels.remove(dirtyInsns);
break;
}
}
}

Pair<InsnList, InsnList> patchRange = findPatchHunkRange(cleanLabel, cleanLabelsOriginal, matchedLabels);
if (patchRange == null) {
return Patch.Result.PASS;
}

ClassNode cleanTargetClass = methodContext.findCleanInjectionTarget().classNode();
List<InsnList> hunkLabels = dirtyLabelsOriginal.subList(dirtyLabelsOriginal.indexOf(patchRange.getFirst()) + 1, dirtyLabelsOriginal.indexOf(patchRange.getSecond()));
for (InsnList insns : hunkLabels) {
for (AbstractInsnNode insn : insns) {
if (insn instanceof MethodInsnNode minsn && minsn.getOpcode() == Opcodes.INVOKESTATIC && !minsn.owner.equals(cleanTargetClass.name)) {
// Looks like some code was moved into a static method outside this class
// Attempt extracting mixin
// TODO Create universal transform builder class so that we don't need to reference transformer classes directly
ClassNode targetClass = methodContext.patchContext().environment().dirtyClassLookup().getClass(minsn.owner).orElse(null);
if (targetClass != null) {
MethodNode targetMethod = methodContext.patchContext().environment().dirtyClassLookup().findMethod(minsn.owner, minsn.name, minsn.desc).orElse(null);
if (targetMethod != null && !methodContext.findInjectionTargetInsns(new MethodContext.TargetPair(targetClass, targetMethod)).isEmpty()) {
Patch.Result extractResult = new ExtractMixin(minsn.owner).apply(classNode, methodNode, methodContext);
if (extractResult != Patch.Result.PASS) {
return extractResult.or(new ModifyInjectionTarget(List.of(minsn.name + minsn.desc), ModifyInjectionTarget.Action.OVERWRITE).apply(classNode, methodNode, methodContext));
}
}
}
}
}
}

// If extraction fails, fall back to injecting at the nearest method call in dirty code
for (InsnList insns : hunkLabels) {
for (AbstractInsnNode insn : insns) {
if (insn instanceof MethodInsnNode minsn) {
String newInjectionPoint = Type.getObjectType(minsn.owner).getDescriptor() + minsn.name + minsn.desc;
return new ModifyInjectionPoint("INVOKE", newInjectionPoint, true, false).apply(classNode, methodNode, methodContext);
}
}
}

return Patch.Result.PASS;
}

@Nullable
private static Pair<InsnList, InsnList> findPatchHunkRange(InsnList cleanLabel, List<InsnList> cleanLabels, Map<InsnList, InsnList> matchedLabels) {
// Find last matched dirty label BEFORE the injection point
InsnList dirtyLabelBefore = Stream.iterate(cleanLabels.indexOf(cleanLabel), i -> i > 0, i -> i - 1)
.map(i -> matchedLabels.get(cleanLabels.get(i)))
.filter(Objects::nonNull)
.findFirst()
.orElse(null);
if (dirtyLabelBefore == null) {
return null;
}

// Find first matched dirty label AFTER the injection point
InsnList dirtyLabelAfter = Stream.iterate(cleanLabels.indexOf(cleanLabel), i -> i < cleanLabels.size(), i -> i + 1)
.map(i -> matchedLabels.get(cleanLabels.get(i)))
.filter(Objects::nonNull)
.findFirst()
.orElse(null);
if (dirtyLabelAfter == null) {
return null;
}

return Pair.of(dirtyLabelBefore, dirtyLabelAfter);
}

private static List<InsnList> getLabelsInMethod(MethodNode methodNode) {
Map<LabelNode, LabelNode> labels = StreamSupport.stream(methodNode.instructions.spliterator(), false)
.map(a -> a instanceof LabelNode label ? label : null)
.filter(Objects::nonNull)
.collect(Collectors.toMap(Function.identity(), l -> new LabelNode()));
List<InsnList> list = new ArrayList<>();
InsnList workingList = null;
for (AbstractInsnNode insn : methodNode.instructions) {
if (insn instanceof FrameNode) {
continue;
}
if (insn instanceof LabelNode) {
if (workingList != null) {
list.add(workingList);
}
workingList = new InsnList();
}
workingList.add(insn.clone(labels));
}
return list;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ public class DynamicInjectionPointPatch implements MethodTransform {
new DynFixResolveAmbigousTarget(),
new DynFixSplitMethod(),
// Have this one always come last
new DynFixMethodComparison(),
new DynFixArbitraryInjectionPoint()
);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,18 @@ void testModifiedSliceTarget() throws Exception {
);
}

@Test
void testCompareModifiedMethod() throws Exception {
// TODO This can correctly determine the injection point in the extracted method now,
// but fails to extract because the mixin calls an injected unique method.
assertSameCode(
"org/sinytra/adapter/test/mixin/LivingEntityMixin",
"onUnderwater",
assertTargetMethod(),
assertInjectionPoint()
);
}

@Override
protected LoadResult load(String className) throws Exception {
final ClassNode patched = loadClass(className);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,12 @@

import net.minecraft.world.entity.LivingEntity;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.Unique;
import org.spongepowered.asm.mixin.injection.At;
import org.spongepowered.asm.mixin.injection.Inject;
import org.spongepowered.asm.mixin.injection.ModifyVariable;
import org.spongepowered.asm.mixin.injection.Slice;
import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;

@Mixin(LivingEntity.class)
public class LivingEntityMixin {
Expand Down Expand Up @@ -42,4 +45,22 @@ private float getSlipperinessForIceSkates(float slipperiness) {
private float getSlipperinessForIceSkatesExpected(float slipperiness) {
return slipperiness;
}

// https://github.com/Earthcomputer/clientcommands/blob/b5ed9155bdab5606498f6dc6538e5a6bdfad3b70/src/main/java/net/earthcomputer/clientcommands/mixin/rngevents/LivingEntityMixin.java#L60
@Inject(method = "baseTick",
slice = @Slice(from = @At(value = "FIELD", target = "Lnet/minecraft/tags/FluidTags;WATER:Lnet/minecraft/tags/TagKey;", ordinal = 0)),
at = @At(value = "INVOKE", target = "Lnet/minecraft/world/level/Level;getBlockState(Lnet/minecraft/core/BlockPos;)Lnet/minecraft/world/level/block/state/BlockState;", ordinal = 0))
public void onUnderwater(CallbackInfo ci) {
ourUniqueMethod(); // Prevent extraction
}

@Inject(method = "baseTick", at = @At(value = "INVOKE", target = "Lnet/minecraft/world/entity/LivingEntity;getAirSupply()I"))
public void onUnderwaterExpected(CallbackInfo ci) {
ourUniqueMethod();
}

@Unique
private void ourUniqueMethod() {
// Noop
}
}

0 comments on commit 9818239

Please sign in to comment.