Enhance Java Agent to intercept Runtime::halt

CHANGELOG.md

@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Security Manager Replacement] Create initial Java Agent to intercept Socket::connect calls ([#17724](https://github.com/opensearch-project/OpenSearch/pull/17724))
 - Add ingestion management APIs for pause, resume and get ingestion state ([#17631](https://github.com/opensearch-project/OpenSearch/pull/17631))
 - [Security Manager Replacement] Enhance Java Agent to intercept System::exit ([#17746](https://github.com/opensearch-project/OpenSearch/pull/17746))
+- [Security Manager Replacement] Enhance Java Agent to intercept Runtime::halt ([#17757](https://github.com/opensearch-project/OpenSearch/pull/17757))
 - Support AutoExpand for SearchReplica ([#17741](https://github.com/opensearch-project/OpenSearch/pull/17741))
 
 ### Changed

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

@@ -64,8 +64,10 @@
         ClassInjector.UsingUnsafe.ofBootLoader()
             .inject(
                 Map.of(
-                    new TypeDescription.ForLoadedType(StackCallerChainExtractor.class),
-                    ClassFileLocator.ForClassLoader.read(StackCallerChainExtractor.class),
+                    new TypeDescription.ForLoadedType(StackCallerProtectionDomainChainExtractor.class),
+                    ClassFileLocator.ForClassLoader.read(StackCallerProtectionDomainChainExtractor.class),
+                    new TypeDescription.ForLoadedType(StackCallerClassChainExtractor.class),
+                    ClassFileLocator.ForClassLoader.read(StackCallerClassChainExtractor.class),
                     new TypeDescription.ForLoadedType(AgentPolicy.class),
                     ClassFileLocator.ForClassLoader.read(AgentPolicy.class)
                 )
@@ -83,6 +85,12 @@
                 (b, typeDescription, classLoader, module, pd) -> b.visit(
                     Advice.to(SystemExitInterceptor.class).on(ElementMatchers.named("exit"))
                 )
+            )
+            .type(ElementMatchers.is(java.lang.Runtime.class))
+            .transform(
+                (b, typeDescription, classLoader, module, pd) -> b.visit(
+                    Advice.to(RuntimeHaltInterceptor.class).on(ElementMatchers.named("halt"))
+                )
             );
 
         return agentBuilder;

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

@@ -0,0 +1,49 @@
+/*
+ * 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;
+
+import org.opensearch.javaagent.bootstrap.AgentPolicy;
+
+import java.lang.StackWalker.Option;
+import java.security.Policy;
+import java.util.stream.Stream;
+
+import net.bytebuddy.asm.Advice;
+
+/**
+ * {@link Runtime#halt} interceptor
+ */
+public class RuntimeHaltInterceptor {
+    /**
+     * RuntimeHaltInterceptor
+     */
+    public RuntimeHaltInterceptor() {}
+
+    /**
+     * Interceptor
+     * @param code exit code
+     * @throws Exception exceptions
+     */
+    @Advice.OnMethodEnter
+    @SuppressWarnings("removal")
+    public static void intercept(int code) throws Exception {
+        final Policy policy = AgentPolicy.getPolicy();
+        if (policy == null) {
+            return; /* noop */
+        }
+
+        final StackWalker walker = StackWalker.getInstance(Option.RETAIN_CLASS_REFERENCE);
+        final Class<?> caller = walker.getCallerClass();
+        final Stream<Class<?>> chain = walker.walk(StackCallerClassChainExtractor.INSTANCE);
+
+        if (AgentPolicy.isChainThatCanExit(caller, chain) == false) {
+            throw new SecurityException("The class " + caller + " is not allowed to call Runtime::halt(" + code + ")");
+        }
+    }
+}

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

@@ -47,7 +47,7 @@
         }
 
         final StackWalker walker = StackWalker.getInstance(Option.RETAIN_CLASS_REFERENCE);
-        final Stream<ProtectionDomain> callers = walker.walk(StackCallerChainExtractor.INSTANCE);
+        final Stream<ProtectionDomain> callers = walker.walk(StackCallerProtectionDomainChainExtractor.INSTANCE);
 
         if (args[0] instanceof InetSocketAddress address) {
             if (!AgentPolicy.isTrustedHost(address.getHostString())) {

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

@@ -0,0 +1,42 @@
+/*
+ * 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;
+
+import java.lang.StackWalker.StackFrame;
+import java.util.function.Function;
+import java.util.stream.Stream;
+
+/**
+ * Stack Caller Class Chain Extractor
+ */
+public final class StackCallerClassChainExtractor implements Function<Stream<StackFrame>, Stream<Class<?>>> {
+    /**
+     * Single instance of stateless class.
+     */
+    public static final StackCallerClassChainExtractor INSTANCE = new StackCallerClassChainExtractor();
+
+    /**
+     * Constructor
+     */
+    private StackCallerClassChainExtractor() {}
+
+    /**
+     * Folds the stack
+     * @param frames stack frames
+     */
+    @Override
+    public Stream<Class<?>> apply(Stream<StackFrame> frames) {
+        return cast(frames);
+    }
+
+    @SuppressWarnings("unchecked")
+    private static <A> Stream<A> cast(Stream<StackFrame> frames) {
+        return (Stream<A>) frames.map(StackFrame::getDeclaringClass).filter(c -> !c.isHidden()).distinct();
+    }
+}

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

@@ -16,16 +16,16 @@
 /**
  * Stack Caller Chain Extractor
  */
-public final class StackCallerChainExtractor implements Function<Stream<StackFrame>, Stream<ProtectionDomain>> {
+public final class StackCallerProtectionDomainChainExtractor implements Function<Stream<StackFrame>, Stream<ProtectionDomain>> {
     /**
      * Single instance of stateless class.
      */
-    public static final StackCallerChainExtractor INSTANCE = new StackCallerChainExtractor();
+    public static final StackCallerProtectionDomainChainExtractor INSTANCE = new StackCallerProtectionDomainChainExtractor();
 
     /**
      * Constructor
      */
-    private StackCallerChainExtractor() {}
+    private StackCallerProtectionDomainChainExtractor() {}
 
     /**
      * Folds the stack

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

@@ -11,6 +11,8 @@
 import org.opensearch.javaagent.bootstrap.AgentPolicy;
 
 import java.lang.StackWalker.Option;
+import java.security.Policy;
+import java.util.stream.Stream;
 
 import net.bytebuddy.asm.Advice;
 
@@ -29,11 +31,18 @@
      * @throws Exception exceptions
      */
     @Advice.OnMethodEnter()
+    @SuppressWarnings("removal")
     public static void intercept(int code) throws Exception {
+        final Policy policy = AgentPolicy.getPolicy();
+        if (policy == null) {
+            return; /* noop */
+        }
+
         final StackWalker walker = StackWalker.getInstance(Option.RETAIN_CLASS_REFERENCE);
         final Class<?> caller = walker.getCallerClass();
+        final Stream<Class<?>> chain = walker.walk(StackCallerClassChainExtractor.INSTANCE);
 
-        if (!AgentPolicy.isClassThatCanExit(caller.getName())) {
+        if (AgentPolicy.isChainThatCanExit(caller, chain) == false) {
             throw new SecurityException("The class " + caller + " is not allowed to call System::exit(" + code + ")");
         }
     }

libs/agent-sm/agent/src/test/java/org/opensearch/javaagent/SystemExitInterceptorTests.java → libs/agent-sm/agent/src/test/java/org/opensearch/javaagent/AgentTests.java

@@ -15,16 +15,21 @@
 import java.security.Policy;
 import java.util.Set;
 
-public class SystemExitInterceptorTests {
+public class AgentTests {
     @SuppressWarnings("removal")
     @BeforeClass
     public static void setUp() {
         AgentPolicy.setPolicy(new Policy() {
-        }, Set.of(), new String[] { "worker.org.gradle.process.internal.worker.GradleWorkerMain" });
+        }, Set.of(), (caller, chain) -> caller.getName().equalsIgnoreCase("worker.org.gradle.process.internal.worker.GradleWorkerMain"));
     }
 
     @Test(expected = SecurityException.class)
     public void testSystemExitIsForbidden() {
         System.exit(0);
     }
+
+    @Test(expected = SecurityException.class)
+    public void testRuntimeHaltIsForbidden() {
+        Runtime.getRuntime().halt(0);
+    }
 }

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

@@ -13,12 +13,13 @@
 import java.security.Permission;
 import java.security.Policy;
 import java.security.ProtectionDomain;
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
+import java.util.function.BiFunction;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * Agent Policy
@@ -28,7 +29,92 @@
     private static final Logger LOGGER = Logger.getLogger(AgentPolicy.class.getName());
     private static volatile Policy policy;
     private static volatile Set<String> trustedHosts;
-    private static volatile Set<String> classesThatCanExit;
+    private static volatile BiFunction<Class<?>, Stream<Class<?>>, Boolean> classesThatCanExit;
+
+    /**
+     * None of the classes is allowed to call {@link System#exit} or {@link Runtime#halt}
+     */
+    public static final class NoneCanExit implements BiFunction<Class<?>, Stream<Class<?>>, Boolean> {
+        /**
+         * NoneCanExit
+         */
+        public NoneCanExit() {}
+
+        /**
+         * Check if class is allowed to call {@link System#exit}, {@link Runtime#halt}
+         * @param caller caller class
+         * @param chain chain of call classes
+         * @return is class allowed to call {@link System#exit}, {@link Runtime#halt} or not
+         */
+        @Override
+        public Boolean apply(Class<?> caller, Stream<Class<?>> chain) {
+            return true;
+        }
+    }
+
+    /**
+     * Only caller is allowed to call {@link System#exit} or {@link Runtime#halt}
+     */
+    public static final class CallerCanExit implements BiFunction<Class<?>, Stream<Class<?>>, Boolean> {
+        private final String[] classesThatCanExit;
+
+        /**
+         * CallerCanExit
+         * @param classesThatCanExit classes that can exit
+         */
+        public CallerCanExit(final String[] classesThatCanExit) {
+            this.classesThatCanExit = classesThatCanExit;
+        }
+
+        /**
+         * Check if class is allowed to call {@link System#exit}, {@link Runtime#halt}
+         * @param caller caller class
+         * @param chain chain of call classes
+         * @return is class allowed to call {@link System#exit}, {@link Runtime#halt} or not
+         */
+        @Override
+        public Boolean apply(Class<?> caller, Stream<Class<?>> chain) {
+            for (final String classThatCanExit : classesThatCanExit) {
+                if (caller.getName().equalsIgnoreCase(classThatCanExit)) {
+                    return true;
+                }
+            }
+            return false;
+        }
+    }
+
+    /**
+     * Any caller in the chain is allowed to call {@link System#exit} or {@link Runtime#halt}
+     */
+    public static final class AnyCanExit implements BiFunction<Class<?>, Stream<Class<?>>, Boolean> {
+        private final String[] classesThatCanExit;
+
+        /**
+         * AnyCanExit
+         * @param classesThatCanExit classes that can exit
+         */
+        public AnyCanExit(final String[] classesThatCanExit) {
+            this.classesThatCanExit = classesThatCanExit;
+        }
+
+        /**
+         * Check if class is allowed to call {@link System#exit}, {@link Runtime#halt}
+         * @param caller caller class
+         * @param chain chain of call classes
+         * @return is class allowed to call {@link System#exit}, {@link Runtime#halt} or not
+         */
+        @Override
+        public Boolean apply(Class<?> caller, Stream<Class<?>> chain) {
+            return chain.anyMatch(clazz -> {
+                for (final String classThatCanExit : classesThatCanExit) {
+                    if (clazz.getName().matches(classThatCanExit)) {
+                        return true;
+                    }
+                }
+                return false;
+            });
+        }
+    }
 
     private AgentPolicy() {}
 
@@ -37,20 +123,24 @@
      * @param policy policy
      */
     public static void setPolicy(Policy policy) {
-        setPolicy(policy, Set.of(), new String[0]);
+        setPolicy(policy, Set.of(), new NoneCanExit());
     }
 
     /**
      * Set Agent policy
      * @param policy policy
      * @param trustedHosts trusted hosts
-     * @param classesThatCanExit classed that are allowed to call {@link System#exit}
+     * @param classesThatCanExit classed that are allowed to call {@link System#exit}, {@link Runtime#halt}
      */
-    public static void setPolicy(Policy policy, final Set<String> trustedHosts, final String[] classesThatCanExit) {
+    public static void setPolicy(
+        Policy policy,
+        final Set<String> trustedHosts,
+        final BiFunction<Class<?>, Stream<Class<?>>, Boolean> classesThatCanExit
+    ) {
         if (AgentPolicy.policy == null) {
             AgentPolicy.policy = policy;
             AgentPolicy.trustedHosts = Collections.unmodifiableSet(trustedHosts);
-            AgentPolicy.classesThatCanExit = Arrays.stream(classesThatCanExit).collect(Collectors.toSet());
+            AgentPolicy.classesThatCanExit = classesThatCanExit;
             LOGGER.info("Policy attached successfully: " + policy);
         } else {
             throw new SecurityException("The Policy has been set already: " + AgentPolicy.policy);
@@ -92,11 +182,12 @@
     }
 
     /**
-     * Check if class is allowed to call {@link System#exit}
-     * @param name class name
-     * @return is class allowed to call {@link System#exit} or not
+     * Check if class is allowed to call {@link System#exit}, {@link Runtime#halt}
+     * @param caller caller class
+     * @param chain chain of call classes
+     * @return is class allowed to call {@link System#exit}, {@link Runtime#halt} or not
      */
-    public static boolean isClassThatCanExit(String name) {
-        return AgentPolicy.classesThatCanExit.contains(name);
+    public static boolean isChainThatCanExit(Class<?> caller, Stream<Class<?>> chain) {
+        return classesThatCanExit.apply(caller, chain);
     }
 }