001/*
002 * Copyright 2011 Atteo.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
005 * in compliance with the License. You may obtain a copy of the License at
006 *
007 * http://www.apache.org/licenses/LICENSE-2.0
008 *
009 * Unless required by applicable law or agreed to in writing, software distributed under the License
010 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
011 * or implied. See the License for the specific language governing permissions and limitations under
012 * the License.
013 */
014package org.atteo.classindex;
015
016import java.io.BufferedReader;
017import java.io.ByteArrayInputStream;
018import java.io.ByteArrayOutputStream;
019import java.io.IOException;
020import java.io.InputStream;
021import java.io.StringReader;
022import java.nio.charset.StandardCharsets;
023import java.util.HashMap;
024import java.util.List;
025import java.util.Map;
026import java.util.jar.JarEntry;
027import java.util.jar.JarOutputStream;
028
029import org.apache.maven.plugins.shade.relocation.Relocator;
030import org.apache.maven.plugins.shade.resource.ResourceTransformer;
031import org.codehaus.plexus.util.IOUtil;
032
033public class ClassIndexTransformer implements ResourceTransformer {
034    public static final String SUBCLASS_INDEX_PREFIX = "META-INF/services/";
035    public static final String ANNOTATED_INDEX_PREFIX = "META-INF/annotations/";
036    public static final String PACKAGE_INDEX_NAME = "jaxb.index";
037    private final Map<String, ByteArrayOutputStream> serviceEntries = new HashMap<>();
038    private List<Relocator> relocators;
039
040    @Override
041    public boolean canTransformResource(String resource) {
042        return resource.startsWith(SUBCLASS_INDEX_PREFIX)
043                || resource.startsWith(ANNOTATED_INDEX_PREFIX)
044                || resource.endsWith("/" + PACKAGE_INDEX_NAME);
045    }
046
047    @Override
048    public void processResource(String resource, InputStream is, List<Relocator> relocators)
049            throws IOException {
050        if (this.relocators == null) {
051            this.relocators = relocators;
052        }
053
054        ByteArrayOutputStream data = serviceEntries.get(resource);
055        if (data == null) {
056            data = new ByteArrayOutputStream();
057            serviceEntries.put(resource, data);
058        }
059
060        try {
061            String content = IOUtil.toString(is, StandardCharsets.UTF_8.name());
062            StringReader reader = new StringReader(content);
063            BufferedReader lineReader = new BufferedReader(reader);
064            String line;
065            while ((line = lineReader.readLine()) != null) {
066                String qualifiedClassName = relocateIfNeeded(line);
067                data.write(qualifiedClassName.getBytes(StandardCharsets.UTF_8));
068                data.write("\n".getBytes(StandardCharsets.UTF_8));
069            }
070        } finally {
071            is.close();
072        }
073    }
074
075    @Override
076    public boolean hasTransformedResource() {
077        return serviceEntries.size() > 0;
078    }
079
080    @Override
081    public void modifyOutputStream(JarOutputStream jos)
082            throws IOException {
083        for (Map.Entry<String, ByteArrayOutputStream> entry : serviceEntries.entrySet()) {
084            String key = entry.getKey();
085            ByteArrayOutputStream stream = entry.getValue();
086            jos.putNextEntry(new JarEntry(relocateFileName(key)));
087            IOUtil.copy(new ByteArrayInputStream(stream.toByteArray()), jos);
088            stream.reset();
089        }
090    }
091
092    private String relocateFileName(String key) {
093        String prefix = "";
094        if (key.startsWith(SUBCLASS_INDEX_PREFIX)) {
095            prefix = SUBCLASS_INDEX_PREFIX;
096        } else if (key.startsWith(ANNOTATED_INDEX_PREFIX)) {
097            prefix = ANNOTATED_INDEX_PREFIX;
098        }
099        return prefix + relocateIfNeeded(key.substring(prefix.length()));
100    }
101
102    private String relocateIfNeeded(String key) {
103        if (relocators == null) {
104            return key;
105        }
106
107        for (Relocator relocator : relocators) {
108            if (relocator.canRelocateClass(key)) {
109                return relocator.relocateClass(key);
110            }
111        }
112
113        return key;
114    }
115}