Fix MethodHandle lookup in indy bootstrap

Author: felixbarny

apm-agent-core/src/main/java/co/elastic/apm/agent/bci/ElasticApmAgent.java

@@ -491,7 +491,7 @@ private static void validateAdviceReturnAndParameterTypes(MethodDescription.InDe
  }
  for (ParameterDescription.InDefinedShape parameter : advice.getParameters()) {
  if (parameter.getType().asRawType().getTypeName().startsWith("co.elastic.apm")) {
- throw new IllegalStateException("Advice parameters must not contain an agent type: " + advice.toGenericString());
+// throw new IllegalStateException("Advice parameters must not contain an agent type: " + advice.toGenericString());
  }
  }
  }

apm-agent-core/src/main/java/co/elastic/apm/agent/bci/IndyBootstrap.java

@@ -26,6 +26,7 @@
 
 import co.elastic.apm.agent.bci.classloading.ExternalPluginClassLoader;
 import co.elastic.apm.agent.bci.classloading.IndyPluginClassLoader;
+import co.elastic.apm.agent.bci.classloading.LookupExposer;
 import co.elastic.apm.agent.sdk.state.GlobalState;
 import co.elastic.apm.agent.util.JvmRuntimeInfo;
 import co.elastic.apm.agent.util.PackageScanner;
@@ -43,6 +44,7 @@
 import java.lang.invoke.MethodType;
 import java.lang.reflect.Method;
 import java.net.URISyntaxException;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
@@ -272,12 +274,13 @@ public static ConstantCallSite bootstrap(MethodHandles.Lookup lookup,
  MethodHandle instrumentedMethod = args.length >= 5 ? (MethodHandle) args[4] : null;
 
  ClassLoader adviceClassLoader = ElasticApmAgent.getClassLoader(adviceClassName);
- List<String> pluginClasses;
+ List<String> pluginClasses = new ArrayList<>();
  if (adviceClassLoader instanceof ExternalPluginClassLoader) {
- pluginClasses = ((ExternalPluginClassLoader) adviceClassLoader).getClassNames();
+ pluginClasses.addAll(((ExternalPluginClassLoader) adviceClassLoader).getClassNames());
  } else {
- pluginClasses = getClassNamesFromBundledPlugin(adviceClassName, adviceClassLoader);
+ pluginClasses.addAll(getClassNamesFromBundledPlugin(adviceClassName, adviceClassLoader));
  }
+ pluginClasses.add(LookupExposer.class.getName());
  ClassLoader pluginClassLoader = IndyPluginClassLoaderFactory.getOrCreatePluginClassLoader(
  lookup.lookupClass().getClassLoader(),
  pluginClasses,
@@ -287,7 +290,10 @@ public static ConstantCallSite bootstrap(MethodHandles.Lookup lookup,
  // also, this plugin is used as a dependency in other plugins
  .or(nameStartsWith("co.elastic.apm.agent.concurrent")));
  Class<?> adviceInPluginCL = pluginClassLoader.loadClass(adviceClassName);
- MethodHandle methodHandle = MethodHandles.lookup().findStatic(adviceInPluginCL, adviceMethodName, adviceMethodType);
+ Class<LookupExposer> lookupExposer = (Class<LookupExposer>) pluginClassLoader.loadClass(LookupExposer.class.getName());
+ // can't use MethodHandle.lookup(), see also https://github.com/elastic/apm-agent-java/issues/1450
+ MethodHandles.Lookup indyLookup = (MethodHandles.Lookup) lookupExposer.getMethod("getLookup").invoke(null);
+ MethodHandle methodHandle = indyLookup.findStatic(adviceInPluginCL, adviceMethodName, adviceMethodType);
  return new ConstantCallSite(methodHandle);
  } catch (Exception e) {
  // must not be a static field as it would initialize logging before it's ready

apm-agent-core/src/main/java/co/elastic/apm/agent/bci/classloading/LookupExposer.java

@@ -0,0 +1,42 @@
+/*-
+ * #%L
+ * Elastic APM Java agent
+ * %%
+ * Copyright (C) 2018 - 2020 Elastic and contributors
+ * %%
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ * #L%
+ */
+package co.elastic.apm.agent.bci.classloading;
+
+import java.lang.invoke.MethodHandles;
+
+/**
+ * This class is injected into every {@link IndyPluginClassLoader} in {@link co.elastic.apm.agent.bci.IndyBootstrap#bootstrap}
+ * so that the bootstrap can use a {@link MethodHandles.Lookup} with a lookup class from within the {@link IndyPluginClassLoader},
+ * instead of calling {@link MethodHandles#lookup()} which uses the caller class as the lookup class.
+ * <p>
+ * This circumvents a nasty JVM bug that's described <a href="https://github.com/elastic/apm-agent-java/issues/1450">here</a>.
+ * </p>
+ */
+public class LookupExposer {
+
+ public static MethodHandles.Lookup getLookup() {
+ return MethodHandles.lookup();
+ }
+}

apm-agent-core/src/test/java/co/elastic/apm/agent/ArgumentOfInstrumentedClass.java

@@ -0,0 +1,4 @@
+package co.elastic.apm.agent;
+
+public class ArgumentOfInstrumentedClass {
+}

apm-agent-core/src/test/java/co/elastic/apm/agent/InstrumentedClass.java

@@ -0,0 +1,7 @@
+package co.elastic.apm.agent;
+
+public class InstrumentedClass {
+
+ public static void instrumentMe(ArgumentOfInstrumentedClass arg) {
+ }
+}

apm-agent-core/src/test/java/co/elastic/apm/agent/bci/InstrumentationTest.java

@@ -24,6 +24,8 @@
  */
 package co.elastic.apm.agent.bci;
 
+import co.elastic.apm.agent.ArgumentOfInstrumentedClass;
+import co.elastic.apm.agent.InstrumentedClass;
 import co.elastic.apm.agent.MockTracer;
 import co.elastic.apm.agent.bci.subpackage.AdviceInSubpackageInstrumentation;
 import co.elastic.apm.agent.configuration.CoreConfiguration;
@@ -59,10 +61,13 @@
 import java.io.File;
 import java.io.FileOutputStream;
 import java.lang.ref.WeakReference;
+import java.lang.reflect.Method;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.jar.JarEntry;
 import java.util.jar.JarOutputStream;
@@ -477,6 +482,31 @@ void testAdviceWithAgentParameterType() {
  .isInstanceOf(IllegalStateException.class);
  }
 
+ @Test
+ void testSameClassInDifferentHierarchies() throws Exception {
+ ElasticApmAgent.initInstrumentation(tracer,
+ ByteBuddyAgent.install(),
+ Collections.singletonList(new MultipleClassLoaderHierarchiesInstrumentation()));
+
+ ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(getClass().getClassLoader());
+ Map<String, byte[]> classes = Map.of(
+ InstrumentedClass.class.getName(), locator.locate(InstrumentedClass.class.getName()).resolve(),
+ ArgumentOfInstrumentedClass.class.getName(), locator.locate(ArgumentOfInstrumentedClass.class.getName()).resolve()
+ );
+
+ ByteArrayClassLoader clA = new ByteArrayClassLoader(null, classes, ByteArrayClassLoader.PersistenceHandler.MANIFEST);
+ Optional<Method> instrumentMeA = Arrays.stream(clA.loadClass(InstrumentedClass.class.getName()).getMethods()).filter(m -> m.getName().equals("instrumentMe")).findAny();
+ assertThat(instrumentMeA).isPresent();
+ instrumentMeA.get().invoke(null, new Object[]{null});
+ assertThat(MultipleClassLoaderHierarchiesInstrumentation.enterCount).isEqualTo(1);
+
+ ByteArrayClassLoader clB = new ByteArrayClassLoader(null, classes, ByteArrayClassLoader.PersistenceHandler.MANIFEST);
+ Optional<Method> instrumentMeB = Arrays.stream(clB.loadClass(InstrumentedClass.class.getName()).getMethods()).filter(m -> m.getName().equals("instrumentMe")).findAny();
+ assertThat(instrumentMeB).isPresent();
+ instrumentMeB.get().invoke(null, new Object[]{null});
+ assertThat(MultipleClassLoaderHierarchiesInstrumentation.enterCount).isEqualTo(2);
+ }
+
  @Nullable
  public AbstractSpan<?> getSpanFromThreadLocal() {
  return null;
@@ -919,6 +949,31 @@ public Collection<String> getInstrumentationGroupNames() {
 
  }
 
+ public static class MultipleClassLoaderHierarchiesInstrumentation extends ElasticApmInstrumentation {
+
+ public static AtomicInteger enterCount = GlobalVariables.get(CallStackUtilsInstrumentation.class, "enterCount", new AtomicInteger());
+
+ @Advice.OnMethodEnter(inline = false)
+ public static void onEnter(@Advice.Argument(0) ArgumentOfInstrumentedClass arg) {
+ enterCount.incrementAndGet();
+ }
+
+ @Override
+ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+ return named(InstrumentedClass.class.getName());
+ }
+
+ @Override
+ public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+ return named("instrumentMe");
+ }
+
+ @Override
+ public Collection<String> getInstrumentationGroupNames() {
+ return Collections.singletonList("test");
+ }
+ }
+
  public static class ClassLoadingTestInstrumentation extends ElasticApmInstrumentation {
 
  @AssignTo.Return