name.remal.annotation.bytecode.BytecodeAnnotationsScanner Maven / Gradle / Ivy

Go to download
package name.remal.annotation.bytecode;

import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import me.nallar.whocalled.WhoCalled;
import name.remal.annotation.AnnotationAttributeAlias;
import name.remal.gradle_plugins.api.RelocatePackages;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.Type;
import org.objectweb.asm.tree.AnnotationNode;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.FieldNode;
import org.objectweb.asm.tree.MethodNode;

import java.io.*;
import java.lang.annotation.Annotation;
import java.lang.annotation.Inherited;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.stream.Stream;

import static java.util.Collections.reverse;
import static java.util.Collections.singletonList;
import static java.util.stream.Collectors.*;
import static name.remal.SneakyThrow.sneakyThrow;
import static name.remal.reflection.HierarchyUtils.getHierarchy;
import static org.objectweb.asm.ClassReader.*;
import static org.objectweb.asm.Opcodes.ACC_ANNOTATION;
import static org.objectweb.asm.Type.*;

@RelocatePackages("org.objectweb.asm")
public class BytecodeAnnotationsScanner {

    @NotNull
    private final BytecodeRetriever bytecodeRetriever;

    public BytecodeAnnotationsScanner(@NotNull BytecodeRetriever bytecodeRetriever) {
        this.bytecodeRetriever = bytecodeRetriever;
    }

    public BytecodeAnnotationsScanner(@NotNull ClassLoader classLoader) {
        this(className -> {
            String resourceName = className.replace('.', '/') + ".class";
            try (InputStream inputStream = classLoader.getResourceAsStream(resourceName)) {
                return inputStream != null ? toByteArray(inputStream) : null;
            }
        });
    }

    public BytecodeAnnotationsScanner() {
        this(WhoCalled.$.getCallingClass().getClassLoader());
    }


    @NotNull
    public List<@NotNull BytecodeAnnotationAnnotationValue> getMetaAnnotations(@NotNull String className, @NotNull String annotationClassName) {
        Collection infos = isAnnotationInherited(annotationClassName)
            ? getAllInfos(className)
            : singletonList(getBytecodeAnnotationsInfo(className));

        return infos.stream()
            .flatMap(info -> info.annotationValues.stream())
            .filter(annotationValue -> annotationValue.getClassName().equals(annotationClassName))
            .distinct()
            .collect(toList());
    }

    @NotNull
    public List<@NotNull BytecodeAnnotationAnnotationValue> getMetaAnnotations(@NotNull Class clazz, @NotNull Class annotationClass) {
        return getMetaAnnotations(clazz.getName(), annotationClass.getName());
    }

    @Nullable
    public BytecodeAnnotationAnnotationValue getMetaAnnotation(@NotNull String className, @NotNull String annotationClassName) {
        List annotationValues = getMetaAnnotations(className, annotationClassName);
        return !annotationValues.isEmpty() ? annotationValues.get(0) : null;
    }

    @Nullable
    public BytecodeAnnotationAnnotationValue getMetaAnnotation(@NotNull Class clazz, @NotNull Class annotationClass) {
        return getMetaAnnotation(clazz.getName(), annotationClass.getName());
    }


    private static final ClassNode NULL_CLASS_NODE = new ClassNode();

    private final ConcurrentMap classNodesCache = new ConcurrentHashMap<>();

    @Nullable
    private ClassNode getClassNode(@NotNull String className) {
        ClassNode cached = classNodesCache.computeIfAbsent(className, curClassName -> {
            try {
                byte[] bytecode = bytecodeRetriever.retrieve(className);
                if (bytecode == null) {
                    return NULL_CLASS_NODE;
                }
                ClassNode classNode = new ClassNode();
                new ClassReader(bytecode).accept(classNode, SKIP_CODE | SKIP_DEBUG | SKIP_FRAMES);
                return classNode;

            } catch (Exception e) {
                throw sneakyThrow(e);
            }
        });
        return cached != NULL_CLASS_NODE ? cached : null;
    }


    private static final Set CORE_CLASS_NAMES = Stream.of(
        Object.class,
        Enum.class,
        Annotation.class,
        Comparable.class,
        Cloneable.class,
        Serializable.class,
        Externalizable.class,
        Closeable.class
    )
        .flatMap(clazz -> getHierarchy(clazz).stream())
        .map(Class::getName)
        .collect(toSet());

    @NotNull
    private final ConcurrentMap bytecodeAnnotationsInfosCache = new ConcurrentHashMap<>();

    {
        CORE_CLASS_NAMES.forEach(className -> bytecodeAnnotationsInfosCache.put(className, new BytecodeAnnotationsInfo(className, true)));
    }

    @NotNull
    private BytecodeAnnotationsInfo getBytecodeAnnotationsInfo(@NotNull String className) {
        BytecodeAnnotationsInfo result = bytecodeAnnotationsInfosCache.computeIfAbsent(className, BytecodeAnnotationsInfo::new);
        if (!result.isInitialized) {
            synchronized (result) {
                if (!result.isInitialized) {

                    ClassNode classNode = getClassNode(className);
                    if (classNode == null) {
                        throw new IllegalStateException("Bytecode can't be loaded: " + className);
                    }

                    if ((classNode.access & ACC_ANNOTATION) != 0) {
                        if (isLangCoreClass(className)) {
                            result.isInitialized = true;
                            return result;
                        }
                    }


                    {
                        // Collect parents:
                        Set parentClassNames = new LinkedHashSet<>();
                        if (classNode.superName != null) {
                            parentClassNames.add(classNode.superName.replace('/', '.'));
                        }
                        if (classNode.interfaces != null) {
                            classNode.interfaces.forEach(internalClassName -> parentClassNames.add(internalClassName.replace('/', '.')));
                        }
                        parentClassNames.forEach(parentClassName -> {
                            BytecodeAnnotationsInfo info = getBytecodeAnnotationsInfo(parentClassName);
                            if (info.isInitialized) { // break circular dependencies
                                result.parents.add(info);
                            }
                        });
                    }


                    { // Collect class annotations:
                        collectAnnotations(result.annotationValues, classNode.visibleAnnotations, classNode.invisibleAnnotations);
                    }


                    { // Collect fields annotations:
                        if (classNode.fields != null) {
                            classNode.fields.forEach(fieldNode -> {
                                Set annotations = new LinkedHashSet<>();
                                collectAnnotations(annotations, fieldNode.visibleAnnotations, fieldNode.invisibleAnnotations);
                                result.fieldsAnnotations.put(computeKey(fieldNode), annotations);
                            });
                        }
                    }


                    { // Collect methods annotations:
                        if (classNode.methods != null) {
                            classNode.methods.forEach(methodNode -> {
                                Set annotations = new LinkedHashSet<>();
                                collectAnnotations(annotations, methodNode.visibleAnnotations, methodNode.invisibleAnnotations);
                                result.methodsAnnotations.put(computeKey(methodNode), annotations);
                            });
                        }
                    }


                    result.isInitialized = true;

                }
            }
        }
        return result;
    }

    private static final String ANNOTATION_ATTRIBUTE_ALIAS_DESCR = getDescriptor(AnnotationAttributeAlias.class);

    @SafeVarargs
    private final void collectAnnotations(@NotNull Collection container, @Nullable List... annotationNodesArray) {
        if (annotationNodesArray == null) return;

        for (List annotationNodes : annotationNodesArray) {
            if (annotationNodes == null) continue;

            for (AnnotationNode annotationNode : annotationNodes) {
                if (annotationNode == null) continue;

                Set annotationValues = new LinkedHashSet<>();
                Queue queue = new LinkedList<>();
                queue.add(new ExpandingElement(toAnnotationValue(annotationNode)));
                while (true) {
                    ExpandingElement expandingElement = queue.poll();
                    if (expandingElement == null) break;

                    BytecodeAnnotationAnnotationValue annotationValue = expandingElement.annotationValue;
                    expandingElement.applyAttributes(annotationValue);
                    if (!annotationValues.add(annotationValue)) continue;

                    {
                        // Expand repeatable annotations:
                        if (annotationValue.getFields().size() == 1) {
                            BytecodeAnnotationValue value = annotationValue.getField("value");
                            if (value != null && value.isAnnotationsArray()) {
                                BytecodeAnnotationAnnotationValue[] items = value.asAnnotationsArray().getValue();
                                for (int i = items.length - 1; i >= 0; --i) {
                                    BytecodeAnnotationAnnotationValue item = items[i];
                                    expandingElement.applyAttributes(item);
                                    queue.add(new ExpandingElement(expandingElement, item));
                                }
                            }
                        }
                    }

                    if (!isLangCoreClass(annotationValue.getClassName())) {
                        ClassNode annotationClassNode = getClassNode(annotationValue.getClassName());
                        if (annotationClassNode != null) {
                            // Parse annotations on annotations:
                            Stream.of(annotationClassNode.visibleAnnotations, annotationClassNode.invisibleAnnotations)
                                .filter(Objects::nonNull)
                                .flatMap(Collection::stream)
                                .filter(node -> !isLangCoreClass(node.desc))
                                .map(this::toAnnotationValue)
                                .collect(toCollection(LinkedList::new))
                                .descendingIterator()
                                .forEachRemaining(nextAnnotationValue -> {
                                    ExpandingElement nextExpandingElement = new ExpandingElement(expandingElement, nextAnnotationValue);
                                    if (annotationClassNode.methods == null) return;
                                    annotationClassNode.methods.forEach(methodNode -> {
                                        if (!methodNode.desc.startsWith("()")) return;
                                        BytecodeAnnotationValue fieldValue = annotationValue.getField(methodNode.name);
                                        if (fieldValue == null) return;
                                        Stream.of(methodNode.visibleAnnotations, methodNode.invisibleAnnotations)
                                            .filter(Objects::nonNull)
                                            .flatMap(Collection::stream)
                                            .filter(node -> ANNOTATION_ATTRIBUTE_ALIAS_DESCR.equals(node.desc))
                                            .map(this::toAnnotationValue)
                                            .forEach(methodAnnotationValue -> {
                                                String annotationClassName = Optional.ofNullable(methodAnnotationValue.getField("annotationClass")).map(BytecodeAnnotationValue::asClass).map(BytecodeAnnotationClassValue::getClassName).orElse(null);
                                                if (annotationClassName == null) return;
                                                BytecodeAnnotationAnnotationValue attr = new BytecodeAnnotationAnnotationValue(annotationClassName);

                                                String attributeName = Optional.ofNullable(methodAnnotationValue.getField("attributeName")).map(BytecodeAnnotationValue::asString).map(BytecodeAnnotationStringValue::getValue).orElse(null);
                                                if (attributeName == null) return;
                                                attr.setField(attributeName, fieldValue);
                                                nextExpandingElement.attributes.add(attr);
                                            });
                                    });
                                    queue.add(nextExpandingElement);
                                });
                        }
                    }
                }

                if (annotationValues.size() == 1) {
                    container.add(annotationValues.iterator().next());
                } else {
                    List list = new ArrayList<>(annotationValues);
                    reverse(list);
                    container.addAll(list);
                }
            }
        }
    }

    private boolean isAnnotationInherited(@NotNull String annotationClassName) {
        BytecodeAnnotationsInfo info = getBytecodeAnnotationsInfo(annotationClassName);
        for (BytecodeAnnotationAnnotationValue annotationValue : info.annotationValues) {
            if (Inherited.class.getName().equals(annotationValue.getClassName())) {
                return true;
            }
        }
        return false;
    }

    @NotNull
    private Collection getAllInfos(@NotNull String rootClassName) {
        Map result = new LinkedHashMap<>();
        Queue queue = new LinkedList<>();
        queue.add(getBytecodeAnnotationsInfo(rootClassName));
        while (true) {
            BytecodeAnnotationsInfo info = queue.poll();
            if (info == null) break;
            if (!result.containsKey(info.className)) {
                result.put(info.className, info);
                info.parents.forEach(queue::add);
            }
        }
        return result.values();
    }


    @NotNull
    private BytecodeAnnotationAnnotationValue toAnnotationValue(@NotNull AnnotationNode annotationNode) {
        BytecodeAnnotationAnnotationValue result = new BytecodeAnnotationAnnotationValue(getType(annotationNode.desc).getClassName());
        if (annotationNode.values != null) {
            for (int index = 0; index < annotationNode.values.size(); index += 2) {
                result.setField(
                    (String) annotationNode.values.get(index),
                    toAnnotationValue(annotationNode.values.get(index + 1))
                );
            }
        }

        ClassNode annotationClassNode = getClassNode(result.getClassName());
        if (annotationClassNode != null && annotationClassNode.methods != null) {
            annotationClassNode.methods.forEach(methodNode -> {
                if (methodNode.annotationDefault != null) {
                    if (result.getField(methodNode.name) == null) {
                        result.setField(methodNode.name, toAnnotationValue(methodNode.annotationDefault));
                    }
                }
            });
        }

        return result;
    }

    @NotNull
    @SuppressFBWarnings("CLI_CONSTANT_LIST_INDEX")
    private BytecodeAnnotationValue toAnnotationValue(@NotNull Object value) {
        if (value instanceof Byte) return new BytecodeAnnotationByteValue((Byte) value);
        if (value instanceof Boolean) return new BytecodeAnnotationBooleanValue((Boolean) value);
        if (value instanceof Character) return new BytecodeAnnotationCharValue((Character) value);
        if (value instanceof Short) return new BytecodeAnnotationShortValue((Short) value);
        if (value instanceof Integer) return new BytecodeAnnotationIntValue((Integer) value);
        if (value instanceof Long) return new BytecodeAnnotationLongValue((Long) value);
        if (value instanceof Float) return new BytecodeAnnotationFloatValue((Float) value);
        if (value instanceof Double) return new BytecodeAnnotationDoubleValue((Double) value);
        if (value instanceof String) return new BytecodeAnnotationStringValue((String) value);
        if (value instanceof Type) return new BytecodeAnnotationClassValue(((Type) value).getClassName());
        if (value instanceof String[]) return new BytecodeAnnotationEnumValue(getType(((String[]) value)[0]).getClassName(), ((String[]) value)[1]);
        if (value instanceof AnnotationNode) return toAnnotationValue((AnnotationNode) value);
        if (value instanceof List) return toAnnotationValue((List) value);
        throw new IllegalArgumentException("Unsupported annotation value: " + value);
    }

    @NotNull
    @SuppressFBWarnings({"CLI_CONSTANT_LIST_INDEX", "UTA_USE_TO_ARRAY"})
    private BytecodeAnnotationValue toAnnotationValue(@NotNull List values) {
        Object firstValue = values.get(0);
        if (firstValue instanceof Byte) {
            byte[] typedValues = new byte[values.size()];
            for (int i = 0; i < values.size(); ++i) {
                typedValues[i] = (Byte) values.get(i);
            }
            return new BytecodeAnnotationBytesArrayValue(typedValues);
        }
        if (firstValue instanceof Boolean) {
            boolean[] typedValues = new boolean[values.size()];
            for (int i = 0; i < values.size(); ++i) {
                typedValues[i] = (Boolean) values.get(i);
            }
            return new BytecodeAnnotationBooleansArrayValue(typedValues);
        }
        if (firstValue instanceof Character) {
            char[] typedValues = new char[values.size()];
            for (int i = 0; i < values.size(); ++i) {
                typedValues[i] = (Character) values.get(i);
            }
            return new BytecodeAnnotationCharsArrayValue(typedValues);
        }
        if (firstValue instanceof Short) {
            short[] typedValues = new short[values.size()];
            for (int i = 0; i < values.size(); ++i) {
                typedValues[i] = (Short) values.get(i);
            }
            return new BytecodeAnnotationShortsArrayValue(typedValues);
        }
        if (firstValue instanceof Integer) {
            int[] typedValues = new int[values.size()];
            for (int i = 0; i < values.size(); ++i) {
                typedValues[i] = (Integer) values.get(i);
            }
            return new BytecodeAnnotationIntsArrayValue(typedValues);
        }
        if (firstValue instanceof Long) {
            long[] typedValues = new long[values.size()];
            for (int i = 0; i < values.size(); ++i) {
                typedValues[i] = (Long) values.get(i);
            }
            return new BytecodeAnnotationLongsArrayValue(typedValues);
        }
        if (firstValue instanceof Float) {
            float[] typedValues = new float[values.size()];
            for (int i = 0; i < values.size(); ++i) {
                typedValues[i] = (Float) values.get(i);
            }
            return new BytecodeAnnotationFloatsArrayValue(typedValues);
        }
        if (firstValue instanceof Double) {
            double[] typedValues = new double[values.size()];
            for (int i = 0; i < values.size(); ++i) {
                typedValues[i] = (Double) values.get(i);
            }
            return new BytecodeAnnotationDoublesArrayValue(typedValues);
        }
        if (firstValue instanceof String) {
            String[] typedValues = new String[values.size()];
            for (int i = 0; i < values.size(); ++i) {
                typedValues[i] = (String) values.get(i);
            }
            return new BytecodeAnnotationStringsArrayValue(typedValues);
        }
        if (firstValue instanceof Type) {
            String[] typedValues = new String[values.size()];
            for (int i = 0; i < values.size(); ++i) {
                typedValues[i] = ((Type) values.get(i)).getClassName();
            }
            return new BytecodeAnnotationClassesArrayValue(typedValues);
        }
        if (firstValue instanceof String[]) {
            BytecodeAnnotationEnumValue[] typedValues = new BytecodeAnnotationEnumValue[values.size()];
            for (int i = 0; i < values.size(); ++i) {
                Object value = values.get(i);
                typedValues[i] = new BytecodeAnnotationEnumValue(getType(((String[]) value)[0]).getClassName(), ((String[]) value)[1]);
            }
            return new BytecodeAnnotationEnumsArrayValue(typedValues);
        }
        if (firstValue instanceof AnnotationNode) {
            BytecodeAnnotationAnnotationValue[] typedValues = new BytecodeAnnotationAnnotationValue[values.size()];
            for (int i = 0; i < values.size(); ++i) {
                typedValues[i] = toAnnotationValue((AnnotationNode) values.get(i));
            }
            return new BytecodeAnnotationAnnotationsArrayValue(typedValues);
        }
        throw new IllegalArgumentException("Unsupported annotation value: " + values);
    }


    @NotNull
    private static byte[] toByteArray(@NotNull InputStream inputStream) throws IOException {
        try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
            int read;
            byte[] buffer = new byte[8192];
            while ((read = inputStream.read(buffer)) >= 0) {
                outputStream.write(buffer, 0, read);
            }
            return outputStream.toByteArray();
        }
    }

    private static boolean isLangCoreClass(@NotNull String className) {
        if (className.startsWith("java.lang.")) return true;
        if (className.startsWith("java/lang/")) return true;
        if (className.startsWith("Ljava/lang/")) return true;
        if (className.startsWith("kotlin.")) return true;
        if (className.startsWith("kotlin/")) return true;
        if (className.startsWith("Lkotlin/")) return true;
        if (className.startsWith("groovy.")) return true;
        if (className.startsWith("groovy/")) return true;
        if (className.startsWith("Lgroovy/")) return true;
        if (className.startsWith("scala.")) return true;
        if (className.startsWith("scala/")) return true;
        if (className.startsWith("Lscala/")) return true;
        return false;
    }

    @NotNull
    private static String computeKey(@NotNull FieldNode fieldNode) {
        return fieldNode.name;
    }

    @NotNull
    private static String computeKey(@NotNull MethodNode methodNode) {
        StringBuilder sb = new StringBuilder();
        sb.append(methodNode.name);
        sb.append('(');
        Type[] paramTypes = getArgumentTypes(methodNode.desc);
        for (int i = 0; i < paramTypes.length; ++i) {
            if (i == 1) sb.append(", ");
            sb.append(paramTypes[i].getDescriptor());
        }
        sb.append(')');
        return sb.toString();
    }

}


class BytecodeAnnotationsInfo {

    volatile boolean isInitialized;

    @NotNull
    final String className;

    BytecodeAnnotationsInfo(@NotNull String className, boolean isInitialized) {
        this.className = className;
        this.isInitialized = isInitialized;
    }

    BytecodeAnnotationsInfo(@NotNull String className) {
        this(className, false);
    }

    @NotNull
    @Override
    public String toString() {
        return BytecodeAnnotationsInfo.class.getSimpleName() + '{'
            + "className='" + className + '\''
            + ", isInitialized=" + isInitialized
            + '}';
    }


    @NotNull
    final List parents = new ArrayList<>();

    @NotNull
    final Set annotationValues = new LinkedHashSet<>();

    @NotNull
    final Map> fieldsAnnotations = new LinkedHashMap<>();

    @NotNull
    final Map> methodsAnnotations = new LinkedHashMap<>();

}


class ExpandingElement {

    @NotNull
    final BytecodeAnnotationAnnotationValue annotationValue;

    @NotNull
    final List attributes;

    ExpandingElement(@NotNull BytecodeAnnotationAnnotationValue annotationValue) {
        this.annotationValue = annotationValue;
        this.attributes = new ArrayList<>();
    }

    ExpandingElement(@NotNull ExpandingElement parent, @NotNull BytecodeAnnotationAnnotationValue annotationValue) {
        this(annotationValue);
        this.attributes.addAll(parent.attributes);
    }

    void applyAttributes(@NotNull BytecodeAnnotationAnnotationValue annotationValue) {
        if (attributes.isEmpty()) return;
        String className = annotationValue.getClassName();
        Map fields = annotationValue.getFields();
        for (int i = attributes.size() - 1; i >= 0; --i) {
            BytecodeAnnotationAnnotationValue attribute = attributes.get(i);
            if (attribute.getClassName().equals(className)) {
                attribute.getFields().forEach(fields::put);
            }
        }
    }

    @NotNull
    @Override
    public String toString() {
        return ExpandingElement.class.getSimpleName() + '{'
            + "annotationValue=" + annotationValue
            + ", attributes=" + attributes
            + '}';
    }

}