Enhance Java Agent to intercept System::exit (#17746)

* Enhance Java Agent to intercept System::exit

Signed-off-by: Andriy Redko <[email protected]>

* Address code review comments

Signed-off-by: Andriy Redko <[email protected]>

---------

Signed-off-by: Andriy Redko <[email protected]>

Author: reta

CHANGELOG.md

@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Added Kinesis support as a plugin for the pull-based ingestion ([#17615](https://github.com/opensearch-project/OpenSearch/pull/17615))
 - [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))
 
 ### Changed
 - Migrate BC libs to their FIPS counterparts ([#14912](https://github.com/opensearch-project/OpenSearch/pull/14912))

gradle/libs.versions.toml

@@ -72,7 +72,7 @@ junit             = "4.13.2"
 hamcrest          = "2.1"
 mockito           = "5.16.1"
 objenesis         = "3.3"
-bytebuddy         = "1.17.4"
+bytebuddy         = "1.17.5"
 
 # benchmark dependencies
 jmh               = "1.35"

libs/agent-sm/agent/build.gradle

@@ -13,6 +13,9 @@ dependencies {
   implementation project(":libs:agent-sm:bootstrap")
   implementation "net.bytebuddy:byte-buddy:${versions.bytebuddy}"
   compileOnly "com.google.code.findbugs:jsr305:3.0.2"
+
+  testImplementation "junit:junit:${versions.junit}"
+  testImplementation "org.hamcrest:hamcrest:${versions.hamcrest}"
 }
 
 var bootClasspath = configurations.bootstrap.incoming.artifactView { }.files
@@ -35,8 +38,8 @@ compileJava {
   options.compilerArgs -= '-Werror'
 }
 
-test.enabled = false
 testingConventions.enabled = false
+tasks.named('forbiddenApisTest').configure { onlyIf { false } }
 
 tasks.named('forbiddenApisMain').configure {
   replaceSignatureFiles 'jdk-signatures'
@@ -62,3 +65,12 @@ thirdPartyAudit {
 tasks.named('validateNebulaPom') {
   dependsOn prepareAgent
 }
+
+tasks.test {
+  dependsOn prepareAgent
+  jvmArgs += ["-javaagent:" + project.jar.archiveFile.get()]
+}
+
+tasks.check {
+  dependsOn test
+}

libs/agent-sm/agent/licenses/byte-buddy-1.17.4.jar.sha1

@@ -0,0 +1 @@
+88450f120903b7e72470462cdbd2b75a3842223c

libs/agent-sm/agent/licenses/byte-buddy-1.17.5.jar.sha1

@@ -77,7 +77,13 @@ private static AgentBuilder createAgentBuilder(Instrumentation inst) throws Exce
             .with(AgentBuilder.TypeStrategy.Default.REDEFINE)
             .ignore(ElementMatchers.none())
             .type(systemType)
-            .transform(transformer);
+            .transform(transformer)
+            .type(ElementMatchers.is(java.lang.System.class))
+            .transform(
+                (b, typeDescription, classLoader, module, pd) -> b.visit(
+                    Advice.to(SystemExitInterceptor.class).on(ElementMatchers.named("exit"))
+                )
+            );
 
         return agentBuilder;
     }

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

@@ -0,0 +1,40 @@
+/*
+ * 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 net.bytebuddy.asm.Advice;
+
+/**
+ * {@link System#exit} interceptor
+ */
+public class SystemExitInterceptor {
+    /**
+     * SystemExitInterceptor
+     */
+    public SystemExitInterceptor() {}
+
+    /**
+     * Interceptor
+     * @param code exit code
+     * @throws Exception exceptions
+     */
+    @Advice.OnMethodEnter()
+    public static void intercept(int code) throws Exception {
+        final StackWalker walker = StackWalker.getInstance(Option.RETAIN_CLASS_REFERENCE);
+        final Class<?> caller = walker.getCallerClass();
+
+        if (!AgentPolicy.isClassThatCanExit(caller.getName())) {
+            throw new SecurityException("The class " + caller + " is not allowed to call System::exit(" + code + ")");
+        }
+    }
+}

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

@@ -0,0 +1,30 @@
+/*
+ * 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 org.junit.BeforeClass;
+import org.junit.Test;
+
+import java.security.Policy;
+import java.util.Set;
+
+public class SystemExitInterceptorTests {
+    @SuppressWarnings("removal")
+    @BeforeClass
+    public static void setUp() {
+        AgentPolicy.setPolicy(new Policy() {
+        }, Set.of(), new String[] { "worker.org.gradle.process.internal.worker.GradleWorkerMain" });
+    }
+
+    @Test(expected = SecurityException.class)
+    public void testSystemExitIsForbidden() {
+        System.exit(0);
+    }
+}

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

@@ -13,6 +13,7 @@
 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;
@@ -27,6 +28,7 @@ public class AgentPolicy {
     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 AgentPolicy() {}
 
@@ -35,18 +37,20 @@ private AgentPolicy() {}
      * @param policy policy
      */
     public static void setPolicy(Policy policy) {
-        setPolicy(policy, Set.of());
+        setPolicy(policy, Set.of(), new String[0]);
     }
 
     /**
      * Set Agent policy
      * @param policy policy
      * @param trustedHosts trusted hosts
+     * @param classesThatCanExit classed that are allowed to call {@link System#exit}
      */
-    public static void setPolicy(Policy policy, final Set<String> trustedHosts) {
+    public static void setPolicy(Policy policy, final Set<String> trustedHosts, final String[] classesThatCanExit) {
         if (AgentPolicy.policy == null) {
             AgentPolicy.policy = policy;
             AgentPolicy.trustedHosts = Collections.unmodifiableSet(trustedHosts);
+            AgentPolicy.classesThatCanExit = Arrays.stream(classesThatCanExit).collect(Collectors.toSet());
             LOGGER.info("Policy attached successfully: " + policy);
         } else {
             throw new SecurityException("The Policy has been set already: " + AgentPolicy.policy);
@@ -86,4 +90,13 @@ public static Policy getPolicy() {
     public static boolean isTrustedHost(String hostname) {
         return AgentPolicy.trustedHosts.contains(hostname);
     }
+
+    /**
+     * 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
+     */
+    public static boolean isClassThatCanExit(String name) {
+        return AgentPolicy.classesThatCanExit.contains(name);
+    }
 }