Javassitを使ってクラスをまるごと入れ替えてみた

こんにちは。インテグレーションテック部の満石です。

以前、Javaのバイトコードを操作してEnumを書き換えるで、
Javassistを使ってEnumを書き換える方法を書きましたが、
今回はJavassitを使ってクラスをまるごと入れ替える方法について紹介します。

背景

モノリシックなアプリケーションをマイクロサービスに分割するプロジェクトで、
単体テストや結合テストで、マイクロサービス化した処理は
モックが動作するようにしてほしいという要請がありました。

その対応には、以下の制約がありました。

既存のテストコードをモックライブラリを使用したコードに変更するのはコストがかかりすぎるし、結合テストでも使うのでモックライブラリは使えない
できればマイクロサービス化前の元のクラスをモックの代わりに使いたい
この対応のためにプロダクションコードには手を入れないで欲しい

以前Javassistを使った経験があり、今回もこの対応にはJavassistでうまく対応できそうだと思いました。

実行環境

サンプルコードの説明に入る前に、まずは今回のプログラムの実行環境について説明しておきます。
Javaのバージョンは1.8.0_181、Javassistは3.12.1.GAです。

サンプルコード

上書き対象のクラス

まず、上書き対象のクラスを説明します。
今回の説明に最低限必要なシンプルなコードにしています。

com.example.phrase.Worldクラス

Hello World!のWorld!部分を返却するだけのクラスです。

package com.example.phrase;

public class World {
    public String getPhrase() {
        return "World!";
    }
}

com.example.Helloクラス

Hello World!を出力するだけのクラスです。
Worldクラスとはあえてパッケージを別にしてimport文が必要になるようにしています。

package com.example;

import com.example.phrase.World;

public class Hello {
    public static void say(World world) {
        System.out.println("Hello " + world.getPhrase());
    }
}

Mainクラス

mainメソッドを実装したクラスで、このクラスだけは今回の上書き対象のクラスではありません。
上記HelloクラスのsayメソッドにWorldクラスのインスタンスを渡しています。

import com.example.Hello;
import com.example.phrase.World;

public class Main {
    public static void main(String[] args) {
        Hello.say(new World());
    }
}

モック側のクラス

モック側のクラスは、上書き対象のクラスと名前は同じですが、
間にmockパッケージを入れて上書き対象のクラスとはパッケージが異なるようにしています。

モッククラスで上書きされた場合、Hello World!が
ようこそ、世界！に出力が変わるように実装しています。

それ以外は、上書き対象のクラスと同じです。

com.example.mock.phrase.Worldクラス

package com.example.mock.phrase;

public class World {
    public String getPhrase() {
        return "世界！";
    }
}

com.example.mock.Hello

package com.example.mock;

import com.example.mock.phrase.World;

public class Hello {
    public static void say(World world) {
        System.out.println("ようこそ、" + world.getPhrase());
    }
}

上記のクラスたちをビルドして、jarファイルを作成しました。
作成されたjarファイルの名前は、hello-world-1.0-SNAPSHOT.jarとします。

クラスをまるごと入れ替える処理

ここからが今回の本題となる、クラスをまるごと入れ替える処理の実装です。

com.example.ClassFileOverrideクラス

このクラスではpremainメソッドを実装しています。

premainメソッドでは、
java.lang.instrument.InstrumentationのaddTransformerメソッドの引数に、
java.lang.instrumentインターフェースの実装をLambda式で実装しています。

これらについては、前書いた
Javaのバイトコードを操作してEnumを書き換える
で詳細に説明しているため、ぜひご覧ください。

他、個別の実装の説明については、詳細なコメントを以下のプログラム中に書いています。

package com.example;

import javassist.*;

import java.io.IOException;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.util.HashMap;
import java.util.Map;

public class ClassFileOverride {

    // どのクラスをどのモッククラスで上書きするかをMapで作成
    // サンプルコードのためMapで作成したが、実際はpropertiesファイル等、外部ファイルから読み込むことを推奨
    private static Map<String, String> overrideClassMap = new HashMap<>();
    static {
        overrideClassMap.put("com.example.Hello", "com.example.mock.Hello");
        overrideClassMap.put("com.example.phrase.World", "com.example.mock.phrase.World");
    }

    // 上書き対象クラスで他の上書き対象クラスが使われている場合、使われ方次第では実行時にエラーが発生することがある
    // モッククラス名を上書き対象クラス名に置き換えることで対応できるため、ここでは置き換え用のClassMapを作成する
    private static ClassMap classMap = new ClassMap();
    static {
        overrideClassMap.forEach((key, value) -> classMap.put(value, key));
    }

    public static void premain(String agentArgs, Instrumentation instrumentation) {
        instrumentation.addTransformer(
                (loader, className, classBeingRedefined, protectionDomain, classfileBuffer) -> {

                    // classNameは"/"区切りで取得されるため、"."に置き換える
                    String originalFqcn = className.replace("/", ".");

                    // 上書き対象のクラスでない場合は書き換えない
                    if (!overrideClassMap.containsKey(originalFqcn)) {
                        // バイトコードを書き換えない場合はnullをreturnする決まりになっている
                        return null;
                    }

                    ClassPool classPool = ClassPool.getDefault();
                    String replaceFqcn = overrideClassMap.get(originalFqcn);

                    try {
                        // モッククラスを取得
                        CtClass replaceClass = classPool.get(replaceFqcn);

                        // モッククラスのFQCNを上書き対象クラスのFQCNに置き換えることで、モッククラスを上書き対象のクラスとして認識させる
                        replaceClass.setName(originalFqcn);

                        // モッククラス内のクラス名を上書き対象のクラス名に置き換える
                        replaceClass.replaceClassName(classMap);

                        return replaceClass.toBytecode();
                    } catch (NotFoundException | CannotCompileException | IOException e) {
                        // 上記例外発生時にスタックトレースを出力する。これを書いておかないと例外発生時にコンソールに何も出力されない
                        e.printStackTrace();
                        IllegalClassFormatException illegalClassFormatException = new IllegalClassFormatException();
                        illegalClassFormatException.initCause(e);
                        // 例外はIllegalClassFormatExceptionをスローするよう決められているが、挙動はnullをreturnした場合と同じで、コンソールには何も出力されない
                        throw illegalClassFormatException;
                    }
                });
    }
}

ビルドと実行

このクラスも別途ビルドして、先ほどとは別のjarファイルを作成する必要がありますが
上書き対象クラスとモッククラスが入ったjarを作成する場合と違い、ビルドに気をつける点があります。

今回はgradleでビルドしたため、build.gradleファイルを説明します。

plugins {
    id 'java'
}

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

sourceCompatibility = 1.8

repositories {
    mavenCentral()
}

dependencies {
    compile group: 'javassist', name: 'javassist', version: '3.12.1.GA'
}

jar {
    manifest {
        attributes 'Premain-Class': 'com.example.ClassFileOverride', 
                   'Boot-Class-Path': '<javassitのjarファイルを配置したPath>/javassist-3.12.1.GA.jar', 
                   'Can-Redefine-Classes': 'true'
    }
}

まず、依存関係としてJavassitを指定したのが以下の部分です。

dependencies {
    compile group: 'javassist', name: 'javassist', version: '3.12.1.GA'
}

さらに、jarファイル内に決められたマニフェスト属性を記述する必要があり

そのための記述が以下の部分です。

jar {
    manifest {
        attributes 'Premain-Class': 'com.example.ClassFileOverride', 
                   'Boot-Class-Path': '<javassitのjarファイルを配置したPath>/javassist-3.12.1.GA.jar', 
                   'Can-Redefine-Classes': 'true'
    }
}

Premain-Classにはpremainメソッドを実装したクラスを指定します。
Boot-Class-Pathにはjavassisistのjarファイルの場所を指定します。

Can-Redefine-Classesバイトコードを操作してクラスの再定義を許すかどうかを指定するため、
今回はtrueを指定します。

これら含めて詳細はjava.lang.instrumentパッケージのJavaDocをご覧ください。

作成されたjarファイルの名前はclass-file-overrider-1.0-SNAPSHOT.jarとします。

実行時は
-javaagent:＜class-file-overrider-1.0-SNAPSHOT.jarを配置したPath＞/class-file-overrider-1.0-SNAPSHOT.jar
を実行時のオプションとして付与して、実行する必要があります。

このオプションを付与せず

java -jar hello-world-1.0-SNAPSHOT.jar

と実行した場合は、当然ながら
Hello World!
が表示されます。

このオプションを付与して、

java -javaagent:＜class-file-overrider-1.0-SNAPSHOT.jarを配置したPath＞/class-file-overrider-1.0-SNAPSHOT.jar -jar hello-world-1.0-SNAPSHOT.jar

と実行した場合は
ようこそ、世界！
が表示されて、意図したとおりクラスが上書きされたことが確認できました。

今回の対応のキモ

今回の対応で苦慮した点は、
モッククラス内のクラス名を上書き対象のクラス名に置き換える必要があった点です。

上記コードだとreplaceClass.replaceClassName(classMap);と
その引数に渡しているclassMapを作成している箇所です。

この対応をしていない場合、今回のコードだと

Exception in thread "main" java.lang.NoSuchMethodError: com.example.Hello.say(Lcom/example/phrase/World;)V
	at Main.main(Main.java:6)

のエラーが発生してしまいます。

終わりに

今どきのフレームワークであれば、
DIを使えばこんなことをする必要がないのでは？と思われるかもしれません。

しかし、DIを利用する場合は、予めDIをする前提で実装しておく必要があります。

ところが今回の方法であれば、DIする前提が無かった処理や
DI不可能なstaticメソッドでさえも、処理を上書きすることができてしまいます。

同じ課題に直面した方に、この情報がお役に立てば幸いです。

CSVIT事業部　IT（インテグレーションテック）部満石昇

アーキテクトとして開発者の方々に技術的貢献をするのが生きがいです。

背景