From 21bcc314c44632c720137307bae77ce8d300d4a4 Mon Sep 17 00:00:00 2001 From: Ryan Ernst Date: Thu, 19 Dec 2024 08:59:23 -0800 Subject: [PATCH] Limit which classes are retransformed (#118786) During entitlements initialization the transformer is added to instrumenation after some classes are already loaded. Currently we end up force loading (though not initializing) all classes that want to transform. This commit simplifies the retransform to only apply to classes which we know are already loaded by the jdk, which Instrumentation provides. --- .../EntitlementInitialization.java | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java index 8e4cddc4d63e..c2ee935e0e5f 100644 --- a/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java +++ b/libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java @@ -34,6 +34,7 @@ import java.lang.reflect.InvocationTargetException; import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.StandardOpenOption; +import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; import java.util.List; @@ -72,17 +73,17 @@ public class EntitlementInitialization { Instrumenter instrumenter = INSTRUMENTER_FACTORY.newInstrumenter(EntitlementChecker.class, checkMethods); inst.addTransformer(new Transformer(instrumenter, classesToTransform), true); - // TODO: should we limit this array somehow? - var classesToRetransform = classesToTransform.stream().map(EntitlementInitialization::internalNameToClass).toArray(Class[]::new); - inst.retransformClasses(classesToRetransform); + inst.retransformClasses(findClassesToRetransform(inst.getAllLoadedClasses(), classesToTransform)); } - private static Class internalNameToClass(String internalName) { - try { - return Class.forName(internalName.replace('/', '.'), false, ClassLoader.getPlatformClassLoader()); - } catch (ClassNotFoundException e) { - throw new RuntimeException(e); + private static Class[] findClassesToRetransform(Class[] loadedClasses, Set classesToTransform) { + List> retransform = new ArrayList<>(); + for (Class loadedClass : loadedClasses) { + if (classesToTransform.contains(loadedClass.getName().replace(".", "/"))) { + retransform.add(loadedClass); + } } + return retransform.toArray(new Class[0]); } private static PolicyManager createPolicyManager() throws IOException {