Create equivalents of JSM's AccessController in the java agent

CHANGELOG.md

@@ -28,6 +28,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Use QueryCoordinatorContext for the rewrite in validate API. ([#18272](https://github.com/opensearch-project/OpenSearch/pull/18272))
 - Upgrade crypto kms plugin dependencies for AWS SDK v2.x. ([#18268](https://github.com/opensearch-project/OpenSearch/pull/18268))
 - Add support for `matched_fields` with the unified highlighter ([#18164](https://github.com/opensearch-project/OpenSearch/issues/18164))
+- Create equivalents of JSM's AccessController in the java agent ([#18346](https://github.com/opensearch-project/OpenSearch/issues/18346))
 
 ### Changed
 - Create generic DocRequest to better categorize ActionRequests ([#18269](https://github.com/opensearch-project/OpenSearch/pull/18269)))

libs/agent-sm/agent/src/main/java/org/opensearch/javaagent/StackCallerProtectionDomainChainExtractor.java

@@ -36,7 +36,9 @@ private StackCallerProtectionDomainChainExtractor() {}
     @Override
     public Collection<ProtectionDomain> apply(Stream<StackFrame> frames) {
         return frames.takeWhile(
-            frame -> !(frame.getClassName().equals("java.security.AccessController") && frame.getMethodName().equals("doPrivileged"))
+            frame -> !((frame.getClassName().equals("java.security.AccessController")
+                || frame.getClassName().equals("org.opensearch.javaagent.bootstrap.AccessController"))
+                && frame.getMethodName().equals("doPrivileged"))
         )
             .map(StackFrame::getDeclaringClass)
             .map(Class::getProtectionDomain)

libs/agent-sm/agent/src/test/java/org/opensearch/javaagent/StackCallerProtectionDomainExtractorTests.java

@@ -115,4 +115,76 @@ public Void run() {
             }
         });
     }
+
+    @Test
+    public void testStackTruncationWithOpenSearchAccessController() throws Exception {
+        org.opensearch.javaagent.bootstrap.AccessController.doPrivileged(new org.opensearch.javaagent.bootstrap.PrivilegedAction<Void>() {
+            @Override
+            public Void run() {
+                StackCallerProtectionDomainChainExtractor extractor = StackCallerProtectionDomainChainExtractor.INSTANCE;
+                Set<ProtectionDomain> protectionDomains = (Set<ProtectionDomain>) extractor.apply(captureStackFrames().stream());
+                assertEquals(1, protectionDomains.size());
+                List<String> simpleNames = protectionDomains.stream().map(pd -> {
+                    try {
+                        return pd.getCodeSource().getLocation().toURI();
+                    } catch (URISyntaxException e) {
+                        throw new RuntimeException(e);
+                    }
+                })
+                    .map(URI::getPath)
+                    .map(Paths::get)
+                    .map(Path::getFileName)
+                    .map(Path::toString)
+                    // strip trailing “-VERSION.jar” if present
+                    .map(name -> name.replaceFirst("-\\d[\\d\\.]*\\.jar$", ""))
+                    // otherwise strip “.jar”
+                    .map(name -> name.replaceFirst("\\.jar$", ""))
+                    .toList();
+                assertThat(
+                    simpleNames,
+                    containsInAnyOrder(
+                        "test"    // from the build/classes/java/test directory
+                    )
+                );
+                return null;
+            }
+        });
+    }
+
+    @Test
+    public void testStackTruncationWithOpenSearchAccessControllerWithCheckedException() throws Exception {
+        org.opensearch.javaagent.bootstrap.AccessController.doPrivileged(
+            new org.opensearch.javaagent.bootstrap.PrivilegedExceptionAction<Void>() {
+                @Override
+                public Void run() {
+                    StackCallerProtectionDomainChainExtractor extractor = StackCallerProtectionDomainChainExtractor.INSTANCE;
+                    Set<ProtectionDomain> protectionDomains = (Set<ProtectionDomain>) extractor.apply(captureStackFrames().stream());
+                    assertEquals(1, protectionDomains.size());
+                    List<String> simpleNames = protectionDomains.stream().map(pd -> {
+                        try {
+                            return pd.getCodeSource().getLocation().toURI();
+                        } catch (URISyntaxException e) {
+                            throw new RuntimeException(e);
+                        }
+                    })
+                        .map(URI::getPath)
+                        .map(Paths::get)
+                        .map(Path::getFileName)
+                        .map(Path::toString)
+                        // strip trailing “-VERSION.jar” if present
+                        .map(name -> name.replaceFirst("-\\d[\\d\\.]*\\.jar$", ""))
+                        // otherwise strip “.jar”
+                        .map(name -> name.replaceFirst("\\.jar$", ""))
+                        .toList();
+                    assertThat(
+                        simpleNames,
+                        containsInAnyOrder(
+                            "test"    // from the build/classes/java/test directory
+                        )
+                    );
+                    return null;
+                }
+            }
+        );
+    }
 }

libs/agent-sm/bootstrap/src/main/java/org/opensearch/javaagent/bootstrap/AccessController.java

@@ -0,0 +1,68 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.javaagent.bootstrap;
+
+/**
+ * Utility class to run code in a privileged block.
+ */
+public final class AccessController {
+    /**
+     * Don't allow instantiation an {@code AccessController}
+     */
+    private AccessController() {}
+
+    /**
+     * Performs the specified action in a privileged block.
+     *
+     * <p> If the action's {@code run} method throws an (unchecked)
+     * exception, it will propagate through this method.
+     *
+     * @param <T> the type of the value returned by the PrivilegedAction's
+     *                  {@code run} method
+     *
+     * @param action the action to be performed
+     *
+     * @return the value returned by the action's {@code run} method
+     */
+    public static <T> T doPrivileged(PrivilegedAction<T> action) {
+        T result = action.run();
+        return result;
+    }
+
+    private static PrivilegedActionException wrapException(Exception e) {
+        return new PrivilegedActionException(e);
+    }
+
+    /**
+     * Performs the specified action.
+     *
+     * <p> If the action's {@code run} method throws an <i>unchecked</i>
+     * exception, it will propagate through this method.
+     *
+     * @param <T> the type of the value returned by the
+     *                  PrivilegedExceptionAction's {@code run} method
+     *
+     * @param action the action to be performed
+     *
+     * @return the value returned by the action's {@code run} method
+     *
+     * @throws    PrivilegedActionException if the specified action's
+     *         {@code run} method threw a <i>checked</i> exception
+     */
+    public static <T> T doPrivileged(PrivilegedExceptionAction<T> action) throws PrivilegedActionException {
+        try {
+            T result = action.run();
+            return result;
+        } catch (RuntimeException e) {
+            throw e;
+        } catch (Exception e) {
+            throw wrapException(e);
+        }
+    }
+}

libs/agent-sm/bootstrap/src/main/java/org/opensearch/javaagent/bootstrap/PrivilegedAction.java

@@ -0,0 +1,35 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.javaagent.bootstrap;
+
+/**
+ * A computation to be performed by invoking
+ * {@code AccessController.doPrivileged} on the
+ * {@code PrivilegedAction} object.  This interface is used only for
+ * computations that do not throw checked exceptions; computations that
+ * throw checked exceptions must use {@code PrivilegedExceptionAction}
+ * instead.
+ * @param <T> the type of the result of running the computation
+ *
+ * @see AccessController
+ * @see AccessController#doPrivileged(PrivilegedAction)
+ * @see PrivilegedExceptionAction
+ */
+@FunctionalInterface
+public interface PrivilegedAction<T> {
+    /**
+     * Performs the computation.  This method will be called by
+     * {@code AccessController.doPrivileged}.
+     *
+     * @return a class-dependent value that may represent the results of the
+     *         computation.
+     * @see AccessController#doPrivileged(PrivilegedAction)
+     */
+    T run();
+}

libs/agent-sm/bootstrap/src/main/java/org/opensearch/javaagent/bootstrap/PrivilegedActionException.java

@@ -0,0 +1,45 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.javaagent.bootstrap;
+
+/**
+ * This exception is thrown by
+ * {@code doPrivileged(PrivilegedExceptionAction)} to indicate
+ * that the action being performed threw a checked exception.  The exception
+ * thrown by the action can be obtained by calling the
+ * {@code getException} method.  In effect, an
+ * {@code PrivilegedActionException} is a "wrapper"
+ * for an exception thrown by a privileged action.
+ *
+ * @see PrivilegedExceptionAction
+ * @see AccessController#doPrivileged(PrivilegedExceptionAction)
+ */
+public class PrivilegedActionException extends Exception {
+    /**
+     * Constructs a new {@code PrivilegedActionException} "wrapping"
+     * the specific Exception.
+     *
+     * @param exception The exception thrown
+     */
+    public PrivilegedActionException(Exception exception) {
+        super(null, exception);  // Disallow initCause
+    }
+
+    /**
+     * Returns a string representation of this exception.
+     *
+     * @return a string representation of this exception.
+     */
+    @Override
+    public String toString() {
+        String s = getClass().getName();
+        Throwable cause = super.getCause();
+        return (cause != null) ? (s + ": " + cause.toString()) : s;
+    }
+}

libs/agent-sm/bootstrap/src/main/java/org/opensearch/javaagent/bootstrap/PrivilegedExceptionAction.java

@@ -0,0 +1,39 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.javaagent.bootstrap;
+
+import java.security.PrivilegedAction;
+
+/**
+ * A computation to be performed with privileges, that throws one or
+ * more checked exceptions.  The computation is performed by invoking
+ * {@code AccessController.doPrivileged} on the
+ * {@code PrivilegedExceptionAction} object.  This interface is
+ * used only for computations that throw checked exceptions;
+ * computations that do not throw
+ * checked exceptions should use {@code PrivilegedAction} instead.
+ * @param <T> the type of the result of running the computation
+ *
+ * @see PrivilegedAction
+ */
+@FunctionalInterface
+public interface PrivilegedExceptionAction<T> {
+    /**
+     * Performs the computation.  This method will be called by
+     * {@code AccessController.doPrivileged}.
+     *
+     * @return a class-dependent value that may represent the results of the
+     *         computation.
+     * @throws Exception an exceptional condition has occurred.  Each class
+     *         that implements {@code PrivilegedExceptionAction} should
+     *         document the exceptions that its run method can throw.
+     */
+
+    T run() throws Exception;
+}