1 /*
2  * Copyright (C) 2018 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 package transformer;
17 
18 import annotations.BootstrapMethod;
19 import annotations.CalledByIndy;
20 import annotations.Constant;
21 import java.io.InputStream;
22 import java.io.OutputStream;
23 import java.lang.invoke.MethodType;
24 import java.lang.reflect.Method;
25 import java.lang.reflect.Modifier;
26 import java.net.URL;
27 import java.net.URLClassLoader;
28 import java.nio.file.Files;
29 import java.nio.file.Path;
30 import java.nio.file.Paths;
31 import java.util.HashMap;
32 import java.util.Map;
33 import org.objectweb.asm.ClassReader;
34 import org.objectweb.asm.ClassVisitor;
35 import org.objectweb.asm.ClassWriter;
36 import org.objectweb.asm.Handle;
37 import org.objectweb.asm.MethodVisitor;
38 import org.objectweb.asm.Opcodes;
39 import org.objectweb.asm.Type;
40 
41 /**
42  * Class for inserting invoke-dynamic instructions in annotated Java class files.
43  *
44  * <p>This class replaces static method invocations of annotated methods with invoke-dynamic
45  * instructions. Suppose a method is annotated as: <code>
46  *
47  * @CalledByIndy(
48  *      bootstrapMethod =
49  *          @BootstrapMethod(
50  *               enclosingType = TestLinkerMethodMinimalArguments.class,
51  *               parameterTypes = {MethodHandles.Lookup.class, String.class, MethodType.class},
52  *               name = "linkerMethod"
53  *     ),
54  *     fieldOdMethodName = "magicAdd",
55  *     returnType = int.class,
56  *     argumentTypes = {int.class, int.class}
57  * )
58  * private int add(int x, int y) {
59  *    throw new UnsupportedOperationException(e);
60  * }
61  *
62  * private int magicAdd(int x, int y) {
63  *    return x + y;
64  * }
65  *
66  * </code> Then invokestatic bytecodes targeting the add() method will be replaced invokedynamic
67  * instructions targetting the CallSite that is construction by the bootstrap method described by
68  * the @CalledByIndy annotation.
69  *
70  * <p>In the example above, this results in add() being replaced by invocations of magicAdd().
71  */
72 public class IndyTransformer {
73 
74     static class BootstrapBuilder extends ClassVisitor {
75 
76         private final Map<String, CalledByIndy> callsiteMap;
77         private final Map<String, Handle> bsmMap = new HashMap<>();
78 
BootstrapBuilder(int api, Map<String, CalledByIndy> callsiteMap)79         public BootstrapBuilder(int api, Map<String, CalledByIndy> callsiteMap) {
80             this(api, null, callsiteMap);
81         }
82 
BootstrapBuilder(int api, ClassVisitor cv, Map<String, CalledByIndy> callsiteMap)83         public BootstrapBuilder(int api, ClassVisitor cv, Map<String, CalledByIndy> callsiteMap) {
84             super(api, cv);
85             this.callsiteMap = callsiteMap;
86         }
87 
88         @Override
visitMethod( int access, String name, String desc, String signature, String[] exceptions)89         public MethodVisitor visitMethod(
90                 int access, String name, String desc, String signature, String[] exceptions) {
91             MethodVisitor mv = cv.visitMethod(access, name, desc, signature, exceptions);
92             return new MethodVisitor(this.api, mv) {
93                 @Override
94                 public void visitMethodInsn(
95                         int opcode, String owner, String name, String desc, boolean itf) {
96                     if (opcode == org.objectweb.asm.Opcodes.INVOKESTATIC) {
97                         CalledByIndy callsite = callsiteMap.get(name);
98                         if (callsite != null) {
99                             insertIndy(callsite.fieldOrMethodName(), desc, callsite);
100                             return;
101                         }
102                     }
103                     mv.visitMethodInsn(opcode, owner, name, desc, itf);
104                 }
105 
106                 private void insertIndy(String name, String desc, CalledByIndy callsite) {
107                     Handle bsm = buildBootstrapMethodHandle(callsite.bootstrapMethod()[0]);
108                     Object[] bsmArgs =
109                             buildBootstrapArguments(callsite.constantArgumentsForBootstrapMethod());
110                     mv.visitInvokeDynamicInsn(name, desc, bsm, bsmArgs);
111                 }
112 
113                 private Handle buildBootstrapMethodHandle(BootstrapMethod bootstrapMethod) {
114                     String className = Type.getInternalName(bootstrapMethod.enclosingType());
115                     String methodName = bootstrapMethod.name();
116                     String methodType =
117                             MethodType.methodType(
118                                             bootstrapMethod.returnType(),
119                                             bootstrapMethod.parameterTypes())
120                                     .toMethodDescriptorString();
121                     return new Handle(
122                             Opcodes.H_INVOKESTATIC,
123                             className,
124                             methodName,
125                             methodType,
126                             false /* itf */);
127                 }
128 
129                 private Object decodeConstant(int index, Constant constant) {
130                     if (constant.booleanValue().length == 1) {
131                         return constant.booleanValue()[0];
132                     } else if (constant.byteValue().length == 1) {
133                         return constant.byteValue()[0];
134                     } else if (constant.charValue().length == 1) {
135                         return constant.charValue()[0];
136                     } else if (constant.shortValue().length == 1) {
137                         return constant.shortValue()[0];
138                     } else if (constant.intValue().length == 1) {
139                         return constant.intValue()[0];
140                     } else if (constant.longValue().length == 1) {
141                         return constant.longValue()[0];
142                     } else if (constant.floatValue().length == 1) {
143                         return constant.floatValue()[0];
144                     } else if (constant.doubleValue().length == 1) {
145                         return constant.doubleValue()[0];
146                     } else if (constant.stringValue().length == 1) {
147                         return constant.stringValue()[0];
148                     } else if (constant.classValue().length == 1) {
149                         return Type.getType(constant.classValue()[0]);
150                     } else {
151                         throw new Error("Bad constant at index " + index);
152                     }
153                 }
154 
155                 private Object[] buildBootstrapArguments(Constant[] bootstrapMethodArguments) {
156                     Object[] args = new Object[bootstrapMethodArguments.length];
157                     for (int i = 0; i < bootstrapMethodArguments.length; ++i) {
158                         args[i] = decodeConstant(i, bootstrapMethodArguments[i]);
159                     }
160                     return args;
161                 }
162             };
163         }
164     }
165 
166     private static void transform(Path inputClassPath, Path outputClassPath) throws Throwable {
167         URL url = inputClassPath.getParent().toUri().toURL();
168         URLClassLoader classLoader =
169                 new URLClassLoader(new URL[] {url}, ClassLoader.getSystemClassLoader());
170         String inputClassName = inputClassPath.getFileName().toString().replace(".class", "");
171         Class<?> inputClass = classLoader.loadClass(inputClassName);
172         Map<String, CalledByIndy> callsiteMap = new HashMap<>();
173 
174         for (Method m : inputClass.getDeclaredMethods()) {
175             CalledByIndy calledByIndy = m.getAnnotation(CalledByIndy.class);
176             if (calledByIndy == null) {
177                 continue;
178             }
179             if (calledByIndy.fieldOrMethodName() == null) {
180                 throw new Error("CallByIndy annotation does not specify a field or method name");
181             }
182             final int PRIVATE_STATIC = Modifier.STATIC | Modifier.PRIVATE;
183             if ((m.getModifiers() & PRIVATE_STATIC) != PRIVATE_STATIC) {
184                 throw new Error(
185                         "Method whose invocations should be replaced should be private and static");
186             }
187             callsiteMap.put(m.getName(), calledByIndy);
188         }
189         ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
190         try (InputStream is = Files.newInputStream(inputClassPath)) {
191             ClassReader cr = new ClassReader(is);
192             cr.accept(new BootstrapBuilder(Opcodes.ASM6, cw, callsiteMap), 0);
193         }
194         try (OutputStream os = Files.newOutputStream(outputClassPath)) {
195             os.write(cw.toByteArray());
196         }
197     }
198 
199     public static void main(String[] args) throws Throwable {
200         transform(Paths.get(args[0]), Paths.get(args[1]));
201     }
202 }
203