Extend public-callers-finder to check if methods are instrumented with entitlements. (elastic#136022)

If `--check-instrumentation` is given, the public-callers-finder will load the dump of instrumented methods and check that against the public callers the tool found. An additional column is added to the output, that is either COVERED if the method is already instrumented or MISSING if not.

Relates to ES-11757

Author: mosche

libs/entitlement/build.gradle

@@ -13,6 +13,7 @@ apply plugin: 'elasticsearch.build'
 apply plugin: 'elasticsearch.publish'
 apply plugin: 'elasticsearch.embedded-providers'
 apply plugin: 'elasticsearch.mrjar'
+apply plugin: 'elasticsearch.internal-test-artifact'
 
 embeddedProviders {
   impl 'entitlement', project(':libs:entitlement:asm-provider')
@@ -27,6 +28,11 @@ dependencies {
     exclude group: 'org.elasticsearch', module: 'entitlement'
   }
 
+  testImplementation 'org.ow2.asm:asm:9.8'
+  testImplementation 'org.ow2.asm:asm-util:9.8'
+  testImplementation 'org.ow2.asm:asm-tree:9.8'
+  testImplementation 'org.ow2.asm:asm-analysis:9.8'
+
   // guarding for intellij
   if (sourceSets.findByName("main23")) {
     main23CompileOnly project(path: ':libs:entitlement:bridge', configuration: 'java23')
@@ -36,3 +42,9 @@ dependencies {
 tasks.withType(CheckForbiddenApisTask).configureEach {
   replaceSignatureFiles 'jdk-signatures'
 }
+
+// invoke using ./gradlew :libs:entitlement:dumpInstrumentedMethods --args="<output file>"
+tasks.register('dumpInstrumentedMethods', JavaExec) {
+  classpath = sourceSets.test.runtimeClasspath
+  mainClass = 'org.elasticsearch.entitlement.initialization.DynamicInstrumentationUtils'
+}

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

@@ -123,7 +123,7 @@ static void initialize(Instrumentation inst, Class<?> checkerInterface, boolean
         }
     }
 
-    private static Map<MethodKey, CheckMethod> getMethodsToInstrument(Class<?> checkerInterface) throws ClassNotFoundException,
+    static Map<MethodKey, CheckMethod> getMethodsToInstrument(Class<?> checkerInterface) throws ClassNotFoundException,
         NoSuchMethodException {
         Map<MethodKey, CheckMethod> checkMethods = new HashMap<>(INSTRUMENTATION_SERVICE.lookupMethods(checkerInterface));
         Stream.of(fileSystemProviderChecks(), fileStoreChecks(), pathChecks(), selectorProviderChecks())

libs/entitlement/src/test/java/org/elasticsearch/entitlement/initialization/DynamicInstrumentationTests.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.entitlement.initialization;
+
+import org.elasticsearch.entitlement.initialization.DynamicInstrumentationUtils.Descriptor;
+import org.elasticsearch.test.ESTestCase;
+import org.hamcrest.Description;
+import org.hamcrest.Matcher;
+import org.hamcrest.TypeSafeMatcher;
+
+import java.util.List;
+
+import static org.elasticsearch.test.LambdaMatchers.transformedMatch;
+import static org.hamcrest.Matchers.anyOf;
+import static org.hamcrest.Matchers.everyItem;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.startsWith;
+
+public class DynamicInstrumentationTests extends ESTestCase {
+
+    public void testInstrumentedMethodsExist() throws Exception {
+        List<Descriptor> descriptors = DynamicInstrumentationUtils.loadInstrumentedMethodDescriptors();
+        assertThat(
+            descriptors,
+            everyItem(
+                anyOf(
+                    isKnownDescriptor(),
+                    // not visible
+                    transformedMatch(Descriptor::className, startsWith("jdk/vm/ci/services/")),
+                    // removed in JDK 24
+                    is(
+                        new Descriptor(
+                            "sun/net/www/protocol/http/HttpURLConnection",
+                            "openConnectionCheckRedirects",
+                            List.of("java/net/URLConnection"),
+                            null
+                        )
+                    )
+                )
+            )
+        );
+    }
+
+    static Matcher<Descriptor> isKnownDescriptor() {
+        return new TypeSafeMatcher<>() {
+            @Override
+            protected boolean matchesSafely(Descriptor item) {
+                return item.methodDescriptor() != null;
+            }
+
+            @Override
+            public void describeTo(Description description) {
+                description.appendText("valid method descriptor");
+            }
+
+            @Override
+            protected void describeMismatchSafely(Descriptor item, Description mismatchDescription) {
+                mismatchDescription.appendText("was ").appendValue(item.toString());
+            }
+        };
+    }
+}

libs/entitlement/src/test/java/org/elasticsearch/entitlement/initialization/DynamicInstrumentationUtils.java

@@ -0,0 +1,91 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.entitlement.initialization;
+
+import org.elasticsearch.common.logging.LogConfigurator;
+import org.elasticsearch.entitlement.bridge.EntitlementChecker;
+import org.elasticsearch.entitlement.instrumentation.CheckMethod;
+import org.elasticsearch.entitlement.instrumentation.MethodKey;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Stream;
+
+import static java.util.Objects.requireNonNull;
+
+public class DynamicInstrumentationUtils {
+
+    /**
+     * This dumps the list of instrumented methods to a file specified by the first argument
+     * or alternatively the system property `es.entitlements.dump`.
+     */
+    public static void main(String[] args) throws Exception {
+        LogConfigurator.loadLog4jPlugins();
+        LogConfigurator.configureESLogging();
+
+        var path = requireNonNull(args.length > 0 ? args[0] : System.getProperty("es.entitlements.dump"), "destination for dump required");
+        var descriptors = loadInstrumentedMethodDescriptors();
+        Files.write(
+            Path.of(path),
+            () -> descriptors.stream().filter(d -> d.methodDescriptor != null).map(Descriptor::toLine).iterator(),
+            StandardCharsets.UTF_8
+        );
+    }
+
+    static List<Descriptor> loadInstrumentedMethodDescriptors() throws Exception {
+        Map<MethodKey, CheckMethod> methodsToInstrument = DynamicInstrumentation.getMethodsToInstrument(
+            EntitlementCheckerUtils.getVersionSpecificCheckerClass(EntitlementChecker.class, Runtime.version().feature())
+        );
+        return methodsToInstrument.keySet().stream().map(DynamicInstrumentationUtils::lookupDescriptor).toList();
+    }
+
+    private static Descriptor lookupDescriptor(MethodKey key) {
+        final String[] foundDescriptor = { null };
+        try {
+            ClassReader reader = new ClassReader(key.className());
+            reader.accept(new ClassVisitor(Opcodes.ASM9) {
+                @Override
+                public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+                    if (name.equals(key.methodName()) == false) {
+                        return null;
+                    }
+                    List<String> argTypes = Stream.of(Type.getArgumentTypes(descriptor)).map(Type::getInternalName).toList();
+                    if (argTypes.equals(key.parameterTypes())) {
+                        foundDescriptor[0] = descriptor;
+                    }
+                    return null;
+                }
+            }, 0);
+        } catch (IOException e) {
+            // nothing to do
+        }
+        return new Descriptor(key.className(), key.methodName(), key.parameterTypes(), foundDescriptor[0]);
+
+    }
+
+    record Descriptor(String className, String methodName, List<String> parameterTypes, String methodDescriptor) {
+
+        private static final String SEPARATOR = "\t";
+
+        CharSequence toLine() {
+            return String.join(SEPARATOR, className, methodName, methodDescriptor);
+        }
+
+    }
+}

libs/entitlement/tools/jdk-api-extractor/README.md

@@ -32,15 +32,17 @@ cat libs/entitlement/tools/jdk-api-extractor/api.diff | grep '^+[^+]' | sed 's/^
 
 # review new additions next for critical ones that should require entitlements
 # once done, remove all lines that are not considered critical and run the public-callers-finder to report
-# the transitive public surface for these additions
-./gradlew :libs:entitlement:tools:public-callers-finder:run -Druntime.java=25 --args="api-jdk25-additions.tsv true"
+# the transitive public surface for these additions to identify methods that are node already covered by entitlements
+./gradlew :libs:entitlement:tools:public-callers-finder:run -Druntime.java=25 --args="api-jdk25-additions.tsv --transitive --check-instrumentation"
 ```
 
 ### Optional arguments:
 
 - `--deprecations-only`: reports public deprecations (by means of `@Deprecated`)
 - `--include-incubator`: include incubator modules (e.g. `jdk.incubator.vector`)
 
+If `-Druntime.java` is not provided, the bundled JDK is used.
+
 ```bash
 ./gradlew :libs:entitlement:tools:jdk-api-extractor:run -Druntime.java=24 --args="deprecations-jdk24.tsv --deprecations-only"
 ```

libs/entitlement/tools/public-callers-finder/README.md

@@ -15,20 +15,30 @@ it treats calls to `super` in `S.m` as regular calls (e.g. `example() -> S.m() -
 
 In order to run the tool, use:
 ```shell
-./gradlew :libs:entitlement:tools:public-callers-finder:run -Druntime.java=25 --args="<input-file> [<bubble-up-from-public>]"
+./gradlew :libs:entitlement:tools:public-callers-finder:run [-Druntime.java=25] --args="<input-file> [--transitive] [--check-instrumentation]"
 ```
-Where `input-file` is a CSV file (columns separated by `TAB`) that contains the following columns:
-1. Module name
-2. unused
-3. unused
-4. Fully qualified class name (ASM style, with `/` separators)
-5. Method name
-6. Method descriptor (ASM signature)
-7. Visibility (PUBLIC/PUBLIC-METHOD/PRIVATE)
 
-And `bubble-up-from-public` is a boolean (`true|false`) indicating if the code should stop at the first public method (`false`: default) or continue to find usages recursively even after reaching the "public surface".
+- `input-file` is a `TAB`-separated TSV file containing the following columns:
+  1. Module name
+  2. unused
+  3. unused
+  4. Fully qualified class name (ASM style, with `/` separators)
+  5. Method name
+  6. Method descriptor (ASM signature)
+  7. Visibility (PUBLIC/PUBLIC-METHOD/PRIVATE)
+
+- optional: `--transitive` to not stop at the first public method, but continue to find the transitive public surface.
 
-The output of the tool is another CSV file, with one line for each entry-point, columns separated by `TAB`
+- optional: `--check-instrumentation` to check if methods are instrumented for entitlements.
+
+If `-Druntime.java` is not provided, the bundled JDK is used.
+
+Examples:
+```bash
+./gradlew :libs:entitlement:tools:public-callers-finder:run --args="sensitive-methods.tsv true" --transitive --check-instrumentation"
+```
+
+The tool writes the following `TAB`-separated columns to standard out:
 
 1. Module name
 2. File name (from source root)
@@ -37,12 +47,13 @@ The output of the tool is another CSV file, with one line for each entry-point,
 5. Method name
 6. Method descriptor (ASM signature)
 7. Visibility (PUBLIC/PUBLIC-METHOD/PRIVATE)
-8. Original caller Module name
-9. Original caller Class name (ASM style, with `/` separators)
-10. Original caller Method name
-11. Original caller Visibility
+8. If using `--check-instrumentation`: `COVERED` if method is instrumented with entitlement checks, `MISSING` otherwise
+9. Original caller Module name
+10. Original caller Class name (ASM style, with `/` separators)
+11. Original caller Method name
+12. Original caller Visibility
 
-Examples:
+Example output:
 ```
 java.base	DeleteOnExitHook.java	50	java/io/DeleteOnExitHook$1	run	()V	PUBLIC	java.base	java/io/File	delete	PUBLIC
 java.base	ZipFile.java	254	java/util/zip/ZipFile	<init>	(Ljava/io/File;ILjava/nio/charset/Charset;)V	PUBLIC	java.base	java/io/File	delete	PUBLIC

libs/entitlement/tools/public-callers-finder/build.gradle

@@ -42,10 +42,24 @@ application {
   ]
 }
 
+configurations {
+  dumpMethods
+}
+
+def instrumentedMethodsDump = project.layout.buildDirectory.file('instrumented-methods.dump').get().asFile
+
 tasks.named("run").configure {
+  dependsOn 'dumpInstrumentedMethods'
+  systemProperty 'es.entitlements.dump', instrumentedMethodsDump
   executable = "${buildParams.runtimeJavaHome.get()}/bin/java" + (OS.current() == OS.WINDOWS ? '.exe' : '')
 }
 
+tasks.register('dumpInstrumentedMethods', JavaExec) {
+  classpath = configurations.dumpMethods
+  mainClass = 'org.elasticsearch.entitlement.initialization.DynamicInstrumentationUtils'
+  systemProperty 'es.entitlements.dump', instrumentedMethodsDump
+}
+
 repositories {
   mavenCentral()
 }
@@ -55,6 +69,8 @@ dependencies {
   implementation 'org.ow2.asm:asm:9.8'
   implementation 'org.ow2.asm:asm-util:9.8'
   implementation(project(':libs:entitlement:tools:common'))
+
+  dumpMethods testArtifact(project(':libs:entitlement'))
 }
 
 tasks.named('forbiddenApisMain').configure {

libs/entitlement/tools/public-callers-finder/src/main/java/org/elasticsearch/entitlement/tools/publiccallersfinder/FindUsagesClassVisitor.java

@@ -31,18 +31,10 @@ class FindUsagesClassVisitor extends ClassVisitor {
 
     record MethodDescriptor(String className, String methodName, String methodDescriptor) {}
 
-    record EntryPoint(
-        String moduleName,
-        String source,
-        int line,
-        String className,
-        String methodName,
-        String methodDescriptor,
-        EnumSet<ExternalAccess> access
-    ) {}
+    record EntryPoint(String moduleName, String source, int line, MethodDescriptor method, EnumSet<ExternalAccess> access) {}
 
     interface CallerConsumer {
-        void accept(String source, int line, String className, String methodName, String methodDescriptor, EnumSet<ExternalAccess> access);
+        void accept(String source, int line, MethodDescriptor method, EnumSet<ExternalAccess> access);
     }
 
     private final Set<String> moduleExports;
@@ -126,7 +118,7 @@ public void visitMethodInsn(int opcode, String owner, String name, String descri
                             (methodAccess & ACC_PUBLIC) != 0,
                             (methodAccess & ACC_PROTECTED) != 0
                         );
-                        callers.accept(source, line, className, methodName, methodDescriptor, externalAccess);
+                        callers.accept(source, line, new MethodDescriptor(className, methodName, methodDescriptor), externalAccess);
                     }
                 }
             }