Custom getCallerClass in entitlement bridge (elastic#125139)

Author: prdoyle

libs/entitlement/asm-provider/src/main/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumenterImpl.java

@@ -37,10 +37,8 @@
 import static org.objectweb.asm.ClassWriter.COMPUTE_FRAMES;
 import static org.objectweb.asm.ClassWriter.COMPUTE_MAXS;
 import static org.objectweb.asm.Opcodes.ACC_STATIC;
-import static org.objectweb.asm.Opcodes.GETSTATIC;
 import static org.objectweb.asm.Opcodes.INVOKEINTERFACE;
 import static org.objectweb.asm.Opcodes.INVOKESTATIC;
-import static org.objectweb.asm.Opcodes.INVOKEVIRTUAL;
 
 public class InstrumenterImpl implements Instrumenter {
     private static final Logger logger = LogManager.getLogger(InstrumenterImpl.class);
@@ -286,22 +284,9 @@ private void pushCallerClass() {
                     false
                 );
             } else {
-                mv.visitFieldInsn(
-                    GETSTATIC,
-                    Type.getInternalName(StackWalker.Option.class),
-                    "RETAIN_CLASS_REFERENCE",
-                    Type.getDescriptor(StackWalker.Option.class)
-                );
                 mv.visitMethodInsn(
                     INVOKESTATIC,
-                    Type.getInternalName(StackWalker.class),
-                    "getInstance",
-                    Type.getMethodDescriptor(Type.getType(StackWalker.class), Type.getType(StackWalker.Option.class)),
-                    false
-                );
-                mv.visitMethodInsn(
-                    INVOKEVIRTUAL,
-                    Type.getInternalName(StackWalker.class),
+                    "org/elasticsearch/entitlement/bridge/Util",
                     "getCallerClass",
                     Type.getMethodDescriptor(Type.getType(Class.class)),
                     false

libs/entitlement/bridge/src/main/java/org/elasticsearch/entitlement/bridge/EntitlementCheckerHandle.java

@@ -18,6 +18,7 @@ public class EntitlementCheckerHandle {
      * This is how the bytecodes injected by our instrumentation access the {@link EntitlementChecker}
      * so they can call the appropriate check method.
      */
+    @SuppressWarnings("unused")
     public static EntitlementChecker instance() {
         return Holder.instance;
     }

libs/entitlement/bridge/src/main/java/org/elasticsearch/entitlement/bridge/Util.java

@@ -0,0 +1,44 @@
+/*
+ * 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.bridge;
+
+import java.util.Optional;
+
+import static java.lang.StackWalker.Option.RETAIN_CLASS_REFERENCE;
+
+public class Util {
+    /**
+     * A special value representing the case where a method <em>has no caller</em>.
+     * This can occur if it's called directly from the JVM.
+     *
+     * @see StackWalker#getCallerClass()
+     */
+    public static final Class<?> NO_CLASS = new Object() {
+    }.getClass();
+
+    /**
+     * Why would we write this instead of using {@link StackWalker#getCallerClass()}?
+     * Because that method throws {@link IllegalCallerException} if called from the "outermost frame",
+     * which includes at least some cases of a method called from a native frame.
+     *
+     * @return the class that called the method which called this; or {@link #NO_CLASS} from the outermost frame.
+     */
+    @SuppressWarnings("unused") // Called reflectively from InstrumenterImpl
+    public static Class<?> getCallerClass() {
+        Optional<Class<?>> callerClassIfAny = StackWalker.getInstance(RETAIN_CLASS_REFERENCE)
+            .walk(
+                frames -> frames.skip(2) // Skip this method and its caller
+                    .findFirst()
+                    .map(StackWalker.StackFrame::getDeclaringClass)
+            );
+        return callerClassIfAny.orElse(NO_CLASS);
+    }
+
+}

libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java

@@ -57,6 +57,7 @@
 import static java.util.stream.Collectors.toUnmodifiableMap;
 import static java.util.zip.ZipFile.OPEN_DELETE;
 import static java.util.zip.ZipFile.OPEN_READ;
+import static org.elasticsearch.entitlement.bridge.Util.NO_CLASS;
 
 public class PolicyManager {
     /**
@@ -712,8 +713,6 @@ Class<?> requestingClass(Class<?> callerClass) {
 
     /**
      * Given a stream of {@link StackFrame}s, identify the one whose entitlements should be checked.
-     *
-     * @throws NullPointerException if the requesting module is {@code null}
      */
     Optional<StackFrame> findRequestingFrame(Stream<StackFrame> frames) {
         return frames.filter(f -> f.getDeclaringClass().getModule() != entitlementsModule) // ignore entitlements library
@@ -732,6 +731,10 @@ private static boolean isTriviallyAllowed(Class<?> requestingClass) {
             generalLogger.debug("Entitlement trivially allowed: no caller frames outside the entitlement library");
             return true;
         }
+        if (requestingClass == NO_CLASS) {
+            generalLogger.debug("Entitlement trivially allowed from outermost frame");
+            return true;
+        }
         if (SYSTEM_LAYER_MODULES.contains(requestingClass.getModule())) {
             generalLogger.debug("Entitlement trivially allowed from system module [{}]", requestingClass.getModule().getName());
             return true;

libs/entitlement/src/test/java/org/elasticsearch/entitlement/bridge/UtilTests.java

@@ -0,0 +1,33 @@
+/*
+ * 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.bridge;
+
+import org.elasticsearch.test.ESTestCase;
+
+import static org.elasticsearch.entitlement.bridge.UtilTests.MockSensitiveClass.mockSensitiveMethod;
+
+@ESTestCase.WithoutSecurityManager
+public class UtilTests extends ESTestCase {
+
+    public void testCallerClass() {
+        assertEquals(UtilTests.class, mockSensitiveMethod());
+    }
+
+    /**
+     * A separate class so the stack walk can discern the sensitive method's own class
+     * from that of its caller.
+     */
+    static class MockSensitiveClass {
+        public static Class<?> mockSensitiveMethod() {
+            return Util.getCallerClass();
+        }
+    }
+
+}