content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.tooling.muzzle;
  
  import static java.util.Collections.emptyList;
  import static java.util.Collections.singletonList;
  
  import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
  import io.opentelemetry.javaagent.tooling.muzzle.references.ClassRef;
  import io.opentelemetry.javaagent.tooling.muzzle.references.FieldRef;
  import io.opentelemetry.javaagent.tooling.muzzle.references.Flag;
  import io.opentelemetry.javaagent.tooling.muzzle.references.MethodRef;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.function.Predicate;
  import java.util.stream.Collectors;
  import net.bytebuddy.description.field.FieldDescription;
  import net.bytebuddy.description.method.MethodDescription;
  import net.bytebuddy.description.type.TypeDescription;
  import net.bytebuddy.pool.TypePool;
  import org.objectweb.asm.Type;
  
  /** Matches a set of references against a class loader. */
  public final class ReferenceMatcher {
  
    private final Map<String, ClassRef> references;
    private final Set<String> helperClassNames;
    private final HelperClassPredicate helperClassPredicate;
  
    public static ReferenceMatcher of(InstrumentationModule instrumentationModule) {
      return new ReferenceMatcher(
          InstrumentationModuleMuzzle.getHelperClassNames(instrumentationModule),
          InstrumentationModuleMuzzle.getMuzzleReferences(instrumentationModule),
          instrumentationModule::isHelperClass);
    }
  
    ReferenceMatcher(
        List<String> helperClassNames,
        Map<String, ClassRef> references,
        Predicate<String> libraryInstrumentationPredicate) {
      this.references = references;
      this.helperClassNames = new HashSet<>(helperClassNames);
      this.helperClassPredicate = new HelperClassPredicate(libraryInstrumentationPredicate);
    }
  
    /**
     * Matcher used by ByteBuddy. Caller is expected to cache the result if this method is called
     * multiple times for given class loader.
     *
     * @param loader Classloader to validate against (cannot be {@code null}, must pass "bootstrap
     *     proxy" instead of bootstrap class loader)
     * @return true if all references match the classpath of loader
     */
    public boolean matches(ClassLoader loader) {
      TypePool typePool = createTypePool(loader);
      for (ClassRef reference : references.values()) {
        if (!checkMatch(reference, typePool, loader).isEmpty()) {
          return false;
        }
      }
      return true;
    }
  
    /**
     * Loads the full list of mismatches. Used in debug contexts only.
     *
     * @param loader Classloader to validate against (cannot be {@code null}, must pass "bootstrap *
     *     proxy" instead of bootstrap class loader)
     * @return A list of all mismatches between this ReferenceMatcher and loader's classpath.
     */
    public List<Mismatch> getMismatchedReferenceSources(ClassLoader loader) {
      TypePool typePool = createTypePool(loader);
  
      List<Mismatch> mismatches = emptyList();
  
      for (ClassRef reference : references.values()) {
        mismatches = addAll(mismatches, checkMatch(reference, typePool, loader));
      }
  
      return mismatches;
    }
  
    // loader cannot be null, must pass "bootstrap proxy" instead of bootstrap class loader
    private static TypePool createTypePool(ClassLoader loader) {
      // ok to use locationStrategy() without fallback bootstrap proxy here since loader is non-null
      return AgentTooling.poolStrategy()
          .typePool(AgentTooling.locationStrategy().classFileLocator(loader), loader);
    }
  
    /**
     * Check a reference against a class loader's classpath.
     *
     * @return A list of mismatched sources. A list of size 0 means the reference matches the class.
     */
    private List<Mismatch> checkMatch(ClassRef reference, TypePool typePool, ClassLoader loader) {
      try {
        if (helperClassPredicate.isHelperClass(reference.getClassName())) {
          // make sure helper class is registered
          if (!helperClassNames.contains(reference.getClassName())) {
            return singletonList(new Mismatch.MissingClass(reference));
          }
          // helper classes get their own check: whether they implement all abstract methods
          return checkHelperClassMatch(reference, typePool);
        } else {
          TypePool.Resolution resolution = typePool.describe(reference.getClassName());
          if (!resolution.isResolved()) {
            return singletonList(new Mismatch.MissingClass(reference));
          }
          return checkThirdPartyTypeMatch(reference, resolution.resolve());
        }
      } catch (RuntimeException e) {
        if (e.getMessage().startsWith("Cannot resolve type description for ")) {
          // bytebuddy throws an illegal state exception with this message if it cannot resolve types
          // TODO: handle missing type resolutions without catching bytebuddy's exceptions
          String className = e.getMessage().replace("Cannot resolve type description for ", "");
          return singletonList(new Mismatch.MissingClass(reference, className));
        } else {
          // Shouldn't happen. Fail the reference check and add a mismatch for debug logging.
          return singletonList(new Mismatch.ReferenceCheckError(e, reference, loader));
        }
      }
    }
  
    // for helper classes we make sure that all abstract methods from super classes and interfaces are
    // implemented and that all accessed fields are defined somewhere in the type hierarchy
    private List<Mismatch> checkHelperClassMatch(ClassRef helperClass, TypePool typePool) {
      List<Mismatch> mismatches = emptyList();
  
      HelperReferenceWrapper helperWrapper =
          new HelperReferenceWrapper.Factory(typePool, references, helperClassPredicate)
              .create(helperClass);
  
      Set<HelperReferenceWrapper.Field> undeclaredFields =
          helperClass.getFields().stream()
              .filter(f -> !f.isDeclared())
              .map(f -> new HelperReferenceWrapper.Field(f.getName(), f.getDescriptor()))
              .collect(Collectors.toSet());
  
      // if there are any fields in this helper class that's not declared here, check the type
      // hierarchy
      if (!undeclaredFields.isEmpty()) {
        Set<HelperReferenceWrapper.Field> superClassFields = new HashSet<>();
        collectFieldsFromTypeHierarchy(helperWrapper, superClassFields);
  
        undeclaredFields.removeAll(superClassFields);
        for (HelperReferenceWrapper.Field missingField : undeclaredFields) {
          mismatches = add(mismatches, new Mismatch.MissingField(helperClass, missingField));
        }
      }
  
      // skip abstract method check if this type does not have super type or is abstract
      if (!helperWrapper.hasSuperTypes() || helperWrapper.isAbstract()) {
        return mismatches;
      }
  
      // treat the helper type as a bag of methods: collect all methods defined in the helper class,
      // all superclasses and interfaces and check if all abstract methods are implemented somewhere
      Set<HelperReferenceWrapper.Method> abstractMethods = new HashSet<>();
      Set<HelperReferenceWrapper.Method> plainMethods = new HashSet<>();
      collectMethodsFromTypeHierarchy(helperWrapper, abstractMethods, plainMethods);
  
      abstractMethods.removeAll(plainMethods);
      for (HelperReferenceWrapper.Method unimplementedMethod : abstractMethods) {
        mismatches = add(mismatches, new Mismatch.MissingMethod(helperClass, unimplementedMethod));
      }
  
      return mismatches;
    }
  
    private static void collectFieldsFromTypeHierarchy(
        HelperReferenceWrapper type, Set<HelperReferenceWrapper.Field> fields) {
  
      type.getFields().forEach(fields::add);
      type.getSuperTypes().forEach(superType -> collectFieldsFromTypeHierarchy(superType, fields));
    }
  
    private static void collectMethodsFromTypeHierarchy(
        HelperReferenceWrapper type,
        Set<HelperReferenceWrapper.Method> abstractMethods,
        Set<HelperReferenceWrapper.Method> plainMethods) {
  
      type.getMethods()
          .forEach(method -> (method.isAbstract() ? abstractMethods : plainMethods).add(method));
  
      type.getSuperTypes()
          .forEach(
              superType -> collectMethodsFromTypeHierarchy(superType, abstractMethods, plainMethods));
    }
  
    private static List<Mismatch> checkThirdPartyTypeMatch(
        ClassRef reference, TypeDescription typeOnClasspath) {
      List<Mismatch> mismatches = Collections.emptyList();
  
      for (Flag flag : reference.getFlags()) {
        if (!flag.matches(typeOnClasspath.getActualModifiers(false))) {
          String desc = reference.getClassName();
          mismatches =
              add(
                  mismatches,
                  new Mismatch.MissingFlag(
                      reference.getSources(), desc, flag, typeOnClasspath.getActualModifiers(false)));
        }
      }
  
      for (FieldRef fieldRef : reference.getFields()) {
        FieldDescription.InDefinedShape fieldDescription = findField(fieldRef, typeOnClasspath);
        if (fieldDescription == null) {
          mismatches = add(mismatches, new Mismatch.MissingField(reference, fieldRef));
        } else {
          for (Flag flag : fieldRef.getFlags()) {
            if (!flag.matches(fieldDescription.getModifiers())) {
              String desc =
                  reference.getClassName()
                      + "#"
                      + fieldRef.getName()
                      + Type.getType(fieldRef.getDescriptor()).getInternalName();
              mismatches =
                  add(
                      mismatches,
                      new Mismatch.MissingFlag(
                          fieldRef.getSources(), desc, flag, fieldDescription.getModifiers()));
            }
          }
        }
      }
  
      for (MethodRef methodRef : reference.getMethods()) {
        MethodDescription.InDefinedShape methodDescription = findMethod(methodRef, typeOnClasspath);
        if (methodDescription == null) {
          mismatches = add(mismatches, new Mismatch.MissingMethod(reference, methodRef));
        } else {
          for (Flag flag : methodRef.getFlags()) {
            if (!flag.matches(methodDescription.getModifiers())) {
              String desc =
                  reference.getClassName() + "#" + methodRef.getName() + methodRef.getDescriptor();
              mismatches =
                  add(
                      mismatches,
                      new Mismatch.MissingFlag(
                          methodRef.getSources(), desc, flag, methodDescription.getModifiers()));
            }
          }
        }
      }
      return mismatches;
    }
  
    private static FieldDescription.InDefinedShape findField(
        FieldRef fieldRef, TypeDescription typeOnClasspath) {
      for (FieldDescription.InDefinedShape fieldType : typeOnClasspath.getDeclaredFields()) {
        if (fieldType.getName().equals(fieldRef.getName())
            && fieldType.getDescriptor().equals(fieldRef.getDescriptor())) {
          return fieldType;
        }
      }
      if (typeOnClasspath.getSuperClass() != null) {
        FieldDescription.InDefinedShape fieldOnSupertype =
            findField(fieldRef, typeOnClasspath.getSuperClass().asErasure());
        if (fieldOnSupertype != null) {
          return fieldOnSupertype;
        }
      }
      for (TypeDescription.Generic interfaceType : typeOnClasspath.getInterfaces()) {
        FieldDescription.InDefinedShape fieldOnSupertype =
            findField(fieldRef, interfaceType.asErasure());
        if (fieldOnSupertype != null) {
          return fieldOnSupertype;
        }
      }
      return null;
    }
  
    private static MethodDescription.InDefinedShape findMethod(
        MethodRef methodRef, TypeDescription typeOnClasspath) {
  
      for (MethodDescription.InDefinedShape methodDescription :
          typeOnClasspath.getDeclaredMethods()) {
        if (methodDescription.getInternalName().equals(methodRef.getName())
            && methodDescription.getDescriptor().equals(methodRef.getDescriptor())) {
          return methodDescription;
        }
      }
  
      // if the method we're looking for is a constructor, we're only checking the direct super type;
      // and skipping all the interfaces and indirect superclasses
      if (methodRef.isConstructor()) {
        return null;
      }
  
      TypeDescription superType = null;
      if (typeOnClasspath.getSuperClass() != null) {
        superType = typeOnClasspath.getSuperClass().asErasure();
      } else if (!"java.lang.Object".equals(typeOnClasspath.getName())) {
        superType = TypeDescription.ForLoadedType.of(Object.class);
      }
      if (superType != null) {
        MethodDescription.InDefinedShape methodOnSupertype = findMethod(methodRef, superType);
        if (methodOnSupertype != null) {
          return methodOnSupertype;
        }
      }
  
      for (TypeDescription.Generic interfaceType : typeOnClasspath.getInterfaces()) {
        MethodDescription.InDefinedShape methodOnSupertype =
            findMethod(methodRef, interfaceType.asErasure());
        if (methodOnSupertype != null) {
          return methodOnSupertype;
        }
      }
      return null;
    }
  
    // optimization to avoid ArrayList allocation in the common case when there are no mismatches
    static List<Mismatch> add(List<Mismatch> mismatches, Mismatch mismatch) {
      List<Mismatch> result = mismatches.isEmpty() ? new ArrayList<>() : mismatches;
      result.add(mismatch);
      return result;
    }
  
    // optimization to avoid ArrayList allocation in the common case when there are no mismatches
    static List<Mismatch> addAll(List<Mismatch> mismatches, List<Mismatch> toAdd) {
      if (!toAdd.isEmpty()) {
        List<Mismatch> result = mismatches.isEmpty() ? new ArrayList<>() : mismatches;
        result.addAll(toAdd);
        return result;
      }
      return mismatches;
    }
  }
  