From 1417136f748617612fed4b291fbe7cdca47d0238 Mon Sep 17 00:00:00 2001
From: 3arthqu4ke <56741599+3arthqu4ke@users.noreply.github.com>
Date: Thu, 12 Jan 2023 17:54:00 +0100
Subject: [PATCH] [1.6.1] Fixed JavaAgent and added unit test

---
 headlessmc-lwjgl/build.gradle                 | 22 +++++++++++++++---
 .../headlessmc/lwjgl/agent/LwjglAgent.java    |  2 +-
 .../headlessmc/lwjgl/agent/AgentTest.java     | 23 +++++++++++++++++++
 .../LaunchWrapperTarget.java                  |  7 ++++--
 4 files changed, 48 insertions(+), 6 deletions(-)
 create mode 100644 headlessmc-lwjgl/src/test/java/me/earth/headlessmc/lwjgl/agent/AgentTest.java

diff --git a/headlessmc-lwjgl/build.gradle b/headlessmc-lwjgl/build.gradle
index 1c47164e..80c3043e 100644
--- a/headlessmc-lwjgl/build.gradle
+++ b/headlessmc-lwjgl/build.gradle
@@ -55,12 +55,28 @@ jar {
     archivesBaseName = 'headlessmc-lwjgl'
     //noinspection GroovyAssignabilityCheck
     manifest {
-        // TODO: test tweaker?
         attributes 'TweakClass': 'me.earth.headlessmc.lwjgl.launchwrapper.LwjglTweaker'
-        // TODO: test java agent
         attributes 'Premain-Class': 'me.earth.headlessmc.lwjgl.agent.LwjglAgent'
         attributes 'Agent-Class': 'me.earth.headlessmc.lwjgl.agent.LwjglAgent'
         attributes 'Can-Redefine-Classes': true
         attributes 'Can-Retransform-Classes': true
     }
-}
\ No newline at end of file
+}
+
+task testAgent(type: Test, dependsOn: jar) {
+    group = 'verification'
+    useJUnitPlatform()
+    testLogging {
+        events "failed"
+        exceptionFormat "full"
+    }
+
+    jvmArgs = [
+        "-javaagent:" + projectDir + "/build/libs/headlessmc-lwjgl.jar",
+        '-Dhmc.lwjgl.agenttest=true'
+    ]
+
+    filter {
+        includeTestsMatching 'me.earth.headlessmc.lwjgl.agent.AgentTest'
+    }
+}
diff --git a/headlessmc-lwjgl/src/main/java/me/earth/headlessmc/lwjgl/agent/LwjglAgent.java b/headlessmc-lwjgl/src/main/java/me/earth/headlessmc/lwjgl/agent/LwjglAgent.java
index 7d8c6f49..ee6332ab 100644
--- a/headlessmc-lwjgl/src/main/java/me/earth/headlessmc/lwjgl/agent/LwjglAgent.java
+++ b/headlessmc-lwjgl/src/main/java/me/earth/headlessmc/lwjgl/agent/LwjglAgent.java
@@ -31,7 +31,7 @@ public byte[] transform(ClassLoader loader, String className,
                             ProtectionDomain protectionDomain,
                             byte[] classfileBuffer)
         throws IllegalClassFormatException {
-        if (className.contains("lwjgl")) {
+        if (className != null && className.startsWith("org/lwjgl")) {
             ClassNode node = AsmUtil.read(classfileBuffer);
             transformer.transform(node);
             // TODO: make writer.getClassLoader() return the given loader?
diff --git a/headlessmc-lwjgl/src/test/java/me/earth/headlessmc/lwjgl/agent/AgentTest.java b/headlessmc-lwjgl/src/test/java/me/earth/headlessmc/lwjgl/agent/AgentTest.java
new file mode 100644
index 00000000..b8942d1b
--- /dev/null
+++ b/headlessmc-lwjgl/src/test/java/me/earth/headlessmc/lwjgl/agent/AgentTest.java
@@ -0,0 +1,23 @@
+package me.earth.headlessmc.lwjgl.agent;
+
+import me.earth.headlessmc.lwjgl.testlaunchwrapper.LaunchWrapperTarget;
+import org.junit.jupiter.api.Assumptions;
+import org.junit.jupiter.api.Test;
+import org.lwjgl.AbstractLwjglClass;
+
+import java.lang.reflect.Modifier;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+
+public class AgentTest {
+    @Test
+    public void testWithAgent() {
+        Assumptions.assumeTrue(Boolean.parseBoolean(
+            System.getProperty("hmc.lwjgl.agenttest")));
+
+        LaunchWrapperTarget.testLwjglClasses();
+        assertFalse(Modifier.isAbstract(
+            AbstractLwjglClass.class.getModifiers()));
+    }
+
+}
diff --git a/headlessmc-lwjgl/src/test/java/me/earth/headlessmc/lwjgl/testlaunchwrapper/LaunchWrapperTarget.java b/headlessmc-lwjgl/src/test/java/me/earth/headlessmc/lwjgl/testlaunchwrapper/LaunchWrapperTarget.java
index cea21d27..cbda0559 100644
--- a/headlessmc-lwjgl/src/test/java/me/earth/headlessmc/lwjgl/testlaunchwrapper/LaunchWrapperTarget.java
+++ b/headlessmc-lwjgl/src/test/java/me/earth/headlessmc/lwjgl/testlaunchwrapper/LaunchWrapperTarget.java
@@ -26,6 +26,11 @@ public static void main(String[] args) {
         assertEquals(Launch.classLoader,
                      LwjglInterface.class.getClassLoader());
 
+        testLwjglClasses();
+        System.setProperty(LaunchWrapperTest.PASSED, "true");
+    }
+
+    public static void testLwjglClasses() {
         val lwjgl = Lwjgl.factoryMethod("test");
         assertNotNull(lwjgl);
         LwjglInstrumentationTest.testRedirections(lwjgl, lwjgl.getClass());
@@ -39,8 +44,6 @@ public static void main(String[] args) {
         val iLwjgl = LwjglInterface.factoryMethod("test");
         assertNotNull(iLwjgl);
         LwjglInstrumentationTest.testRedirections(iLwjgl, LwjglInterface.class);
-
-        System.setProperty(LaunchWrapperTest.PASSED, "true");
     }
 
 }