Work with me to implement compile-time bytecode weaving using ASM

Hello, I’m N0tExpectErr0r, an Android developer who loves technology

My personal blog: blog.n0tExpecterr0r.cn

Demo address: github.com/N0tExpectEr…

This paper has given the public, “Guo Lin” release: mp.weixin.qq.com/s/aBYA1mwUa…

The cause of

I found Hugo by Jake Wharton very interesting. Through this library, I can record some data related to method calls. For example, it can print the method’s input parameters, elapsed time, return value, and so on in Logcat when the method executes by preloading the method with DebugLog annotations.

For example, add a simple comment to your code like this:

@DebugLog
public String getName(String first, String last) {
  SystemClock.sleep(15);
  return first + "" + last;
}
Copy the code

You can print the following log in Logcat:

V/Example: ⇢ getName (first ="Jake", last="Wharton"V/Example: ⇠ getName [16ms] ="Jake Wharton"
Copy the code

The library is very interesting in that it is easy to print debugging information in the form of annotations, which is much less intrusive than directly modifying the code implementation. After consulting materials, WE learned that Hugo was implemented based on AspectJ, and its core principle is the staking of bytecode at compile time. Just two days ago, the author realized the non-trace burying point of View click events through ASM bytecode peg in the project. Therefore, I suddenly thought of implementing a library with similar functions through ASM.

However, Hugo only provided the function of printing methods to execute relevant information, so he began to think about whether it could be based on its ideas to extend the function of executing specified logic before and after method calls.

If such a library could be implemented, then for Hugo’s functionality we would simply record the time before the method call and calculate the time difference after the method call.

If you also need a function to count the number of times a method is called in your application, you only need to perform the logic of counting when the method is called.

The benefit of this implementation is that it is easy to extend, listening before and after method calls, and the specific execution logic can be determined by the consumer. If you are interested in the implementation of this feature, please follow me along.

The basic principle of

First, we need to understand what ASM is. ASM is a Java bytecode level code analysis and modification tool. It has a very easy to use API, through which you can manipulate existing class files to dynamically generate classes, or extend functions based on existing classes.

At this time, there may be readers will ask, ASM is to manipulate the class file, but Apk is not the dex file? Isn’t there any way to apply it to Android?

In fact, during the compilation process of Android, Java files will be compiled into a class file, and then the compiled class file will be packaged into a dex file. We can use the gap before the class is packaged into a dex file. Insert ASM related logic to manipulate class files.

The idea is simple, but how do we execute our ASM-related code before the class file is packaged?

Since Gradle 1.5.0, Google has provided an API called Transform, which allows third-party Gradle plugins to manipulate class files before packaging dex. This time we will implement such a Gradle Plugin using the Transform API.

Implementation approach

With the fundamentals mentioned above, let’s think about a concrete implementation.

The idea is very simple. It’s a classic observer model. Our user subscribes to a method invocation event, notifies the user when the method is invoked, and executes the specified logic.

We need a dispatch center for method invocation events to which subscribers can subscribe for invocation events of a certain type of method, which is notified whenever an invocation event occurs for a method with a specified annotation, and which notifies subscribers of the corresponding type.

In this case, we only need to use ASM to weave in the code to notify the dispatch center before and after the method is called.

Show me the code

With that in mind, we were ready to start coding, and here I set up a project called Elapse. (Don’t ask why, just because it looks good)

The preparatory work

Let’s do some preparatory work — create a Module for the ASM plugin, clear out the automatically generated Gradle code, and write Gradle as follows:

apply plugin: 'groovy'

dependencies {
    implementation gradleApi()
    implementation localGroovy()

    implementation 'com. Android. The tools: gradle: 3.1.2'
}

repositories {
    mavenCentral()
    jcenter()
    google()
}
Copy the code

We also need a note to indicate the method that needs to be staked. We use the following compile-time annotation, which has a tag parameter representing the tag of the method, through which we can implement different processing for different methods.

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.CLASS)
public @interface TrackMethod {
    String tag();
}
Copy the code

Then we create a MethodEventManager to register and distribute method call events:

public class MethodEventManager {

    private static volatile MethodEventManager INSTANCE;
    private Map<String, List<MethodObserver>> mObserverMap = new HashMap<>();

    private MethodEventManager() {
    }

    public static MethodEventManager getInstance() {
        if (INSTANCE == null) {
            synchronized (MethodEventManager.class) {
                if(INSTANCE == null) { INSTANCE = new MethodEventManager(); }}}return INSTANCE;
    }

    public void registerMethodObserver(String tag, MethodObserver listener) {
        if (listener == null) {
            return;
        }

        List<MethodObserver> listeners = mObserverMap.get(tag);
        if (listeners == null) {
            listeners = new ArrayList<>();
        }
        listeners.add(listener);
        mObserverMap.put(tag, listeners);
    }

    public void notifyMethodEnter(String tag, String methodName) {
        List<MethodObserver> listeners = mObserverMap.get(tag);
        if (listeners == null) {
            return;
        }
        for (MethodObserver listener : listeners) {
            listener.onMethodEnter(tag, methodName);
        }
    }

    public void notifyMethodExit(String tag, String methodName) {
        List<MethodObserver> listeners = mObserverMap.get(tag);
        if (listeners == null) {
            return;
        }
        for(MethodObserver listener : listeners) { listener.onMethodExit(tag, methodName); }}}Copy the code

The code here is not very complicated and exposes three main methods:

registerMethodObserver: Used to register a listener for a TAG
notifyMethodEnter: is used to notify listeners of the corresponding TAG of the method call
notifyMethodExit: used to notify listeners of the corresponding TAG of the method exit

With such a class, we just need to weave the corresponding code at the beginning and end of the annotated method at code editing time, as follows:

public void method(String param) { MethodEventManager.getInstance().notifyMethodEnter(tag, methodName); . / / the original code MethodEventManager getInstance () notifyMethodExit (tag, methodName); }Copy the code

The Transform of the writing

Then we create a class ElapseTransform that inherits from Transform:

public class ElapseTransform extends Transform {

    @Override
    public String getName() {
        return ElapseTransform.class.getSimpleName();
    }

    @Override
    public Set<QualifiedContent.ContentType> getInputTypes() {
        return TransformManager.CONTENT_CLASS;
    }

    @Override
    public Set<? super QualifiedContent.Scope> getScopes() {
        return TransformManager.SCOPE_FULL_PROJECT;
    }

    @Override
    public boolean isIncremental() {
        return false; } @Override public void transform(TransformInvocation transformInvocation) throws TransformException, InterruptedException, IOException { super.transform(transformInvocation); / /... Find the class file and deal with it}}Copy the code

Here we need to implement four methods, we will introduce them respectively:

getName: The name of the current Transform
getInputTypes: Transform Specifies the data type to be processedContentTypeSet, whereContentTypeThe following values are available:
- DefaultContentType.CLASSES: Compiled bytecode files (JAR packages or directories) to work with
- DefaultContentType.RESOURCES: To work with standard Java resources
getScopesThe scope of the: Transform is aScopeSet, whereScopeThe values are as follows:
- PROJECT: Only the current project is processed
- SUB_PROJECTS: Only subprojects are processed
- PROJECT_LOCAL_DEPS: Handles only local dependencies of the current project, such as JAR, AAR
- EXTERNAL_LIBRARIES: Only external dependent libraries are handled
- PROVIDED_ONLY: Handles only dependent libraries imported locally or remotely as provided
- TESTED_CODE: Only test code is handled
isIncremental: Indicates whether incremental compilation is supported

Here we specify transformmanager.content_class to process the compiled bytecode file, and transformManager.scope_full_project to process the entire project. They are all sets preconfigured by TransformManager.

When the Transform is called, its Transform method is called, where we can find the class file and process the class file:

@Override public void transform(TransformInvocation transformInvocation) throws TransformException, InterruptedException, IOException { super.transform(transformInvocation); / / get input (consumer input, need to be passed on to the next Transform) Collection < TransformInput > inputs = transformInvocation. GetInputs ();forTransformInput input: inputs) {// Check jar and directory for inputsfor(JarInput JarInput: input.getJarInputs()) {transformJar(JarInput); }for(DirectoryInput DirectoryInput: input. GetDirectoryInputs ()) {transformDirectory(DirectoryInput); }}}Copy the code

Here I pass transformInvocation. Capturing the getInputs input, this input is a consumer, need to be passed on to the next Transform, which contains the jar file and directory file.

Check inputs, obtain the JAR list and directory list, and then pass through them. Invoke the transformJar and transformDirectory methods on the JAR file and directory respectively.

Class file search

jar

For jar files, we need to iterate through the JarEntry, find the class file, write a new temporary JAR file after modifying the class file, and then copy it to the output path after editing it.

private void transformJar(TransformInvocation invocation, JarInput input) throws IOException {
    File tempDir = invocation.getContext().getTemporaryDir();
    String destName = input.getFile().getName();
    String hexName = DigestUtils.md5Hex(input.getFile().getAbsolutePath()).substring(0, 8);
    if (destName.endsWith(".jar")) { destName = destName.substring(0, destName.length() - 4); } / / the output path File dest = invocation. GetOutputProvider () getContentLocation (destName +"_" + hexName, input.getContentTypes(), input.getScopes(), Format.JAR);
    JarFile originJar = new JarFile(input.getFile());
    File outputJar = new File(tempDir, "temp_"+input.getFile().getName()); JarOutputStream output = new JarOutputStream(new FileOutputStream(outputJar)); Enumeration<JarEntry> Enumeration = Originjar.entries ();while (enumeration.hasMoreElements()) {
        JarEntry originEntry = enumeration.nextElement();
        InputStream inputStream = originJar.getInputStream(originEntry);
        String entryName = originEntry.getName();
        if (entryName.endsWith(".class")) {
            JarEntry destEntry = new JarEntry(entryName);
            output.putNextEntry(destEntry);
            byte[] sourceBytes = IOUtils.toByteArray(inputStream); Byte [] modifiedBytes = modifyClass(sourceBytes);
            if (modifiedBytes == null) {
                modifiedBytes = sourceBytes; } output.write(modifiedBytes); output.closeEntry(); } } output.close(); originJar.close(); // copy the modified jar to the output path fileutils.copyfile (outputJar, dest); }Copy the code

As you can see, there are mainly the following steps:

throughgetContentLocationMethod gets the output path,
A temporary output JAR file is built
Iterate through the entry of the original JAR file and call the class file in itmodifyClassMake the changes, and then put in the temporary JAR file
Copy the temporary JAR file to the output path.

This changes all the class files in the JAR file.

Code is woven through ASM

This brings us to the key point where we need to use ASM to modify the specified class. The actual logic for class processing is in the modifyClass method.

private byte[] modifyClass(byte[] classBytes) {
    ClassReader classReader = new ClassReader(classBytes);
    ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_MAXS);
    ClassVisitor classVisitor = new ElapseClassVisitor(classWriter);
    classReader.accept(classVisitor, ClassReader.EXPAND_FRAMES);
    return classWriter.toByteArray();
}
Copy the code

We first need to use the ASM ClassReader to parse some of the information contained in our class files.

Then we need a ClassWriter class that can write bytecode to the class file.

After that, we define an ElapseClassVisitor that uses the previous custom ClassVisitor to “call” the class file via the classReader.accept method, and in the calling, We can insert some logic to edit the class file.

ClassWriter is also an implementation class for ClassVisitor, and we only proxy ClassWriter through ElapseClassVisitor.

And since we’re basically weaving in the method code, we don’t have to do much in that ClassVisitor, just when the visitMethod method is called, when the method is called, Return our own ElapseMethodVisitor to weave in the method:

Here ElapseMethodVisitor is not actually a subclass of MethodVisitor, but a subclass of AdviceAdapter that ASM provides that inherits from MethodVisitor, It allows you to insert your own code at the beginning, end, and so on.

class ElapseMethodVisitor extends AdviceAdapter { private final MethodVisitor methodVisitor; private final String methodName; / /... public ElapseMethodVisitor(MethodVisitor methodVisitor, int access, String name, String desc) { super(Opcodes.ASM6, methodVisitor, access, name, desc); this.methodVisitor = methodVisitor; this.methodName = name; } / /... Other code}Copy the code

Here we save the methodVisitor, which is used to weave code into the class file later, and the methodName, which is passed to the MethodEventManager for notification later.

Annotation processing

Next, we can override the visitAnnotation method to process the annotation of a method when it is accessed to determine whether the method needs to be woven in and to get the tag in the annotation.

private static final String ANNOTATION_TRACK_METHOD = "Lcom/n0texpecterr0r/elapse/TrackMethod;";
private boolean needInject;
private String tag;

@Override
public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
    AnnotationVisitor annotationVisitor = super.visitAnnotation(desc, visible);
    if (desc.equals(ANNOTATION_TRACK_METHOD)) {
        needInject = true;
        return new AnnotationVisitor(Opcodes.ASM6, annotationVisitor) {
            @Override
            public void visit(String name, Object value) {
                super.visit(name, value);
                if (name.equals("tag") && value instanceof String) { tag = (String) value; }}}; }return annotationVisitor;
}
Copy the code

Here we first determine whether the signature of the annotation is the same as the annotation TrackMethod we need (the specific signature rules are not introduced here, you can search by yourself, it is actually the method signature set, pay attention to the semicolon inside).

If the annotation is the one we need, set needInject to true and get the value of tag from the annotation,

In this way, we later just need to check whether needInject is needed to know which methods need to be woven in.

Code weaving

OnMethodEnter and onMethodExit can be used to monitor method entry and exit by overriding onMethodEnter and onMethodExit:

@Override
protected void onMethodEnter() {
    super.onMethodEnter();
    handleMethodEnter();
}

@Override
protected void onMethodExit(int opcode) {
    super.onMethodExit(opcode);
    handleMethodExit();
}
Copy the code

The two pieces of code are very similar, except that the method name at the end is different, so I’ll just use handleMethodEnter as an example.

In the ASM, through MethodWriter. VisitMethodInsn method can call the similar bytecode instruction to invoke the method. Such as

VisitMethodInsn (INVOKESTATIC, class signature, method name, method signature);

This way you can call a static method on a class. If the method requires parameters, we can call directives such as ALOAD to push variables to the stack via the visitVarInsn method. The whole process is actually similar to the form of calls in bytecode.

It would be fine if we just called a static method, but here we need to call a specific method under a singleton class, such as

MethodEventManager.getInstance().notifyMethodEnter(tag, methodName);

It’s probably hard for anyone other than those familiar with bytecodes to think directly about how to represent such code in bytecode. We can solve the problem in the following two ways:

1. View the bytecode in javAP

So we can write a singleton call Demo and then use Javap -v to look at the generated bytecode to get an idea of the order in which the bytecode is called:

As you can see, the getInstance method is called by INVOKESTATIC, the two string constants are placed at the top of the stack by LDC, and the notify method is called by INVOKEVIRTUAL.

We can mimic this process by calling the corresponding method in ASM to complete a similar process, so we write the following code, where visitLdcInsn is similar to LDC in bytecode.

private void handleMethodEnter() {
    if(needInject && tag ! = null) { methodVisitor.visitMethodInsn(INVOKESTATIC, METHOD_EVENT_MANAGER,"getInstance"."()L"+METHOD_EVENT_MANAGER+";");
        methodVisitor.visitLdcInsn(tag);
        methodVisitor.visitLdcInsn(methodName);
        methodVisitor.visitMethodInsn(INVOKEVIRTUAL, METHOD_EVENT_MANAGER, 
                "notifyMethodEnter"."(Ljava/lang/String; Ljava/lang/String;) V"); }}Copy the code

This way, we can weave in the code we want.

2. Use the ASM Bytecode plug-in to view information

There is another way to simplify the process of Bytecode viewing. There is an IDEA plugin named “ASM Bytecode Outline” written by The oracle, which can be used to directly view the corresponding ASM code.

After installing the plug-in, right-click the code you want to view ->Show ByteCode to view the corresponding ASM code. The result is as follows:

Both methods have their advantages and disadvantages, and readers can use different ways to achieve them according to their own needs.

Through the previous series of steps, the core function of ASM knitting has been implemented. If you still need to obtain the parameters of the function and other extensions, just need to know the corresponding bytecode implementation, the rest is easy to implement, I won’t go into details due to the lack of space.

Package as a Gradle plug-in

For the final step, we package the library as a Gradle Plugin. We create a new ElapsePlugin class that inherits from Plugin and register our ElapseTransform in it.

public class ElapsePlugin implements Plugin<Project> {
    @Override
    public void apply(@NotNull Project project) {
        AppExtension appExtension = project.getExtensions().findByType(AppExtension.class);
        assert appExtension != null;
        appExtension.registerTransform(new ElapseTransform(project));
    }
}
Copy the code

Then we add the following gradle code to build.gradle to describe our POM:

apply plugin: 'maven'

uploadArchives {
    repositories.mavenDeployer {
        repository(url: uri('.. /repo'))
        pom.groupId = 'com.n0texpecterr0r.build'
        pom.artifactId = 'elapse-asm'
        pom.version = '1.0.0'}}Copy the code

Create a < plugin name >.properties file in the resources/ meta-INF /gradle-plugins folder SRC /main.

Fill in the file as follows:

Implementation – the class = < Plugin directory >, for example, I here is implementation – class = com. N0texpecterr0r. Elapseasm. ElapsePlugin

This way, we can run uploadArchives Gradle script to generate the corresponding JAR package. At this point, our Gradle Plugin for function calls staking is complete.

Results show

We can add it to the classpath as needed:

repositories {
    //...
    maven {
        url uri("repo")
    }
}

dependencies {
    // ...
    classpath 'com. N0texpecterr0r. Build: elapse - asm: 1.0.0'
}
Copy the code

Then apply it under app Module:

apply plugin: 'com.n0texpecterr0r.elapse-asm'
Copy the code

We can write a Demo to test the effect:

public class MainActivity extends AppCompatActivity {

    private static final String TAG_TEST = "test";

    @Override
    protected void onCreate(Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);
        setContentView(R.layout.activity_main);
        MethodEventManager.getInstance().registerMethodObserver(TAG_TEST, new MethodObserver() {
            @Override
            public void onMethodEnter(String tag, String methodName) {
                Log.d("MethodEvent"."method "+ methodName + " enter at time " + System.currentTimeMillis());
            }

            @Override
            public void onMethodExit(String tag, String methodName) {
                Log.d("MethodEvent"."method "+ methodName + " exit at time "+ System.currentTimeMillis()); }});test(a); } @TrackMethod(tag = TAG_TEST) public voidtest() { try { Thread.sleep(1200); } catch (InterruptedException e) { e.printStackTrace(); }}}Copy the code

Run the program, we can see that Logcat successfully printed the information we need:

In other words, our code has been successfully incorporated into bytecode. Let’s take a look at the compiled bytecode generated, we can open the elapse – demo/build/intermediates/transforms/ElapseTransform/debug / 33 / MainActivitiy. Class:

As you can see, our code has been successfully inserted into the bytecode.

Implement Hugo

We can then use it to try to implement Hugo’s printing method time function by creating a new TimeObserver:

public class TimeObserver implements MethodObserver {
    private static final String TAG_METHOD_TIME = "MethodCost";
    private Map<String, Long> enterTimeMap = new HashMap<>();
    @Override
    public void onMethodEnter(String tag, String methodName) {
        String key = generateKey(tag, methodName);
        Long time = System.currentTimeMillis();
        enterTimeMap.put(key, time);
    }
    @Override
    public void onMethodExit(String tag, String methodName) {
        String key = generateKey(tag, methodName);
        Long enterTime = enterTimeMap.get(key);
        if (enterTime == null) {
            throw new IllegalStateException("method exit without enter");
        }
        long cost = System.currentTimeMillis() - enterTime;
        Log.d(TAG_METHOD_TIME, "method " + methodName + " cost "
                + (double)cost/1000 + "s" + " in thread " + Thread.currentThread().getName());
		enterTimeMap.remove(key);
    }
    private String generateKey(String tag, String methodName) {
        returntag + methodName + Thread.currentThread().getName(); }}Copy the code

Here we use Tag + methodName + CurrentThread. name as key to avoid interference caused by multi-threaded calls. The start time is recorded when the method enters and the time difference is calculated when the method exits.

After we register it in the Application, we can see the effect after it runs:

Let’s open 10 threads to run test separately and see what happens:

private ExecutorService mExecutor = Executors.newFixedThreadPool(10);

@Override
protected void onCreate(Bundle savedInstanceState) {
    super.onCreate(savedInstanceState);
    setContentView(R.layout.activity_main);
    for (int i = 0; i < 10; i++) {
        mExecutor.execute(this::test); }}Copy the code

As you can see, it is still possible to count the method call time normally:

conclusion

The ASM + Transform API makes it easy to edit the bytecode before the class is packaged as a dex file to insert the logic you need anywhere in the code. This is just a small Demo to give you an idea of how powerful ASM can be. What you can do with ASM is actually much richer. At present in the domestic ASM related articles are still relatively scarce, if you want to further understand the functions of ASM, readers can go here to check ASM official documents.

Actually this Demo has more functions can be expanded, such as function arguments and return values of information carrying, inserted on the whole class of methods for pile, and so on, the reader can according to the prior knowledge, try to extend the function, due to the limited space is no longer here, here are essentially inserted into the corresponding bytecode instructions.