アプリ開発備忘録

PlayStationMobile、Android、UWPの開発備忘録

Java Agentで実行時にクラスを差し替える

Java Kotlin

Java Agentとは

java コマンドを実行するとhelpが出力されるわけですが、そこにJava Agentについて書いてあります。

    -javaagent:<jarpath>[=<options>]
                  Javaプログラミング言語エージェントをロードする。java.lang.instrumentを参照

Instrumentとは

既存のアプリケーションを変更、コードの追加をできるものです。コンパイル時と実行時に行うことができます。ここでは実行時に差し替える方法だけを紹介します。

Java Agentの使用方法

既に main.jar があるとします。Agentは別で agent.jar という名前で作るとしましょう。
以下のようにして実行します。

java -javaagent:agent.jar -jar main.jar

Java Agentの作り方

今回は Java8 と Kotlin1.4.10 と Gradle6.3 を使用します。
static な premain という関数を以下の引数と共に作成します。Agent自体の仕様では instrumentation 引数は無しでも大丈夫ですが、今回はクラスの差し替えを行いたいので引数は必要です。
トップレベルに premain を置いても大丈夫です。 build.gradle で PremainKt に置き換え忘れだけ気をつけてください。

Premain.kt

package com.exsample

class Premain {
    companion object {
        @JvmStatic
        fun premain(agentArgs: String?, instrumentation: Instrumentation) {
            println("loading premain")
        }
    }
}

javaの通常のアプリケーションでは Main-Class を書くわけですが、Java Agentでは Premain-Class が必要となります。
build.gradle

plugins {
    id 'java'
    id 'org.jetbrains.kotlin.jvm' version '1.4.10'
}

group 'com.exsample'
version '1.0-SNAPSHOT'

repositories {
    mavenCentral()
}

dependencies {
    implementation "org.jetbrains.kotlin:kotlin-stdlib"
}

jar {
    manifest {
        attributes 'Premain-Class': 'com.exsample.Premain'
    }

    from { configurations.compileClasspath.collect { it.isDirectory() ? it : zipTree(it) } }
}

これで main.jar の Main-Class のmainが実行される前に Premain-Class のpremainが呼ばれます。

Instrumentationでクラスを差し替える

使用例

ここからは具体例を出して実装していきます。 NewRelicというサービスがあります。エラーやメトリクスを送信できます。でもデバッグ時に送られてしまったらまずいですよね。なので、NewRelicのjavaではJava Agentの仕組みを使用しています。

Staticな関数を呼んで使用します。

NewRelic.noticeError(IllegalStateException("error"))

noticeErrorの定義を見ると、以下のように空の定義があります。これを本番ではAgentを使用して中身の実装のあるクラスに差し替えています。

public static void noticeError(Throwable throwable) {
}

https://github.com/newrelic/newrelic-java-agent/blob/c224d2a9e5e2f3559eadc2141f6fe4b777da469a/newrelic-api/src/main/java/com/newrelic/api/agent/NewRelic.java

こちらが実装

public static void noticeError(Throwable throwable) {
    AgentBridge.publicApi.noticeError(throwable);
}

https://github.com/newrelic/newrelic-java-agent/blob/c224d2a9e5e2f3559eadc2141f6fe4b777da469a/agent-bridge/src/main/java/com/newrelic/api/agent/NewRelic.java

これにより、本番時にだけ処理が実行されます。しかし、デバッグ時には別途ログに出力したいとなった場合、どうすれば良いでしょうか。正直 NewRelic クラスをラップして分岐してしまえば良いのですが、 今回は同じAgentの仕組みで差し替えていきます。

今回差し替えたい実装

以下の3つを差し替えてみます。

  • 戻り値が無い、Javaで定義されている引数があるもの
  • 戻り値が無い、ライブラリで定義されている引数があるもの
  • ライブラリで定義されている戻り値があるもの

main.jar

fun main(args: Array<String>) {
    NewRelic.noticeError(IllegalStateException("error"))
    NewRelic.getAgent()
    NewRelic.setRequestAndResponse(null, null)
}

plugins {
    id 'java'
    id 'org.jetbrains.kotlin.jvm' version '1.4.10'
}

group 'com.example'
version '1.0-SNAPSHOT'

repositories {
    mavenCentral()
}

dependencies {
    implementation "org.jetbrains.kotlin:kotlin-stdlib"
    implementation "com.newrelic.agent.java:newrelic-api:3.27.0"
}

jar {
    manifest {
        attributes 'Main-Class': 'MainKt'
    }
    from { configurations.compileClasspath.collect { it.isDirectory() ? it : zipTree(it) } }
}

今回呼んでいる newrelic-api の一部抜粋です。

package com.newrelic.api.agent;

public final class NewRelic {
    public static com.newrelic.api.agent.Agent getAgent() {
        return com.newrelic.api.agent.NoOpAgent.INSTANCE;
    }
    public static void noticeError(Throwable throwable) {
    }
    public static void setRequestAndResponse(com.newrelic.api.agent.Request request, com.newrelic.api.agent.Response response) {
    }
}

差し替える実装

agent.jar の方に差し替えたいクラスを定義します。package nameは同じにします。

package com.newrelic.api.agent;

public final class NewRelic {
    public static com.newrelic.api.agent.Agent getAgent() {
        // 問題なく NoOpAgent.INSTANCE を返したい。(NoOpAgentはprivate)
        return com.newrelic.api.agent.NoOpAgent.INSTANCE;
    }
    public static void noticeError(Throwable throwable) {
        // 標準出力に出すように差し替えてみる
        System.out.println(throwable);
    }
    public static void setRequestAndResponse(com.newrelic.api.agent.Request request, com.newrelic.api.agent.Response response) {
        // 問題なく呼べるか
    }
}

agent.jar の方ではライブラリを読み込んでいないので、上記定義だけではコンパイルできません。コンパイルできるようにダミーの定義をしておきます。

package com.newrelic.api.agent;

class NoOpAgent {
    static final Agent INSTANCE = new Agent() {
    };
}

package com.newrelic.api.agent;

interface Agent {}

package com.newrelic.api.agent;

interface Request {}

package com.newrelic.api.agent;

interface Request {}

差し替え処理を実装するには

クラスを色々いじるには ClassLoader をいじります。
ClassLoader には BootstrapClassLoader と SystemClassLoader があります。
BootstrapClassLoader が先に読み込まれ、次に SystemClassLoader が読み込まれます。
BootstrapClassLoader がjavaの基本的なライブラリ群( rt.jar )、 その他ユーザーが定義したものは SystemClassLoader で読み込まれるようです。
※このへんはそこまで把握ができていませんので解説が適当

Class<*>#getClassLoader で ClassLoader を取得するわけですが、これがnullの場合は BootstrapClassLoader が使用されているようです。
例)

// null
String::class.java.classLoader

差し替えを行いたいので、優先的に読み込ませるために appendToBootstrapClassLoaderSearch の方を使用します。

@JvmStatic
fun premain(agentArgs: String?, instrumentation: Instrumentation) {
    instrumentation.appendToBootstrapClassLoaderSearch(JarFile(/*java.io.File*/))
    instrumentation.appendToSystemClassLoaderSearch(JarFile(/*java.io.File*/))
}

ここに premain.jar を指定すれば良いのですが、必要のないクラスファイルも置き換えられてしまいます。なので必要なクラスファイルだけを取り出してjarを作り、それを指定します。

Jarファイルを作る処理

このへんの処理は https://github.com/newrelic/newrelic-java-agent を参考にしています。

AgentのClassLoaderです。

class JVMAgentClassLoader(
    urls: Array<URL>,
    parent: ClassLoader?
) : URLClassLoader(urls, parent) {
    companion object {
        init {
            try {
                ClassLoader.registerAsParallelCapable()
            } catch (t: Throwable) {
                System.err.println(
                    MessageFormat.format(
                        "Unable to register as parallel-capable: {0}",
                        t
                    )
                )
            }
        }
    }
}

クラスファイルをByteArrayに変換する処理です。

class ClassfileUtil(
    private val loader: ClassLoader
) {
    fun getByteArray(classFullName: String): ByteArray {
        val internalPath = classFullName.replace('.', '/') + ".class"

        val res = loader.getResource(internalPath)!!

        return res.openStream().use {
            it.readBytes()
        }
    }
    
    companion object {
        // AgentのClassLoaderをセットします。
        fun getAgentClassLoader(): ClassfileUtil {
            val agentJarUrl = Premain::class.java.protectionDomain.codeSource.location
            val loader = JVMAgentClassLoader(
                arrayOf(agentJarUrl),
                null
            )
            return ClassfileUtil(loader)
        }

        // main.jarのClassLoaderをセットします。
        fun getClassLoader() : ClassfileUtil {
            return ClassfileUtil(Premain::class.java.classLoader)
        }
    }
}

Jarファイルを作成する処理です。

import java.io.File
import java.io.FileOutputStream
import java.io.IOException
import java.util.jar.JarEntry
import java.util.jar.JarOutputStream
import java.util.jar.Manifest
import kotlin.jvm.Throws

object JarUtil {

    @Throws(IOException::class)
    fun createJarFile(
        prefix: String,
        tmpDir: File,
        classes: Map<String, ByteArray>,
        manifest: Manifest = Manifest()
    ): File {
        val file = File.createTempFile(prefix, ".jar", tmpDir)

        val outStream = JarOutputStream(FileOutputStream(file), manifest)
        writeFilesToJarStream(classes, outStream)
        return file
    }

    @Throws(IOException::class)
    private fun writeFilesToJarStream(
        classes: Map<String, ByteArray>,
        outStream: JarOutputStream
    ) {
        val resources: MutableMap<String, ByteArray> = HashMap()
        for ((key, value) in classes) {
            resources[key.replace('.', '/') + ".class"] = value
        }
        try {
            addJarEntries(outStream, resources)
        } finally {
            try {
                outStream.close()
            } catch (ioe: IOException) {
            }
        }
    }

    @Throws(IOException::class)
    private fun addJarEntries(
        jarStream: JarOutputStream,
        files: Map<String, ByteArray>
    ) {
        for ((key, value) in files) {
            val jarEntry = JarEntry(key)
            jarStream.putNextEntry(jarEntry)
            jarStream.write(value)
            jarStream.closeEntry()
        }
    }
}

差し替え処理を実装する

これで実行時に好きに処理を差し替えられます。今回は3つだけしか関数を定義していませんが、定義していない関数にアクセスした場合、 NoSuchMethodError になります。

// クラスの一覧を取得するために以下のライブラリを使用したので、dependenciesに追加してください
// implementation("com.google.guava:guava:30.0-jre")
import com.google.common.reflect.ClassPath
import java.io.File
import java.lang.instrument.Instrumentation
import java.util.jar.JarFile

@JvmStatic
fun premain(agentArgs: String?, instrumentation: Instrumentation) {
    val agentClassLoaderUtil = ClassfileUtil.getAgentClassLoader()
    val classLoaderUtil = ClassfileUtil.getClassLoader()

    val jar = JarUtil.createJarFile(
        prefix = "NewRelicClass",
        tmpDir = tempDir,
        classes = mutableMapOf<String, ByteArray>().also { map ->
            // Agentの自前で定義したNewRelicクラスを読み込みます。
            map.putAll(
                listOf(
                    "com.newrelic.api.agent.NewRelic"
                )
                    .map { it to agentClassLoaderUtil.getByteArray(it) }
                    .toMap()
            )

            // その他のライブラリのクラスファイルを読み込みます。
            // 無いと以下のようなエラーが出る
            // Exception in thread "main" java.lang.NoClassDefFoundError: com/newrelic/api/agent/NoOpAgent
            map.putAll(
                ClassPath.from(Premain::class.java.classLoader).allClasses
                    .filter { it.name.startsWith("com.newrelic.api.agent.") }
                    .filterNot { it.name == "com.newrelic.api.agent.NewRelic" }
                    .map { it.name }
                    .map { it to classLoaderUtil.getByteArray(it) }.toMap()
            )
        }
    )

    instrumentation.appendToBootstrapClassLoaderSearch(JarFile(jar))
}

private val tempDir: File by lazy {
    val file = File("newrelic_tmp")
    if (file.exists()) {
        file.deleteRecursively()
    }
    file.mkdirs()
    file
}

後はビルドして使用するだけです。

java -javaagent:agent.jar -jar main.jar

注意

クラスを文字列で指定してます。 NewRelic::class.java のように指定すると、参照が壊れる可能性があります。(このへんはよくわかっていない)
上記を行なった時点でClassLoaderにアクセスしているはずなので、まぁ壊れる可能性はあるだろうなとは思いますが。

おわりに

興味本位でJava Agentによる差し替え処理を実装してみましたが、ライブラリ作成者しか使用しないような機能ですね。アプリケーションを作成するのに差し替えたいならラップしたりするのが絶対に良いです。
クラスを差し替えられるという強力な機能ですが、実行時に差し替えているので、間違えれば NoClassDefFoundError NoSuchMethodError になったりしますので。

今回作ったNewRelicをLog4jに差し替えるやつ。
https://github.com/matsudamper/newrelic_log4j_agent