DefineClassHandler.java

/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package io.opentelemetry.javaagent.tooling;

import static java.nio.charset.StandardCharsets.ISO_8859_1;
import static java.util.Collections.emptySet;

import io.opentelemetry.javaagent.bootstrap.DefineClassHelper.Handler;
import java.util.HashSet;
import java.util.Set;
import javax.annotation.Nullable;
import org.objectweb.asm.ClassReader;

public class DefineClassHandler implements Handler {
  public static final DefineClassHandler INSTANCE = new DefineClassHandler();
  private static final ThreadLocal<DefineClassContextImpl> defineClassContext =
      ThreadLocal.withInitial(() -> DefineClassContextImpl.NOP);

  private DefineClassHandler() {}

  @Override
  public DefineClassContext beforeDefineClass(
      ClassLoader classLoader, String className, byte[] classBytes, int offset, int length) {
    // with OpenJ9 class data sharing we don't get real class bytes
    if (classBytes == null
        || (classBytes.length == 40
            && new String(classBytes, ISO_8859_1).startsWith("J9ROMCLASSCOOKIE"))) {
      return DefineClassContextImpl.NOP;
    }

    Set<String> superNames = new HashSet<>();
    DefineClassContextImpl context = DefineClassContextImpl.enter();
    // attempt to load super types of currently loaded class
    // for a class to be loaded all of its super types must be loaded, here we just change the order
    // of operations and load super types before transforming the bytes for current class so that
    // we could use these super types for resolving the advice that needs to be applied to current
    // class
    try {
      ClassReader cr = new ClassReader(classBytes, offset, length);
      String superName = cr.getSuperName();
      if (superName != null) {
        String superDotName = superName.replace('/', '.');
        Class<?> clazz = Class.forName(superDotName, false, classLoader);
        addSuperNames(superNames, clazz);
      }
      String[] interfaces = cr.getInterfaces();
      for (String interfaceName : interfaces) {
        String interfaceDotName = interfaceName.replace('/', '.');
        Class<?> clazz = Class.forName(interfaceDotName, false, classLoader);
        addSuperNames(superNames, clazz);
      }
      context.superDotNames = superNames;
    } catch (Throwable ignored) {
      // loading of super class or interface failed
      // mark current class as failed to skip matching and transforming it
      // we'll let defining the class proceed as usual so that it would throw the same exception as
      // it does when running without the agent
      context.failedClassDotName = className;
    }

    return context;
  }

  @Override
  public DefineClassContext beforeDefineLambdaClass(Class<?> lambdaInterface) {
    DefineClassContextImpl context = DefineClassContextImpl.enter();
    Set<String> superNames = new HashSet<>();
    addSuperNames(superNames, lambdaInterface);
    context.superDotNames = superNames;

    return context;
  }

  private static void addSuperNames(Set<String> superNames, Class<?> clazz) {
    if (clazz == null || !superNames.add(clazz.getName())) {
      return;
    }
    addSuperNames(superNames, clazz.getSuperclass());
    for (Class<?> interfaceClass : clazz.getInterfaces()) {
      addSuperNames(superNames, interfaceClass);
    }
  }

  @Override
  public void afterDefineClass(DefineClassContext context) {
    context.exit();
  }

  /**
   * Detect whether loading the specified class is known to fail.
   *
   * @param dotClassName class being loaded
   * @return true if it is known that loading class with given name will fail
   */
  public static boolean isFailedClass(String dotClassName) {
    DefineClassContextImpl context = defineClassContext.get();
    return context.failedClassDotName != null && context.failedClassDotName.equals(dotClassName);
  }

  public static Set<String> getSuperTypes() {
    Set<String> superNames = defineClassContext.get().superDotNames;
    return superNames == null ? emptySet() : superNames;
  }

  private static class DefineClassContextImpl implements DefineClassContext {
    // NOP is returned from beforeDefineClass without calling enter() (no frame pushed),
    // so exit() must not mutate the ThreadLocal — otherwise a nested J9 defineClass would
    // clobber an outer frame that is still active.
    private static final DefineClassContextImpl NOP =
        new DefineClassContextImpl() {
          @Override
          public void exit() {}
        };

    @Nullable private final DefineClassContextImpl previous;
    @Nullable String failedClassDotName;
    @Nullable Set<String> superDotNames;

    private DefineClassContextImpl() {
      previous = null;
    }

    private DefineClassContextImpl(@Nullable DefineClassContextImpl previous) {
      this.previous = previous;
    }

    static DefineClassContextImpl enter() {
      DefineClassContextImpl previous = defineClassContext.get();
      DefineClassContextImpl context = new DefineClassContextImpl(previous);
      defineClassContext.set(context);
      return context;
    }

    @Override
    public void exit() {
      defineClassContext.set(previous != null ? previous : NOP);
    }
  }
}