AbstractTransformTask.java revision 674060f01e9090cd21b3c5656cc3204912ad17a6 
1/*
2 * Copyright 2003,2004 The Apache Software Foundation
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
5 * use this file except in compliance with the License. You may obtain a copy of
6 * 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, WITHOUT
12 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
13 * License for the specific language governing permissions and limitations under
14 * the License.
15 */
16package org.mockito.cglib.transform;
17
18import java.io.*;
19import java.net.MalformedURLException;
20import java.util.*;
21import java.util.zip.*;
22
23import org.apache.tools.ant.*;
24import org.mockito.asm.*;
25import org.mockito.cglib.core.*;
26
27abstract public class AbstractTransformTask extends AbstractProcessTask {
28    private static final int ZIP_MAGIC = 0x504B0304;
29
30    private static final int CLASS_MAGIC = 0xCAFEBABE;
31
32    private boolean verbose;
33
34    public void setVerbose(boolean verbose) {
35        this.verbose = verbose;
36    }
37
38    /**
39     * returns transformation for source class
40     *
41     * @param classInfo
42     *            class information
43     *            class name := classInfo[ 0 ]
44     *            super class  name := classInfo[ 1 ]
45     *            interfaces := classInfo[ >1 ]
46     */
47    abstract protected ClassTransformer getClassTransformer(String[] classInfo);
48
49    protected Attribute[] attributes() {
50        return null;
51    }
52
53    protected void processFile(File file) throws Exception {
54
55        if (isClassFile(file)) {
56
57            processClassFile(file);
58
59        } else if (isJarFile(file)) {
60
61            processJarFile(file);
62
63        } else {
64
65            log("ignoring " + file.toURL(), Project.MSG_WARN);
66
67        }
68    }
69
70    /**
71     * @param file
72     * @throws Exception
73     * @throws FileNotFoundException
74     * @throws IOException
75     * @throws MalformedURLException
76     */
77    private void processClassFile(File file) throws Exception,
78            FileNotFoundException, IOException, MalformedURLException {
79
80        ClassReader reader = getClassReader(file);
81        String name[] = ClassNameReader.getClassInfo(reader);
82        ClassWriter w = new DebuggingClassWriter(ClassWriter.COMPUTE_MAXS);
83        ClassTransformer t = getClassTransformer(name);
84        if (t != null) {
85
86            if (verbose) {
87                log("processing " + file.toURL());
88            }
89            new TransformingClassGenerator(new ClassReaderGenerator(
90                    getClassReader(file), attributes(), getFlags()), t)
91                    .generateClass(w);
92            FileOutputStream fos = new FileOutputStream(file);
93            try {
94                fos.write(w.toByteArray());
95            } finally {
96                fos.close();
97            }
98
99        }
100
101    }
102
103    protected int getFlags() {
104        return 0;
105    }
106
107    private static ClassReader getClassReader(File file) throws Exception {
108        InputStream in = new BufferedInputStream(new FileInputStream(file));
109        try {
110            ClassReader r = new ClassReader(in);
111            return r;
112        } finally {
113            in.close();
114        }
115
116    }
117
118    protected boolean isClassFile(File file) throws IOException {
119
120        return checkMagic(file, CLASS_MAGIC);
121
122    }
123
124    protected void processJarFile(File file) throws Exception {
125
126        if (verbose) {
127            log("processing " + file.toURL());
128        }
129
130        File tempFile = File.createTempFile(file.getName(), null, new File(file
131                .getAbsoluteFile().getParent()));
132        try{
133
134            ZipInputStream zip = new ZipInputStream(new FileInputStream(file));
135            try {
136                FileOutputStream fout = new FileOutputStream(tempFile, false);
137                try{
138                 ZipOutputStream out = new ZipOutputStream(fout);
139
140                    ZipEntry entry;
141                    while ((entry = zip.getNextEntry()) != null) {
142
143
144                        byte bytes[] = getBytes(zip);
145
146                        if (!entry.isDirectory()) {
147
148                            DataInputStream din = new DataInputStream(
149                                              new ByteArrayInputStream(bytes)
150                                            );
151
152                            if (din.readInt() == CLASS_MAGIC) {
153
154                                bytes = process(bytes);
155
156                            } else {
157                                if (verbose) {
158                                 log("ignoring " + entry.toString());
159                                }
160                            }
161                        }
162
163                        ZipEntry outEntry = new ZipEntry(entry.getName());
164                        outEntry.setMethod(entry.getMethod());
165                        outEntry.setComment(entry.getComment());
166                        outEntry.setSize(bytes.length);
167
168
169                        if(outEntry.getMethod() == ZipEntry.STORED){
170                            CRC32 crc = new CRC32();
171                            crc.update(bytes);
172                            outEntry.setCrc( crc.getValue() );
173                            outEntry.setCompressedSize(bytes.length);
174                        }
175                        out.putNextEntry(outEntry);
176                        out.write(bytes);
177                        out.closeEntry();
178                        zip.closeEntry();
179
180                    }
181                    out.close();
182                }finally{
183                 fout.close();
184                }
185            } finally {
186                zip.close();
187            }
188
189
190            if(file.delete()){
191
192                File newFile = new File(tempFile.getAbsolutePath());
193
194                if(!newFile.renameTo(file)){
195                    throw new IOException("can not rename " + tempFile + " to " + file);
196                }
197
198            }else{
199                throw new IOException("can not delete " + file);
200            }
201
202        }finally{
203
204            tempFile.delete();
205
206        }
207
208    }
209
210    /**
211     * @param bytes
212     * @return
213     * @throws IOException
214     * @throws Exception
215     */
216    private byte[] process(byte[] bytes) throws Exception {
217
218        ClassReader reader = new ClassReader(new ByteArrayInputStream(bytes));
219        String name[] = ClassNameReader.getClassInfo(reader);
220        ClassWriter w = new DebuggingClassWriter(ClassWriter.COMPUTE_MAXS);
221        ClassTransformer t = getClassTransformer(name);
222        if (t != null) {
223            if (verbose) {
224                log("processing " + name[0]);
225            }
226            new TransformingClassGenerator(new ClassReaderGenerator(
227                    new ClassReader(new ByteArrayInputStream(bytes)),
228                    attributes(), getFlags()), t).generateClass(w);
229            ByteArrayOutputStream out = new ByteArrayOutputStream();
230            out.write(w.toByteArray());
231            return out.toByteArray();
232        }
233        return bytes;
234    }
235
236    /**
237     * @param zip
238     * @return
239     * @throws IOException
240     */
241    private byte[] getBytes(ZipInputStream zip) throws IOException {
242
243        ByteArrayOutputStream bout = new ByteArrayOutputStream();
244        InputStream in = new BufferedInputStream(zip);
245        int b;
246        while ((b = in.read()) != -1) {
247            bout.write(b);
248        }
249        return bout.toByteArray();
250    }
251
252    private boolean checkMagic(File file, long magic) throws IOException {
253        DataInputStream in = new DataInputStream(new FileInputStream(file));
254        try {
255            int m = in.readInt();
256            return magic == m;
257        } finally {
258            in.close();
259        }
260    }
261
262    protected boolean isJarFile(File file) throws IOException {
263        return checkMagic(file, ZIP_MAGIC);
264    }
265
266}