content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.tooling.muzzle;
  
  import static java.nio.charset.StandardCharsets.UTF_8;
  import static java.util.Collections.singleton;
  
  import com.google.common.base.Preconditions;
  import com.google.common.base.Strings;
  import com.google.common.graph.Graph;
  import com.google.common.graph.GraphBuilder;
  import com.google.common.graph.Graphs;
  import com.google.common.graph.MutableGraph;
  import io.opentelemetry.javaagent.tooling.muzzle.references.ClassRef;
  import io.opentelemetry.javaagent.tooling.muzzle.references.Flag;
  import java.io.BufferedReader;
  import java.io.IOException;
  import java.io.InputStream;
  import java.io.InputStreamReader;
  import java.net.URLConnection;
  import java.util.ArrayDeque;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashSet;
  import java.util.Iterator;
  import java.util.LinkedHashMap;
  import java.util.LinkedList;
  import java.util.List;
  import java.util.Map;
  import java.util.Queue;
  import java.util.Set;
  import java.util.function.Predicate;
  import java.util.regex.Pattern;
  import javax.annotation.Nullable;
  import net.bytebuddy.description.method.MethodDescription;
  import org.objectweb.asm.ClassReader;
  
  /**
   * {@link LinkedHashMap} is used for reference map to guarantee a deterministic order of iteration,
   * so that bytecode generated based on it would also be deterministic.
   *
   * <p>This class is only called at compile time by the {@code MuzzleCodeGenerationPlugin} ByteBuddy
   * plugin.
   */
  public final class ReferenceCollector {
  
    private final Map<String, ClassRef> references = new LinkedHashMap<>();
    private final MutableGraph<String> helperSuperClassGraph = GraphBuilder.directed().build();
    private final VirtualFieldMappingsBuilderImpl virtualFieldMappingsBuilder =
        new VirtualFieldMappingsBuilderImpl();
    private final Set<String> visitedClasses = new HashSet<>();
    private final HelperClassPredicate helperClassPredicate;
    private final ClassLoader resourceLoader;
  
    // only used by tests
    ReferenceCollector(Predicate<String> libraryInstrumentationPredicate) {
      this(libraryInstrumentationPredicate, ReferenceCollector.class.getClassLoader());
    }
  
    public ReferenceCollector(
        Predicate<String> libraryInstrumentationPredicate, ClassLoader resourceLoader) {
      this.helperClassPredicate = new HelperClassPredicate(libraryInstrumentationPredicate);
      this.resourceLoader = resourceLoader;
    }
  
    /**
     * If passed {@code resource} path points to an SPI file (either Java {@link
     * java.util.ServiceLoader} or AWS SDK {@code ExecutionInterceptor}) reads the file and adds every
     * implementation as a reference, traversing the graph of classes until a non-instrumentation
     * (external) class is encountered.
     *
     * @see HelperClassPredicate
     */
    public void collectReferencesFromResource(HelperResource helperResource) {
      if (!isSpiFile(helperResource.getApplicationPath())) {
        return;
      }
  
      List<String> spiImplementations = new ArrayList<>();
      try (InputStream stream = getResourceStream(helperResource.getAgentPath())) {
        BufferedReader reader = new BufferedReader(new InputStreamReader(stream, UTF_8));
        while (reader.ready()) {
          String line = reader.readLine();
          if (!Strings.isNullOrEmpty(line)) {
            spiImplementations.add(line);
          }
        }
      } catch (IOException e) {
        throw new IllegalStateException("Error reading resource " + helperResource.getAgentPath(), e);
      }
  
      visitClassesAndCollectReferences(spiImplementations, /* startsFromAdviceClass= */ false);
    }
  
    private static final Pattern AWS_SDK_V2_SERVICE_INTERCEPTOR_SPI =
        Pattern.compile("software/amazon/awssdk/services/\\w+(/\\w+)?/execution.interceptors");
  
    private static final Pattern AWS_SDK_V1_SERVICE_INTERCEPTOR_SPI =
        Pattern.compile("com/amazonaws/services/\\w+(/\\w+)?/request.handler2s");
  
    private static boolean isSpiFile(String resource) {
      return resource.startsWith("META-INF/services/")
          || resource.equals("software/amazon/awssdk/global/handlers/execution.interceptors")
          || resource.equals("com/amazonaws/global/handlers/request.handler2s")
          || AWS_SDK_V2_SERVICE_INTERCEPTOR_SPI.matcher(resource).matches()
          || AWS_SDK_V1_SERVICE_INTERCEPTOR_SPI.matcher(resource).matches();
    }
  
    /**
     * Traverse a graph of classes starting from {@code adviceClassName} and collect all references to
     * both internal (instrumentation) and external classes.
     *
     * <p>The graph of classes is traversed until a non-instrumentation (external) class is
     * encountered.
     *
     * @param adviceClassName Starting point for generating references.
     * @see HelperClassPredicate
     */
    public void collectReferencesFromAdvice(String adviceClassName) {
      visitClassesAndCollectReferences(singleton(adviceClassName), /* startsFromAdviceClass= */ true);
    }
  
    private void visitClassesAndCollectReferences(
        Collection<String> startingClasses, boolean startsFromAdviceClass) {
      Queue<String> instrumentationQueue = new ArrayDeque<>(startingClasses);
      boolean isAdviceClass = startsFromAdviceClass;
  
      while (!instrumentationQueue.isEmpty()) {
        String visitedClassName = instrumentationQueue.remove();
        visitedClasses.add(visitedClassName);
  
        try (InputStream in = getClassFileStream(visitedClassName)) {
          // only start from method bodies for the advice class (skips class/method references)
          ReferenceCollectingClassVisitor cv =
              new ReferenceCollectingClassVisitor(helperClassPredicate, isAdviceClass);
          ClassReader reader = new ClassReader(in);
          reader.accept(cv, ClassReader.SKIP_FRAMES);
  
          for (Map.Entry<String, ClassRef> entry : cv.getReferences().entrySet()) {
            String refClassName = entry.getKey();
            ClassRef reference = entry.getValue();
  
            // Don't generate references created outside of the instrumentation package.
            if (!visitedClasses.contains(refClassName)
                && helperClassPredicate.isHelperClass(refClassName)) {
              instrumentationQueue.add(refClassName);
            }
            addReference(refClassName, reference);
          }
          collectHelperClasses(
              isAdviceClass, visitedClassName, cv.getHelperClasses(), cv.getHelperSuperClasses());
  
          virtualFieldMappingsBuilder.registerAll(cv.getVirtualFieldMappings());
        } catch (IOException e) {
          throw new IllegalStateException("Error reading class " + visitedClassName, e);
        }
  
        if (isAdviceClass) {
          isAdviceClass = false;
        }
      }
    }
  
    private InputStream getClassFileStream(String className) throws IOException {
      return getResourceStream(Utils.getResourceName(className));
    }
  
    private InputStream getResourceStream(String resource) throws IOException {
      URLConnection connection =
          Preconditions.checkNotNull(
                  resourceLoader.getResource(resource), "Couldn't find resource %s", resource)
              .openConnection();
  
      // Since the JarFile cache is not per class loader, but global with path as key, using cache may
      // cause the same instance of JarFile being used for consecutive builds, even if the file has
      // been changed. There is still another cache in ZipFile.Source which checks last modified time
      // as well, so the zip index is not scanned again on every class.
      connection.setUseCaches(false);
      return connection.getInputStream();
    }
  
    private void addReference(String refClassName, ClassRef reference) {
      if (references.containsKey(refClassName)) {
        references.put(refClassName, references.get(refClassName).merge(reference));
      } else {
        references.put(refClassName, reference);
      }
    }
  
    private void collectHelperClasses(
        boolean isAdviceClass,
        String className,
        Set<String> helperClasses,
        Set<String> helperSuperClasses) {
      for (String helperClass : helperClasses) {
        helperSuperClassGraph.addNode(helperClass);
      }
      if (!isAdviceClass) {
        for (String helperSuperClass : helperSuperClasses) {
          helperSuperClassGraph.putEdge(className, helperSuperClass);
        }
      }
    }
  
    public Map<String, ClassRef> getReferences() {
      return references;
    }
  
    public void prune() {
      // helper classes that may help another helper class implement an abstract library method
      // must be retained
      // for example if helper class A extends helper class B, and A also implements a library
      // interface L, then B needs to be retained so that it can be used at runtime to verify that A
      // implements all of L's methods.
      // Super types of A that are not also helper classes do not need to be retained because they can
      // be looked up on the classpath at runtime, see HelperReferenceWrapper.create().
      Set<ClassRef> helperClassesParticipatingInLibrarySuperType =
          getHelperClassesParticipatingInLibrarySuperType();
  
      for (Iterator<ClassRef> i = references.values().iterator(); i.hasNext(); ) {
        ClassRef reference = i.next();
        if (helperClassPredicate.isLibraryClass(reference.getClassName())) {
          // these are the references to library classes which need to be checked at runtime
          continue;
        }
        if (helperClassesParticipatingInLibrarySuperType.contains(reference)) {
          // these need to be kept in order to check that abstract methods are implemented,
          // and to check that declared super class fields are present
          //
          // can at least prune constructors, private, and static methods, since those cannot be used
          // to help implement an abstract library method
          reference
              .getMethods()
              .removeIf(
                  method ->
                      method.getName().equals(MethodDescription.CONSTRUCTOR_INTERNAL_NAME)
                          || method.getFlags().contains(Flag.VisibilityFlag.PRIVATE)
                          || method.getFlags().contains(Flag.OwnershipFlag.STATIC));
          continue;
        }
        i.remove();
      }
    }
  
    private Set<ClassRef> getHelperClassesParticipatingInLibrarySuperType() {
      Set<ClassRef> helperClassesParticipatingInLibrarySuperType = new HashSet<>();
      for (ClassRef reference : getHelperClassesWithLibrarySuperType()) {
        addSuperTypesThatAreAlsoHelperClasses(
            reference.getClassName(), helperClassesParticipatingInLibrarySuperType);
      }
      return helperClassesParticipatingInLibrarySuperType;
    }
  
    private Set<ClassRef> getHelperClassesWithLibrarySuperType() {
      Set<ClassRef> helperClassesWithLibrarySuperType = new HashSet<>();
      for (ClassRef reference : references.values()) {
        if (helperClassPredicate.isHelperClass(reference.getClassName())
            && hasLibrarySuperType(reference.getClassName())) {
          helperClassesWithLibrarySuperType.add(reference);
        }
      }
      return helperClassesWithLibrarySuperType;
    }
  
    private void addSuperTypesThatAreAlsoHelperClasses(
        @Nullable String className, Set<ClassRef> superTypes) {
      if (className != null && helperClassPredicate.isHelperClass(className)) {
        ClassRef reference = references.get(className);
        superTypes.add(reference);
  
        addSuperTypesThatAreAlsoHelperClasses(reference.getSuperClassName(), superTypes);
        // need to keep interfaces too since they may have default methods
        for (String superType : reference.getInterfaceNames()) {
          addSuperTypesThatAreAlsoHelperClasses(superType, superTypes);
        }
      }
    }
  
    private boolean hasLibrarySuperType(@Nullable String typeName) {
      if (typeName == null || typeName.startsWith("java.")) {
        return false;
      }
      if (helperClassPredicate.isLibraryClass(typeName)) {
        return true;
      }
      ClassRef reference = references.get(typeName);
      if (hasLibrarySuperType(reference.getSuperClassName())) {
        return true;
      }
      for (String type : reference.getInterfaceNames()) {
        if (hasLibrarySuperType(type)) {
          return true;
        }
      }
      return false;
    }
  
    // see https://en.wikipedia.org/wiki/Topological_sorting#Kahn's_algorithm
    public List<String> getSortedHelperClasses() {
      MutableGraph<String> dependencyGraph = Graphs.copyOf(Graphs.transpose(helperSuperClassGraph));
      List<String> helperClasses = new ArrayList<>(dependencyGraph.nodes().size());
  
      Queue<String> helpersWithNoDeps = findAllHelperClassesWithoutDependencies(dependencyGraph);
  
      while (!helpersWithNoDeps.isEmpty()) {
        String helperClass = helpersWithNoDeps.remove();
        helperClasses.add(helperClass);
  
        Set<String> dependencies = new HashSet<>(dependencyGraph.successors(helperClass));
        for (String dependency : dependencies) {
          dependencyGraph.removeEdge(helperClass, dependency);
          if (dependencyGraph.predecessors(dependency).isEmpty()) {
            helpersWithNoDeps.add(dependency);
          }
        }
      }
  
      return helperClasses;
    }
  
    private static Queue<String> findAllHelperClassesWithoutDependencies(
        Graph<String> dependencyGraph) {
      Queue<String> helpersWithNoDeps = new LinkedList<>();
      for (String helperClass : dependencyGraph.nodes()) {
        if (dependencyGraph.predecessors(helperClass).isEmpty()) {
          helpersWithNoDeps.add(helperClass);
        }
      }
      return helpersWithNoDeps;
    }
  
    public VirtualFieldMappings getVirtualFieldMappings() {
      return virtualFieldMappingsBuilder.build();
    }
  }
  