|
1
8
9 package com.sleepycat.persist.model;
10
11 import java.io.File  ;
12 import java.io.FileInputStream ;
13 import java.io.FileOutputStream ;
14 import java.io.IOException ;
15 import java.lang.instrument.ClassFileTransformer ;
16 import java.lang.instrument.Instrumentation ;
17 import java.security.ProtectionDomain ;
18 import java.util.ArrayList ;
19 import java.util.HashSet ;
20 import java.util.List ;
21 import java.util.Set ;
22 import java.util.StringTokenizer ;
23
24 import com.sleepycat.asm.ClassReader;
25 import com.sleepycat.asm.ClassVisitor;
26 import com.sleepycat.asm.ClassWriter;
27
28
65 public class ClassEnhancer implements ClassFileTransformer {
66
67 private static final String AGENT_PREFIX = "enhance:";
68
69 private Set <String > packagePrefixes;
70 private boolean verbose;
71
72
87 public static void main(String [] args) throws Exception {
88 try {
89 boolean verbose = false;
90 List <File > fileList = new ArrayList <File >();
91 for (int i = 0; i < args.length; i += 1) {
92 String arg = args[i];
93 if (arg.startsWith("-")) {
94 if ("-v".equals(args[i])) {
95 verbose = true;
96 } else {
97 throw new IllegalArgumentException
98 ("Unknown arg: " + arg);
99 }
100 } else {
101 fileList.add(new File (arg));
102 }
103 }
104 ClassEnhancer enhancer = new ClassEnhancer();
105 enhancer.setVerbose(verbose);
106 int nFiles = 0;
107 for (File file : fileList) {
108 nFiles += enhancer.enhanceFile(file);
109 }
110 if (nFiles > 0) {
111 System.out.println("Enhanced: " + nFiles + " files");
112 }
113 } catch (Exception e) {
114 e.printStackTrace();
115 throw e;
116 }
117 }
118
119
124 public static void premain(String args, Instrumentation inst) {
125 if (!args.startsWith(AGENT_PREFIX)) {
126 throw new IllegalArgumentException
127 ("Unknown javaagent args: " + args +
128 " Args must start with: \"" + AGENT_PREFIX + '"');
129 }
130 args = args.substring(AGENT_PREFIX.length());
131 Set <String > packageNames = null;
132 boolean verbose = false;
133 if (args.length() > 0) {
134 packageNames = new HashSet <String >();
135 StringTokenizer tokens = new StringTokenizer (args, ",");
136 while (tokens.hasMoreTokens()) {
137 String token = tokens.nextToken();
138 if (token.startsWith("-")) {
139 if (token.equals("-v")) {
140 verbose = true;
141 } else {
142 throw new IllegalArgumentException
143 ("Unknown javaagent arg: " + token);
144 }
145 } else {
146 packageNames.add(token);
147 }
148 }
149 }
150 ClassEnhancer enhancer = new ClassEnhancer(packageNames);
151 enhancer.setVerbose(verbose);
152 inst.addTransformer(enhancer);
153 }
154
155
158 public ClassEnhancer() {
159 }
160
161
167 public void setVerbose(boolean verbose) {
168 this.verbose = verbose;
169 }
170
171
176 public boolean getVerbose() {
177 return verbose;
178 }
179
180
187 public ClassEnhancer(Set <String > packageNames) {
188 if (packageNames != null) {
189 packagePrefixes = new HashSet <String >();
190 for (String name : packageNames) {
191 packagePrefixes.add(name + '.');
192 }
193 }
194 }
195
196 public byte[] transform(ClassLoader loader,
197 String className,
198 Class <?> classBeingRedefined,
199 ProtectionDomain protectionDomain,
200 byte[] classfileBuffer) {
201 className = className.replace('/', '.');
202 byte[] bytes = enhance(className, classfileBuffer);
203 if (verbose && bytes != null) {
204 System.out.println("Enhanced: " + className);
205 }
206 return bytes;
207 }
208
209
221 public byte[] enhance(String className, byte[] classBytes) {
222 if (className != null && packagePrefixes != null) {
223 for (String prefix : packagePrefixes) {
224 if (className.startsWith(prefix)) {
225 return enhanceBytes(classBytes);
226 }
227 }
228 return null;
229 } else {
230 return enhanceBytes(classBytes);
231 }
232 }
233
234 int enhanceFile(File file)
235 throws IOException {
236
237 int nFiles = 0;
238 if (file.isDirectory()) {
239 String [] names = file.list();
240 if (names != null) {
241 for (int i = 0; i < names.length; i += 1) {
242 nFiles += enhanceFile(new File (file, names[i]));
243 }
244 }
245 } else if (file.getName().endsWith(".class")) {
246 byte[] newBytes = enhanceBytes(readFile(file));
247 if (newBytes != null) {
248 long modified = file.lastModified();
249 writeFile(file, newBytes);
250 file.setLastModified(modified);
251 nFiles += 1;
252 if (verbose) {
253 System.out.println("Enhanced: " + file);
254 }
255 }
256 }
257 return nFiles;
258 }
259
260 private byte[] readFile(File file)
261 throws IOException {
262
263 byte[] bytes = new byte[(int) file.length()];
264 FileInputStream in = new FileInputStream (file);
265 try {
266 in.read(bytes);
267 } finally {
268 in.close();
269 }
270 return bytes;
271 }
272
273 private void writeFile(File file, byte[] bytes)
274 throws IOException {
275
276 FileOutputStream out = new FileOutputStream (file);
277 try {
278 out.write(bytes);
279 } finally {
280 out.close();
281 }
282 }
283
284 private byte[] enhanceBytes(byte[] bytes) {
285
286
290 ClassWriter writer = new ClassWriter(true);
291 ClassVisitor visitor = writer;
292
293
294 visitor = new BytecodeEnhancer(visitor);
295
296
297 ClassReader reader = new ClassReader(bytes);
298 try {
299
300
304 reader.accept(visitor, false);
305 return writer.toByteArray();
306 } catch (BytecodeEnhancer.NotPersistentException e) {
307
308 return null;
309 }
310 }
311 }
312 |
Popular Tags |
Java API By Example, From Geeks To Geeks.