Add support for agent self-testing, to confirm JVM compat

Author: pimterry

build.gradle

@@ -164,4 +164,4 @@ task quickTest(type: Test) {
 task distTest(type: Test) {
     environment 'TEST_JAR', tasks.distJar.getArchiveFile().get().asFile.toString()
     dependsOn('distJar')
-}
+}

src/main/kotlin/tech/httptoolkit/javaagent/AgentMain.kt

@@ -51,6 +51,14 @@ fun agentmain(arguments: String?, instrumentation: Instrumentation) {
     if (arguments.isNullOrEmpty()) {
         throw Error("Can't attach proxy agent without configuration arguments")
     }
+
+    // If attached as a test, we don't intercept anything, we're just checking that it's
+    // possible to attach in the first place with the current VM.
+    if (arguments == "attach-test") {
+        println("Agent attach test successful")
+        return
+    };
+
     val config = getConfigFromArg(arguments)
     interceptAllHttps(config, instrumentation)
 }

src/main/kotlin/tech/httptoolkit/javaagent/AttachMain.kt

@@ -6,6 +6,7 @@ import com.sun.tools.attach.*
 import kotlin.system.exitProcess
 import java.lang.management.ManagementFactory
 import java.io.File
+import java.lang.IllegalArgumentException
 
 // This file is the only one that uses com.sun.tools.attach.VirtualMachine. That's important because that
 // requires tools.jar to be in your classpath (i.e. it requires a JDK not a JRE). If you run this without
@@ -24,41 +25,67 @@ val x: Class<*> = try { // Var declaration required for static init for some rea
 
 // If run directly, can either list potential targets (list-targets) or attach to a target (pid, ...config)
 fun main(args: Array<String>) {
-    if (args.size == 1 && args[0] == "list-targets") {
-        // This isn't guaranteed to work everywhere, but it should work in most places:
-        val (pid) = ManagementFactory.getRuntimeMXBean().name.split("@")
-
-        val vms = VirtualMachine.list()
-        if (vms.isEmpty()) {
-            // VMs should never be empty, because at the very least _we_ should be in there! If it's empty then
-            // scanning isn't working at all, and we should fail clearly.
-            System.err.println("Can't scan for attachable JVMs. Are we running in a JRE instead of a JDK?")
-            exitProcess(4)
+    // Self-test ensures that the JVM we're using is capable of scanning & attachment. It *doesn't* fully
+    // test its ability to transform classes as we'd like.
+    if (args.size == 1 && args[0] == "self-test") {
+        val selfAttachAllowed = System.getProperty("jdk.attach.allowAttachSelf")
+        if (selfAttachAllowed != "true") {
+            throw IllegalArgumentException("Cannot run self-test without -Djdk.attach.allowAttachSelf=true")
         }
 
+        getTargets() // Check we can scan for targets
+        attachAgent(getOwnPid(), "attach-test") // Check we can attach (against ourselves)
+    } else if (args.size == 1 && args[0] == "list-targets") {
+        // List-targets prints a list of pid:name target paids
+        val pid = getOwnPid()
+        val vms = getTargets()
         vms.forEach { vmd ->
             if (vmd.id() != pid) {
                 println("${vmd.id()}:${vmd.displayName()}")
             }
         }
 
         exitProcess(0)
-    } else if (args.size != 4) {
+    } else if (args.size == 4) {
+        // 4-args format attaches to a target pid with the given config values
+        val (pid, proxyHost, proxyPort, certPath) = args
+        attachAgent(pid, formatConfigArg(proxyHost, proxyPort, certPath))
+    } else {
         System.err.println("Usage: java -jar <agent.jar> <target-PID> <proxyHost> <proxyPort> <path-to-certificate>")
+        System.err.println("Or pass a single 'self-test' or 'list-target' arg to check capabilities or scan for pids")
         exitProcess(2)
     }
+}
+
+fun getOwnPid(): String {
+    // This should work in general, but it's implementation dependent:
+    return ManagementFactory.getRuntimeMXBean().name.split("@")[0]
+}
 
-    val (pid, proxyHost, proxyPort, certPath) = args
+fun getTargets(): List<VirtualMachineDescriptor> {
+    val vms = VirtualMachine.list()
+    if (vms.isEmpty()) {
+        // VMs should never be empty, because at the very least _we_ should be in there! If it's empty then
+        // scanning isn't working at all, and we should fail clearly.
+        System.err.println("Can't scan for attachable JVMs. Are we running in a JRE instead of a JDK?")
+        exitProcess(4)
+    }
+    return vms
+}
 
+fun attachAgent(
+    pid: String,
+    agentArg: String
+) {
     val jarPath = File(
         ConstantProxySelector::class.java // Any arbitrary class defined inside this JAR
             .protectionDomain.codeSource.location.path
     ).absolutePath
 
-    // Inject the agent with our config arguments into the target VM
+    // Inject the agent into the target VM
     try {
         val vm: VirtualMachine = VirtualMachine.attach(pid)
-        vm.loadAgent(jarPath, formatConfigArg(proxyHost, proxyPort, certPath))
+        vm.loadAgent(jarPath, agentArg)
         vm.detach()
     } catch (e: AgentLoadException) {
         if (e.message == "0") {

src/test/kotlin/IntegrationTests.kt

@@ -42,6 +42,21 @@ val wireMockServer = WireMockServer(options()
 val runningProcs = arrayListOf<Process>()
 
 class IntegrationTests : StringSpec({
+    "Launching a self test should return successfully" {
+        val proc = ProcessBuilder(
+            javaPath,
+            "-Djdk.attach.allowAttachSelf=true",
+            "-jar", AGENT_JAR_PATH,
+            "self-test"
+        ).start()
+        runningProcs.add(proc)
+
+        proc.waitFor(10, TimeUnit.SECONDS)
+
+        proc.isAlive.shouldBe(false)
+        proc.exitValue().shouldBe(0)
+    }
+
     "Launching with list-targets should return successfully" {
         val proc = ProcessBuilder(
             javaPath,