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.ClassReader;
20 import org.objectweb.asm.ClassWriter;
21 
22 import java.io.BufferedInputStream;
23 import java.io.FileOutputStream;
24 import java.io.IOException;
25 import java.io.InputStream;
26 import java.io.OutputStream;
27 import java.util.Enumeration;
28 import java.util.zip.ZipEntry;
29 import java.util.zip.ZipFile;
30 import java.util.zip.ZipOutputStream;
31 
32 public class Main {
main(String[] args)33     public static void main(String[] args) throws IOException {
34         String inJar = null;
35         String outJar = null;
36         String annotation = null;
37         String traceStart = null;
38         String traceEnd = null;
39 
40         // All arguments require a value currently, so just make sure we have an even number and
41         // then process them all two at a time.
42         if (args.length % 2 != 0) {
43             throw new IllegalArgumentException("Argument is missing corresponding value");
44         }
45         for (int i = 0; i < args.length - 1; i += 2) {
46             final String arg = args[i].trim();
47             final String argValue = args[i + 1].trim();
48             if ("-i".equals(arg)) {
49                 inJar = argValue;
50             } else if ("-o".equals(arg)) {
51                 outJar = argValue;
52             } else if ("--annotation".equals(arg)) {
53                 annotation = argValue;
54             } else if ("--start".equals(arg)) {
55                 traceStart = argValue;
56             } else if ("--end".equals(arg)) {
57                 traceEnd = argValue;
58             } else {
59                 throw new IllegalArgumentException("Unknown argument: " + arg);
60             }
61         }
62 
63         if (inJar == null) {
64             throw new IllegalArgumentException("input jar is required");
65         }
66 
67         if (outJar == null) {
68             throw new IllegalArgumentException("output jar is required");
69         }
70 
71         if (annotation == null) {
72             throw new IllegalArgumentException("trace annotation is required");
73         }
74 
75         if (traceStart == null) {
76             throw new IllegalArgumentException("start trace method is required");
77         }
78 
79         if (traceEnd == null) {
80             throw new IllegalArgumentException("end trace method is required");
81         }
82 
83         TraceInjectionConfiguration params =
84                 new TraceInjectionConfiguration(annotation, traceStart, traceEnd);
85 
86         try (
87                 ZipFile zipSrc = new ZipFile(inJar);
88                 ZipOutputStream zos = new ZipOutputStream(new FileOutputStream(outJar));
89         ) {
90             Enumeration<? extends ZipEntry> srcEntries = zipSrc.entries();
91             while (srcEntries.hasMoreElements()) {
92                 ZipEntry entry = srcEntries.nextElement();
93                 ZipEntry newEntry = new ZipEntry(entry.getName());
94                 newEntry.setTime(entry.getTime());
95                 zos.putNextEntry(newEntry);
96                 BufferedInputStream bis = new BufferedInputStream(zipSrc.getInputStream(entry));
97 
98                 if (entry.getName().endsWith(".class")) {
99                     convert(bis, zos, params);
100                 } else {
101                     while (bis.available() > 0) {
102                         zos.write(bis.read());
103                     }
104                     zos.closeEntry();
105                     bis.close();
106                 }
107             }
108             zos.finish();
109         }
110     }
111 
convert(InputStream in, OutputStream out, TraceInjectionConfiguration params)112     private static void convert(InputStream in, OutputStream out,
113             TraceInjectionConfiguration params) throws IOException {
114         ClassReader cr = new ClassReader(in);
115         ClassWriter cw = new ClassWriter(0);
116         TraceInjectionClassVisitor cv = new TraceInjectionClassVisitor(cw, params);
117         cr.accept(cv, ClassReader.EXPAND_FRAMES);
118         byte[] data = cw.toByteArray();
119         out.write(data);
120     }
121 }
122