Fully initialize policy checker before instrumenting (elastic#128703)

Entitlement instrumentation works by reflectively calling back into the
entitlements lib to grab the checker. It must be fully in place before
any classes are instrumented. This commit fixes a bug that was
introduced by refactoring which caused the checker to not be set until
after all classes were instrumented. In some situations this could lead
the checker to being null when it is grab (and statically cached) by the
entitlement bridge.

Author: rjernst

libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java

@@ -68,7 +68,9 @@ public static EntitlementChecker checker() {
      * @param inst the JVM instrumentation class instance
      */
     public static void initialize(Instrumentation inst) throws Exception {
-        checker = initChecker(inst, createPolicyManager());
+        // the checker _MUST_ be set before _any_ instrumentation is done
+        checker = initChecker(createPolicyManager());
+        initInstrumentation(inst);
     }
 
     /**
@@ -146,7 +148,7 @@ private static void ensureClassesSensitiveToVerificationAreInitialized() {
         }
     }
 
-    static ElasticsearchEntitlementChecker initChecker(Instrumentation inst, PolicyManager policyManager) throws Exception {
+    static ElasticsearchEntitlementChecker initChecker(PolicyManager policyManager) {
         final PolicyChecker policyChecker = createPolicyChecker(policyManager);
         final Class<?> clazz = EntitlementCheckerUtils.getVersionSpecificCheckerClass(
             ElasticsearchEntitlementChecker.class,
@@ -167,17 +169,20 @@ static ElasticsearchEntitlementChecker initChecker(Instrumentation inst, PolicyM
             throw new AssertionError(e);
         }
 
+        return checker;
+    }
+
+    static void initInstrumentation(Instrumentation instrumentation) throws Exception {
         var verifyBytecode = Booleans.parseBoolean(System.getProperty("es.entitlements.verify_bytecode", "false"));
         if (verifyBytecode) {
             ensureClassesSensitiveToVerificationAreInitialized();
         }
 
         DynamicInstrumentation.initialize(
-            inst,
+            instrumentation,
             EntitlementCheckerUtils.getVersionSpecificCheckerClass(EntitlementChecker.class, Runtime.version().feature()),
             verifyBytecode
         );
 
-        return checker;
     }
 }

muted-tests.yml

@@ -492,18 +492,12 @@ tests:
 - class: org.elasticsearch.packaging.test.DockerTests
   method: test085EnvironmentVariablesAreRespectedUnderDockerExec
   issue: https://github.com/elastic/elasticsearch/issues/128115
-- class: org.elasticsearch.test.apmintegration.TracesApmIT
-  method: testApmIntegration
-  issue: https://github.com/elastic/elasticsearch/issues/128651
 - class: org.elasticsearch.xpack.esql.expression.function.scalar.string.WildcardLikeTests
   method: testEvaluateInManyThreads {TestCase=100 random code points matches self case insensitive with keyword}
   issue: https://github.com/elastic/elasticsearch/issues/128676
 - class: org.elasticsearch.xpack.esql.expression.function.scalar.string.WildcardLikeTests
   method: testEvaluateInManyThreads {TestCase=100 random code points matches self case insensitive with text}
   issue: https://github.com/elastic/elasticsearch/issues/128677
-- class: org.elasticsearch.test.apmintegration.MetricsApmIT
-  method: testApmIntegration
-  issue: https://github.com/elastic/elasticsearch/issues/128678
 
 # Examples:
 #

test/framework/src/main/java/org/elasticsearch/entitlement/initialization/TestEntitlementInitialization.java

@@ -31,6 +31,8 @@
 import java.util.List;
 import java.util.Map;
 
+import static org.elasticsearch.entitlement.initialization.EntitlementInitialization.initInstrumentation;
+
 /**
  * Test-specific version of {@code EntitlementInitialization}
  */
@@ -45,7 +47,8 @@ public static EntitlementChecker checker() {
     }
 
     public static void initialize(Instrumentation inst) throws Exception {
-        checker = EntitlementInitialization.initChecker(inst, createPolicyManager(initializeArgs.pathLookup()));
+        checker = EntitlementInitialization.initChecker(createPolicyManager(initializeArgs.pathLookup()));
+        initInstrumentation(inst);
     }
 
     public record InitializeArgs(PathLookup pathLookup) {}