1 /*
2  * Copyright (C) 2022 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.traceinjection;
18 
19 import org.objectweb.asm.AnnotationVisitor;
20 import org.objectweb.asm.Label;
21 import org.objectweb.asm.MethodVisitor;
22 import org.objectweb.asm.Opcodes;
23 import org.objectweb.asm.Type;
24 import org.objectweb.asm.commons.AdviceAdapter;
25 import org.objectweb.asm.commons.Method;
26 
27 /**
28  * Adapter that injects tracing code to methods annotated with the configured annotation.
29  *
30  * Assuming the configured annotation is {@code @Trace} and the configured methods are
31  * {@code Tracing.begin()} and {@code Tracing.end()}, it effectively transforms:
32  *
33  * <pre>{@code
34  * @Trace
35  * void method() {
36  *     doStuff();
37  * }
38  * }</pre>
39  *
40  * into:
41  * <pre>{@code
42  * @Trace
43  * void method() {
44  *     Tracing.begin();
45  *     try {
46  *         doStuff();
47  *     } finally {
48  *         Tracing.end();
49  *     }
50  * }
51  * }</pre>
52  */
53 public class TraceInjectionMethodAdapter extends AdviceAdapter {
54     private final TraceInjectionConfiguration mParams;
55     private final Label mStartFinally = newLabel();
56     private final boolean mIsConstructor;
57 
58     private boolean mShouldTrace;
59     private long mTraceId;
60     private String mTraceLabel;
61 
TraceInjectionMethodAdapter(MethodVisitor methodVisitor, int access, String name, String descriptor, TraceInjectionConfiguration params)62     public TraceInjectionMethodAdapter(MethodVisitor methodVisitor, int access,
63             String name, String descriptor, TraceInjectionConfiguration params) {
64         super(Opcodes.ASM9, methodVisitor, access, name, descriptor);
65         mParams = params;
66         mIsConstructor = "<init>".equals(name);
67     }
68 
69     @Override
visitCode()70     public void visitCode() {
71         super.visitCode();
72         if (mShouldTrace) {
73             visitLabel(mStartFinally);
74         }
75     }
76 
77     @Override
onMethodEnter()78     protected void onMethodEnter() {
79         if (!mShouldTrace) {
80             return;
81         }
82         Type type = Type.getType(toJavaSpecifier(mParams.startMethodClass));
83         Method trace = Method.getMethod("void " + mParams.startMethodName + " (long, String)");
84         push(mTraceId);
85         push(getTraceLabel());
86         invokeStatic(type, trace);
87     }
88 
getTraceLabel()89     private String getTraceLabel() {
90         return !isEmpty(mTraceLabel) ? mTraceLabel : getName();
91     }
92 
93     @Override
onMethodExit(int opCode)94     protected void onMethodExit(int opCode) {
95         // Any ATHROW exits will be caught as part of our exception-handling block, so putting it
96         // here would cause us to call the end trace method multiple times.
97         if (opCode != ATHROW) {
98             onFinally();
99         }
100     }
101 
onFinally()102     private void onFinally() {
103         if (!mShouldTrace) {
104             return;
105         }
106         Type type = Type.getType(toJavaSpecifier(mParams.endMethodClass));
107         Method trace = Method.getMethod("void " + mParams.endMethodName + " (long)");
108         push(mTraceId);
109         invokeStatic(type, trace);
110     }
111 
112     @Override
visitMaxs(int maxStack, int maxLocals)113     public void visitMaxs(int maxStack, int maxLocals) {
114         final int minStackSize;
115         if (mShouldTrace) {
116             Label endFinally = newLabel();
117             visitLabel(endFinally);
118             catchException(mStartFinally, endFinally, null);
119             // The stack will always contain exactly one element: the exception we caught
120             final Object[] stack = new Object[]{ "java/lang/Throwable"};
121             // Because we use EXPAND_FRAMES, the frame type must always be F_NEW.
122             visitFrame(F_NEW, /* numLocal= */ 0, /* local= */ null, stack.length, stack);
123             onFinally();
124             // Rethrow the exception that we caught in the finally block.
125             throwException();
126 
127             // Make sure we have at least enough stack space to push the trace arguments
128             // (long, String)
129             minStackSize = Type.LONG_TYPE.getSize() + Type.getType(String.class).getSize();
130         } else {
131             // We didn't inject anything, so no need for additional stack space.
132             minStackSize = 0;
133         }
134 
135         super.visitMaxs(Math.max(minStackSize, maxStack), maxLocals);
136     }
137 
138     @Override
visitAnnotation(String descriptor, boolean visible)139     public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
140         AnnotationVisitor av = super.visitAnnotation(descriptor, visible);
141         if (descriptor.equals(toJavaSpecifier(mParams.annotation))) {
142             if (mIsConstructor) {
143                 // TODO: Support constructor tracing. At the moment, constructors aren't supported
144                 //  because you can't put an exception handler around a super() call within the
145                 //  constructor itself.
146                 throw new IllegalStateException("Cannot trace constructors");
147             }
148             av = new TracingAnnotationVisitor(av);
149         }
150         return av;
151     }
152 
153     /**
154      * An AnnotationVisitor that pulls the trace ID and label information from the configured
155      * annotation.
156      */
157     class TracingAnnotationVisitor extends AnnotationVisitor {
158 
TracingAnnotationVisitor(AnnotationVisitor annotationVisitor)159         TracingAnnotationVisitor(AnnotationVisitor annotationVisitor) {
160             super(Opcodes.ASM9, annotationVisitor);
161         }
162 
163         @Override
visit(String name, Object value)164         public void visit(String name, Object value) {
165             if ("tag".equals(name)) {
166                 mTraceId = (long) value;
167                 // If we have a trace annotation and ID, then we have everything we need to trace
168                 mShouldTrace = true;
169             } else if ("label".equals(name)) {
170                 mTraceLabel = (String) value;
171             }
172             super.visit(name, value);
173         }
174     }
175 
toJavaSpecifier(String klass)176     private static String toJavaSpecifier(String klass) {
177         return "L" + klass + ";";
178     }
179 
isEmpty(String str)180     private static boolean isEmpty(String str) {
181         return str == null || "".equals(str);
182     }
183 }
184