content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.testing.bytebuddy;
  
  import static java.util.logging.Level.SEVERE;
  
  import io.opentelemetry.javaagent.extension.ignore.IgnoredTypesConfigurer;
  import io.opentelemetry.javaagent.tooling.EmptyConfigProperties;
  import io.opentelemetry.javaagent.tooling.SafeServiceLoader;
  import io.opentelemetry.javaagent.tooling.Utils;
  import io.opentelemetry.javaagent.tooling.ignore.AdditionalLibraryIgnoredTypesConfigurer;
  import io.opentelemetry.javaagent.tooling.ignore.GlobalIgnoredTypesConfigurer;
  import io.opentelemetry.javaagent.tooling.ignore.IgnoreAllow;
  import io.opentelemetry.javaagent.tooling.ignore.IgnoredTypesBuilderImpl;
  import io.opentelemetry.javaagent.tooling.instrumentation.MuzzleFailureCounter;
  import io.opentelemetry.javaagent.tooling.util.Trie;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.List;
  import java.util.Set;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.function.BiFunction;
  import java.util.function.Function;
  import java.util.logging.Logger;
  import net.bytebuddy.agent.builder.AgentBuilder;
  import net.bytebuddy.description.type.TypeDescription;
  import net.bytebuddy.dynamic.DynamicType;
  import net.bytebuddy.utility.JavaModule;
  
  public class TestAgentListener implements AgentBuilder.Listener {
  
    private static final Logger logger = Logger.getLogger(TestAgentListener.class.getName());
  
    private static final Trie<IgnoreAllow> ADDITIONAL_LIBRARIES_TRIE;
    private static final Trie<IgnoreAllow> OTHER_IGNORES_TRIE;
  
    static {
      ADDITIONAL_LIBRARIES_TRIE = buildAdditionalLibraryIgnores();
      OTHER_IGNORES_TRIE = buildOtherConfiguredIgnores();
    }
  
    private static Trie<IgnoreAllow> buildAdditionalLibraryIgnores() {
      IgnoredTypesBuilderImpl builder = new IgnoredTypesBuilderImpl();
      new AdditionalLibraryIgnoredTypesConfigurer().configure(builder);
      return builder.buildIgnoredTypesTrie();
    }
  
    private static Trie<IgnoreAllow> buildOtherConfiguredIgnores() {
      IgnoredTypesBuilderImpl builder = new IgnoredTypesBuilderImpl();
      for (IgnoredTypesConfigurer configurer :
          SafeServiceLoader.loadOrdered(
              IgnoredTypesConfigurer.class, Utils.getExtensionsClassLoader())) {
        // skip built-in agent ignores
        if (configurer instanceof AdditionalLibraryIgnoredTypesConfigurer
            || configurer instanceof GlobalIgnoredTypesConfigurer) {
          continue;
        }
        configurer.configure(builder, EmptyConfigProperties.INSTANCE);
      }
      return builder.buildIgnoredTypesTrie();
    }
  
    public static void reset() {
      INSTANCE.transformedClassesNames.clear();
      INSTANCE.instrumentationErrorCount.set(0);
      INSTANCE.skipTransformationConditions.clear();
      INSTANCE.skipErrorConditions.clear();
    }
  
    public static int getInstrumentationErrorCount() {
      return INSTANCE.instrumentationErrorCount.get();
    }
  
    public static int getAndResetMuzzleFailureCount() {
      return MuzzleFailureCounter.getAndReset();
    }
  
    public static List<String> getIgnoredButTransformedClassNames() {
      List<String> names = new ArrayList<>();
      for (String name : INSTANCE.transformedClassesNames) {
        // only record those types that weren't explicitly marked as either ignored or allowed by the
        // instrumentation authors
        if (ADDITIONAL_LIBRARIES_TRIE.getOrNull(name) == IgnoreAllow.IGNORE
            && OTHER_IGNORES_TRIE.getOrNull(name) == null) {
          names.add(name);
        }
      }
      return names;
    }
  
    public static void addSkipTransformationCondition(Function<String, Boolean> condition) {
      INSTANCE.skipTransformationConditions.add(condition);
    }
  
    public static void addSkipErrorCondition(BiFunction<String, Throwable, Boolean> condition) {
      INSTANCE.skipErrorConditions.add(condition);
    }
  
    static final TestAgentListener INSTANCE = new TestAgentListener();
  
    private final Set<String> transformedClassesNames =
        Collections.newSetFromMap(new ConcurrentHashMap<>());
    private final AtomicInteger instrumentationErrorCount = new AtomicInteger(0);
    private final Set<Function<String, Boolean>> skipTransformationConditions =
        Collections.newSetFromMap(new ConcurrentHashMap<>());
    private final Set<BiFunction<String, Throwable, Boolean>> skipErrorConditions =
        Collections.newSetFromMap(new ConcurrentHashMap<>());
  
    @Override
    public void onDiscovery(
        String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {
      for (Function<String, Boolean> skipCondition : skipTransformationConditions) {
        if (skipCondition.apply(typeName)) {
          throw new AbortTransformationException(
              "Aborting transform for class name = " + typeName + ", loader = " + classLoader);
        }
      }
    }
  
    @Override
    public void onTransformation(
        TypeDescription typeDescription,
        ClassLoader classLoader,
        JavaModule module,
        boolean loaded,
        DynamicType dynamicType) {
      transformedClassesNames.add(typeDescription.getActualName());
    }
  
    @Override
    public void onIgnored(
        TypeDescription typeDescription,
        ClassLoader classLoader,
        JavaModule module,
        boolean loaded) {}
  
    @Override
    public void onError(
        String typeName,
        ClassLoader classLoader,
        JavaModule module,
        boolean loaded,
        Throwable throwable) {
      for (BiFunction<String, Throwable, Boolean> condition : skipErrorConditions) {
        if (condition.apply(typeName, throwable)) {
          return;
        }
      }
      if (!(throwable instanceof AbortTransformationException)) {
        logger.log(
            SEVERE,
            "Unexpected instrumentation error when instrumenting " + typeName + " on " + classLoader,
            throwable);
        instrumentationErrorCount.incrementAndGet();
      }
    }
  
    @Override
    public void onComplete(
        String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {}
  
    /** Used to signal that a transformation was intentionally aborted and is not an error. */
    private static class AbortTransformationException extends RuntimeException {
      private static final long serialVersionUID = 1L;
  
      public AbortTransformationException(String message) {
        super(message);
      }
    }
  }
  