AtlassianSpringByteCodeScanner

package com.atlassian.plugin.spring.scanner.core;

import com.atlassian.plugin.spring.scanner.annotation.component.ClasspathComponent;
import com.atlassian.plugin.spring.scanner.annotation.export.ExportAsService;
import com.atlassian.plugin.spring.scanner.annotation.export.ModuleType;
import com.atlassian.plugin.spring.scanner.annotation.imports.ComponentImport;
import com.atlassian.plugin.spring.scanner.util.CommonConstants;
import javassist.bytecode.AnnotationsAttribute;
import javassist.bytecode.ClassFile;
import javassist.bytecode.Descriptor;
import javassist.bytecode.FieldInfo;
import javassist.bytecode.MethodInfo;
import javassist.bytecode.annotation.Annotation;
import org.reflections.Reflections;
import org.reflections.scanners.AbstractScanner;
import org.reflections.util.ConfigurationBuilder;
import org.reflections.util.FilterBuilder;
import org.slf4j.Logger;
import org.springframework.stereotype.Component;

import javax.inject.Named;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import static com.google.common.base.Strings.isNullOrEmpty;
import static java.lang.String.format;

/**
 * A reflections scanner than knows about well-known injection annotations such as {@link Component},
 * {@link Named} etc..
 * <p>
 * This scanner will be passed every class file in the build directory and it will examine it for Annotations of
 * interest.
 * <p>
 * This class uses a mix of the higher level Reflections code to get class information and also uses the lower level
 * Javassist byte code helpers to get the rest of the information.  We prefer the former but end up having to use the
 * latter to get the job done.
 */
@SuppressWarnings("unchecked")
public class AtlassianSpringByteCodeScanner extends AbstractScanner {
    private final Logger log;
    private final Stats stats;
    private final Errors errors;
    private final SpringIndexWriter springIndexWriter;
    private final ProfileFinder profileFinder;
    private final JavassistHelper javassistHelper;
    private final AnnotationValidator annotationValidator;
    private boolean dbg;

    public AtlassianSpringByteCodeScanner(final ByteCodeScannerConfiguration configuration) {
        this.log = configuration.getLog();
        this.springIndexWriter = new SpringIndexWriter(configuration.getOutputDirectory());
        this.javassistHelper = new JavassistHelper();
        this.profileFinder = new ProfileFinder(configuration.getClassPathUrls(), log);
        this.annotationValidator = new AnnotationValidator();
        this.stats = new Stats();
        this.errors = new Errors();
        this.dbg = log.isDebugEnabled() || configuration.isVerbose();

        go(configuration);

        if (! configuration.isPermitDuplicateImports()) {
            this.annotationValidator.validate(this.errors);
        }
    }

    private void go(final ByteCodeScannerConfiguration configuration) {
        //
        ConfigurationBuilder config = new ConfigurationBuilder();

        config.setUrls(configuration.getClassPathUrls());

        if (!isNullOrEmpty(configuration.getIncludeExclude())) {
            config.filterInputsBy(FilterBuilder.parse(configuration.getIncludeExclude()));
        }

        //
        // we use our own specific scanner which means that the underlying Reflections code will
        // not be able to consume our output.  But that's ok.
        //
        config.setScanners(this);

        // we don want reflections logs thank you
        try {
            Reflections.log = null;
        } catch (Error e) {
            //ignore
        }

        // this will cause the scanner code to run!
        new Reflections(config);

        // and write the results
        springIndexWriter.writeIndexes();
    }

    public Stats getStats() {
        return stats;
    }


    public Errors getErrors() {
        return errors;
    }

    @Override
    public void scan(final Object cls) {
        ClassFile classFile = (ClassFile) cls;
        try {
            scanClass(classFile);
        } catch (Exception e) {
            log.error(format("Unable to run byte code scanner on class %s. Continuing to the next class...", cls));
        }
    }

    private void scanClass(final ClassFile classFile) {
        stats.classesEncountered++;

        Set<String> profiles = profileFinder.getProfiles(classFile);

        String className = getMetadataAdapter().getClassName(classFile);
        List<String> classAnnotationNames = getMetadataAdapter().getClassAnnotationNames(classFile);

        for (String annotationType : classAnnotationNames) {
            if (isSuitableClassAnnotation(annotationType)) {
                // quick check to see that the class makes sense
                if (!isSuitableClass(classFile)) {
                    debug(format("\t(X) Class not suitable '%s' for annotation '%s'", className, annotationType));
                    return;
                }
                stats.componentClassesEncountered++;
                String nameFromAnnotation = javassistHelper.getAnnotationMember(classFile, annotationType, "value");


                debug(format("(/) Found annotation '%s' inside class '%s' with name '%s'", annotationType, className, nameFromAnnotation));

                springIndexWriter.encounteredAnnotation(profiles, annotationType, nameFromAnnotation, className);
                annotationValidator.encounteredAnnotation(annotationType, className, className);

                //
                // if its a @ModuleType then we need to add extra fix-up to the host container so that
                // can work more easily out of the box.  We just slip in a Component with the right class.  Easy!
                if (ModuleType.class.getCanonicalName().equals(annotationType)) {
                    springIndexWriter.encounteredAnnotation(profiles, Component.class.getCanonicalName(), "", CommonConstants.HOST_CONTAINER_CLASS);
                }
            }
        }

        // find constructor parameter imports etc...
        visitConstructors(classFile, profiles);
        //
        // find field annotated imports etc..
        visitFields(classFile, profiles);
    }

    private void debug(String msg) {
        if (dbg) {
            log.info("\t" + msg);
        }
    }

    private boolean isSuitableClass(final ClassFile classFile) {
        String className = classFile.getName();
        if (classFile.isInterface()) {
            log.error(format("Found a type [%s] annotated as a component, but the type is not a concrete class. NOT adding to index file!!", className));
            return false;
        }
        if (classFile.isAbstract()) {
            log.error(format("Found a type [%s] annotated as a component, but the type is abstract. NOT adding to index file!!", className));
            return false;
        }
        // package-info don't count either but its not an error to encounter one
        return !profileFinder.isPackageClass(classFile);
    }

    private boolean isSuitableClassAnnotation(final String annotationType) {
        return super.acceptResult(annotationType) && springIndexWriter.isInteresting(annotationType);
    }

    /**
     * This will visit the constructors of the class and see if they have a annotations that may be interesting for the
     * scanner
     *
     * @param classFile the class we are inspecting
     * @param profiles  the profiles that are in play for the index
     */
    private void visitConstructors(final ClassFile classFile, final Set<String> profiles) {
        String className = classFile.getName();
        List<MethodInfo> methods = classFile.getMethods();
        for (MethodInfo method : methods) {
            String methodName = method.getName();
            if (method.isConstructor()) {
                // parameter 'names' in this case is actually parameter types
                List<String> parameterTypes = getMetadataAdapter().getParameterNames(method);
                for (int i = 0; i < parameterTypes.size(); i++) {
                    String parameterType = parameterTypes.get(i);
                    List<Annotation> parameterAnnotationNames = javassistHelper.getParameterAnnotations(method, i);

                    for (Annotation annotation : parameterAnnotationNames) {
                        String annotationType = annotation.getTypeName();
                        if (acceptResult(annotationType) && springIndexWriter.isParameterOrFieldAnnotation(annotationType)) {
                            String nameFromAnnotation = javassistHelper.getAnnotationMember(annotation, "value");

                            debug(format("(/) Found '%s' inside class '%s' method '%s' parameter '%s'", annotationType, className, methodName, parameterType));

                            springIndexWriter.encounteredAnnotation(profiles, annotationType, nameFromAnnotation, parameterType);
                            annotationValidator.encounteredAnnotation(annotationType, parameterType, className);
                        }
                    }

                }
            }
        }
    }

    private void visitFields(final ClassFile classFile, final Set<String> profiles) {
        String className = classFile.getName();
        List<FieldInfo> fields = classFile.getFields();
        for (FieldInfo field : fields) {
            final List<String> annotationTypes = new LinkedList<>();
            String fieldName = field.getName();
            AnnotationsAttribute annotations = (AnnotationsAttribute) field.getAttribute(AnnotationsAttribute.visibleTag);
            if (annotations != null) {
                for (Annotation annotation : annotations.getAnnotations()) {
                    String annotationType = annotation.getTypeName();
                    annotationTypes.add(annotationType);

                    if (acceptResult(annotationType) && springIndexWriter.isParameterOrFieldAnnotation(annotationType)) {
                        String fieldType = Descriptor.toClassName(field.getDescriptor());
                        String nameFromAnnotation = javassistHelper.getAnnotationMember(annotation, "value");

                        debug(format("(/) Found '%s' inside class '%s' on field '%s' of type '%s'", annotationType, className, fieldName, fieldType));

                        springIndexWriter.encounteredAnnotation(profiles, annotationType, nameFromAnnotation, fieldType);
                        annotationValidator.encounteredAnnotation(annotationType, fieldType, className);
                    }
                }
            }

            if (annotationTypes.contains(ComponentImport.class.getCanonicalName())) {

                final List<String> productImportsPresentOnField = new ArrayList<>(SpringIndexWriter.KNOWN_PRODUCT_IMPORT_ANNOTATIONS);

                productImportsPresentOnField.retainAll(annotationTypes);

                if (!productImportsPresentOnField.isEmpty()) {
                    errors.addError(String.format("ComponentImport annotation cannot be used with product specific component imports: %s found on %s.%s",
                            Arrays.toString(productImportsPresentOnField.toArray()), classFile.getName(), fieldName));
                }
            }
        }
    }

    public static class Stats {
        private int classesEncountered;
        private int componentClassesEncountered;

        public int getClassesEncountered() {
            return classesEncountered;
        }

        public int getComponentClassesEncountered() {
            return componentClassesEncountered;
        }
    }

    public static class Errors {
        private final List<String> errorsEncountered = new ArrayList<>();

        public void addError(String error) {
            errorsEncountered.add(error);
        }

        public List<String> getErrorsEncountered() {
            return errorsEncountered;
        }
    }

    /**
     * Mantains a collection of annotated component names, and for each annotation,
     * the locations where the annotation was found.
     */
    public static class AnnotationReferences {
        private Map<String, List<String>> references;

        public AnnotationReferences() {
            this.references = new HashMap<>();
        }

        public void addReference(String component, String whereReferenced) {
            List<String> sources = references.computeIfAbsent(component, k -> new ArrayList());
            sources.add(whereReferenced);
        }

        public List<String> getReferences(String component) {
            return references.get(component);
        }

        public Set<String> intersect(AnnotationReferences that) {
            Set intersectSet = new HashSet(this.references.keySet());
            intersectSet.retainAll(that.references.keySet());
            return intersectSet;
        }
    }

    public static class AnnotationValidator {
        private final Set<String> declaredComponentAnnotationTypes;
        private final Set<String> importedComponentAnnotationTypes;

        private AnnotationReferences declaredComponents;
        private AnnotationReferences importedComponents;

        public AnnotationValidator() {
            this.declaredComponents = new AnnotationReferences();

            // All annotation types that cause us to create a component
            this.declaredComponentAnnotationTypes = new HashSet<>();
            this.declaredComponentAnnotationTypes.add(ExportAsService.class.getCanonicalName());
            this.declaredComponentAnnotationTypes.add(Component.class.getCanonicalName());
            this.declaredComponentAnnotationTypes.add(Named.class.getCanonicalName());

            this.importedComponents = new AnnotationReferences();

            // All annotation types that cause us to import a component
            this.importedComponentAnnotationTypes = new HashSet<>();
            this.importedComponentAnnotationTypes.add(ComponentImport.class.getCanonicalName());
            this.importedComponentAnnotationTypes.add(ClasspathComponent.class.getCanonicalName());
            this.importedComponentAnnotationTypes.addAll(SpringIndexWriter.KNOWN_PRODUCT_IMPORT_ANNOTATIONS);
        }

        public void encounteredAnnotation(String annotationType, String component, String whereReferenced) {
            if (importedComponentAnnotationTypes.contains(annotationType)) {
                importedComponents.addReference(component, whereReferenced);
            } else if (declaredComponentAnnotationTypes.contains(annotationType)) {
                declaredComponents.addReference(component, whereReferenced);
            }
        }

        public void validate(Errors errors) {
            Set<String> conflictingComponents = importedComponents.intersect(declaredComponents);

            for (String component : conflictingComponents) {
                List<String> imports = importedComponents.getReferences(component);
                List<String> declarations = declaredComponents.getReferences(component);

                errors.addError(format("Cannot import a component within a plugin that declares the same component. " +
                                        "Component '%s' was declared in %s but imported in %s",
                                        component,
                                        Arrays.toString(declarations.toArray()),
                                        Arrays.toString(imports.toArray())));
            }
        }
    }
}