From 7c755717332d718a2d1f0fa2e4844f55e8275915 Mon Sep 17 00:00:00 2001
From: Gabriel Tofvesson <gabriel.tofvesson@gmail.com>
Date: Thu, 16 Apr 2020 22:19:29 +0200
Subject: [PATCH] Implement field injection

---
 src/dev/w1zzrd/asm/Inject.java |  2 +-
 src/dev/w1zzrd/asm/Merger.java | 56 ++++++++++++++++++++++++++++------
 2 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/src/dev/w1zzrd/asm/Inject.java b/src/dev/w1zzrd/asm/Inject.java
index 5e6a147..bf440bc 100644
--- a/src/dev/w1zzrd/asm/Inject.java
+++ b/src/dev/w1zzrd/asm/Inject.java
@@ -6,5 +6,5 @@ import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
 @Retention(RetentionPolicy.RUNTIME)
-@Target({ElementType.CONSTRUCTOR, ElementType.METHOD})
+@Target({ElementType.CONSTRUCTOR, ElementType.METHOD, ElementType.FIELD})
 public @interface Inject { }
diff --git a/src/dev/w1zzrd/asm/Merger.java b/src/dev/w1zzrd/asm/Merger.java
index efc6798..2c5d4b3 100644
--- a/src/dev/w1zzrd/asm/Merger.java
+++ b/src/dev/w1zzrd/asm/Merger.java
@@ -14,7 +14,8 @@ import java.util.stream.Collectors;
 public class Merger {
 
     protected final ClassNode targetNode;
-    protected final List<MethodNode> inject = new ArrayList<>();
+    protected final List<MethodNode> injectMethods = new ArrayList<>();
+    protected final List<FieldNode> injectFields = new ArrayList<>();
 
 
     public Merger(String targetClass) throws IOException {
@@ -40,18 +41,23 @@ public class Merger {
 
     public void inject(MethodNode inject, String injectOwner) {
         transformInjection(inject, injectOwner);
-        this.inject.add(inject);
+        injectMethods.add(inject);
+    }
+
+    public void inject(FieldNode inject) {
+        injectFields.add(inject);
     }
 
     public void inject(ClassNode inject) {
         inject.methods.stream().filter(Merger::shouldInject).forEach(mNode -> inject(mNode, inject.name));
+        inject.fields.stream().filter(Merger::shouldInject).forEach(this::inject);
 
         if (inject.visibleAnnotations != null && inject.interfaces != null) {
 
             AsmAnnotation annot = getAnnotation("Ldev/w1zzrd/asm/InjectClass;", inject);
 
-            // If there is not inject annotation or there is an
-            // explicit request to not inject interfaces, just return
+            // If there is not injectMethods annotation or there is an
+            // explicit request to not injectMethods interfaces, just return
             if (annot == null || (annot.hasEntry("injectInterfaces") && !annot.getEntry("injectInterfaces", Boolean.class)))
                 return;
 
@@ -130,12 +136,24 @@ public class Merger {
     public byte[] toByteArray() {
         ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
 
-        List<MethodNode> original = targetNode.methods;
+        // Adapt nodes as necessary
+        List<MethodNode> originalMethods = targetNode.methods;
         targetNode.methods = targetNode.methods.stream().filter(this::isNotInjected).collect(Collectors.toList());
-        targetNode.accept(writer);
-        targetNode.methods = original;
 
-        inject.forEach(node -> node.accept(writer));
+        List<FieldNode> originalFields = targetNode.fields;
+        targetNode.fields = targetNode.fields.stream().filter(this::isNotInjected).collect(Collectors.toList());
+
+
+        // Accept writer
+        targetNode.accept(writer);
+
+        // Restore originals
+        targetNode.methods = originalMethods;
+        targetNode.fields = originalFields;
+
+        // Inject methods and fields
+        injectMethods.forEach(node -> node.accept(writer));
+        injectFields.forEach(node -> node.accept(writer));
 
         return writer.toByteArray();
     }
@@ -167,13 +185,20 @@ public class Merger {
 
 
     protected boolean isNotInjected(MethodNode node) {
-        for (MethodNode mNode : inject)
+        for (MethodNode mNode : injectMethods)
             if (methodNodeEquals(node, mNode))
                 return false;
 
         return true;
     }
 
+    protected boolean isNotInjected(FieldNode node) {
+        for (FieldNode mNode : injectFields)
+            if (fieldNodeEquals(node, mNode))
+                return false;
+
+        return true;
+    }
 
 
     // To be used instead of referencing object constructs
@@ -208,6 +233,10 @@ public class Merger {
     }
 
     protected static boolean methodNodeEquals(MethodNode a, MethodNode b) {
+        return a.name.equals(b.name) && Objects.equals(a.desc, b.desc);
+    }
+
+    protected static boolean fieldNodeEquals(FieldNode a, FieldNode b) {
         return a.name.equals(b.name) && Objects.equals(a.signature, b.signature);
     }
 
@@ -220,6 +249,15 @@ public class Merger {
         return false;
     }
 
+    protected static boolean shouldInject(FieldNode node) {
+        if (node.visibleAnnotations == null) return false;
+        for (AnnotationNode aNode : node.visibleAnnotations)
+            if (aNode.desc.equals("Ldev/w1zzrd/asm/Inject;"))
+                return true;
+
+        return false;
+    }
+
     public static ClassNode getClassNode(URL url) throws IOException {
         return readClass(getClassBytes(url));
     }