Skip to content

Commit cf31931

Browse files
retareta
authored
Enhance Java Agent to intercept Runtime::halt (#17757)
Signed-off-by: Andriy Redko <[email protected]>
1 parent 5f1b4ab commit cf31931

File tree

9 files changed

+225
-20
lines changed

9 files changed

+225
-20
lines changed

‎CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
1717
- [Security Manager Replacement] Create initial Java Agent to intercept Socket::connect calls ([#17724](https://github.com/opensearch-project/OpenSearch/pull/17724))
1818
- Add ingestion management APIs for pause, resume and get ingestion state ([#17631](https://github.com/opensearch-project/OpenSearch/pull/17631))
1919
- [Security Manager Replacement] Enhance Java Agent to intercept System::exit ([#17746](https://github.com/opensearch-project/OpenSearch/pull/17746))
20+
- [Security Manager Replacement] Enhance Java Agent to intercept Runtime::halt ([#17757](https://github.com/opensearch-project/OpenSearch/pull/17757))
2021
- Support AutoExpand for SearchReplica ([#17741](https://github.com/opensearch-project/OpenSearch/pull/17741))
2122
- Implement fixed interval refresh task scheduling ([#17777](https://github.com/opensearch-project/OpenSearch/pull/17777))
2223

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,10 @@ private static AgentBuilder createAgentBuilder(Instrumentation inst) throws Exce
6464
ClassInjector.UsingUnsafe.ofBootLoader()
6565
.inject(
6666
Map.of(
67-
new TypeDescription.ForLoadedType(StackCallerChainExtractor.class),
68-
ClassFileLocator.ForClassLoader.read(StackCallerChainExtractor.class),
67+
new TypeDescription.ForLoadedType(StackCallerProtectionDomainChainExtractor.class),
68+
ClassFileLocator.ForClassLoader.read(StackCallerProtectionDomainChainExtractor.class),
69+
new TypeDescription.ForLoadedType(StackCallerClassChainExtractor.class),
70+
ClassFileLocator.ForClassLoader.read(StackCallerClassChainExtractor.class),
6971
new TypeDescription.ForLoadedType(AgentPolicy.class),
7072
ClassFileLocator.ForClassLoader.read(AgentPolicy.class)
7173
)
@@ -83,6 +85,12 @@ private static AgentBuilder createAgentBuilder(Instrumentation inst) throws Exce
8385
(b, typeDescription, classLoader, module, pd) -> b.visit(
8486
Advice.to(SystemExitInterceptor.class).on(ElementMatchers.named("exit"))
8587
)
88+
)
89+
.type(ElementMatchers.is(java.lang.Runtime.class))
90+
.transform(
91+
(b, typeDescription, classLoader, module, pd) -> b.visit(
92+
Advice.to(RuntimeHaltInterceptor.class).on(ElementMatchers.named("halt"))
93+
)
8694
);
8795

8896
return agentBuilder;

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

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.javaagent;
10+
11+
import org.opensearch.javaagent.bootstrap.AgentPolicy;
12+
13+
import java.lang.StackWalker.Option;
14+
import java.security.Policy;
15+
import java.util.stream.Stream;
16+
17+
import net.bytebuddy.asm.Advice;
18+
19+
/**
20+
* {@link Runtime#halt} interceptor
21+
*/
22+
public class RuntimeHaltInterceptor {
23+
/**
24+
* RuntimeHaltInterceptor
25+
*/
26+
public RuntimeHaltInterceptor() {}
27+
28+
/**
29+
* Interceptor
30+
* @param code exit code
31+
* @throws Exception exceptions
32+
*/
33+
@Advice.OnMethodEnter
34+
@SuppressWarnings("removal")
35+
public static void intercept(int code) throws Exception {
36+
final Policy policy = AgentPolicy.getPolicy();
37+
if (policy == null) {
38+
return; /* noop */
39+
}
40+
41+
final StackWalker walker = StackWalker.getInstance(Option.RETAIN_CLASS_REFERENCE);
42+
final Class<?> caller = walker.getCallerClass();
43+
final Stream<Class<?>> chain = walker.walk(StackCallerClassChainExtractor.INSTANCE);
44+
45+
if (AgentPolicy.isChainThatCanExit(caller, chain) == false) {
46+
throw new SecurityException("The class " + caller + " is not allowed to call Runtime::halt(" + code + ")");
47+
}
48+
}
49+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ public static void intercept(@Advice.AllArguments Object[] args, @Origin Method
4747
}
4848

4949
final StackWalker walker = StackWalker.getInstance(Option.RETAIN_CLASS_REFERENCE);
50-
final Stream<ProtectionDomain> callers = walker.walk(StackCallerChainExtractor.INSTANCE);
50+
final Stream<ProtectionDomain> callers = walker.walk(StackCallerProtectionDomainChainExtractor.INSTANCE);
5151

5252
if (args[0] instanceof InetSocketAddress address) {
5353
if (!AgentPolicy.isTrustedHost(address.getHostString())) {

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

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.javaagent;
10+
11+
import java.lang.StackWalker.StackFrame;
12+
import java.util.function.Function;
13+
import java.util.stream.Stream;
14+
15+
/**
16+
* Stack Caller Class Chain Extractor
17+
*/
18+
public final class StackCallerClassChainExtractor implements Function<Stream<StackFrame>, Stream<Class<?>>> {
19+
/**
20+
* Single instance of stateless class.
21+
*/
22+
public static final StackCallerClassChainExtractor INSTANCE = new StackCallerClassChainExtractor();
23+
24+
/**
25+
* Constructor
26+
*/
27+
private StackCallerClassChainExtractor() {}
28+
29+
/**
30+
* Folds the stack
31+
* @param frames stack frames
32+
*/
33+
@Override
34+
public Stream<Class<?>> apply(Stream<StackFrame> frames) {
35+
return cast(frames);
36+
}
37+
38+
@SuppressWarnings("unchecked")
39+
private static <A> Stream<A> cast(Stream<StackFrame> frames) {
40+
return (Stream<A>) frames.map(StackFrame::getDeclaringClass).filter(c -> !c.isHidden()).distinct();
41+
}
42+
}

‎libs/agent-sm/agent/src/main/java/org/opensearch/javaagent/StackCallerChainExtractor.java renamed to ‎libs/agent-sm/agent/src/main/java/org/opensearch/javaagent/StackCallerProtectionDomainChainExtractor.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,16 +16,16 @@
1616
/**
1717
* Stack Caller Chain Extractor
1818
*/
19-
public final class StackCallerChainExtractor implements Function<Stream<StackFrame>, Stream<ProtectionDomain>> {
19+
public final class StackCallerProtectionDomainChainExtractor implements Function<Stream<StackFrame>, Stream<ProtectionDomain>> {
2020
/**
2121
* Single instance of stateless class.
2222
*/
23-
public static final StackCallerChainExtractor INSTANCE = new StackCallerChainExtractor();
23+
public static final StackCallerProtectionDomainChainExtractor INSTANCE = new StackCallerProtectionDomainChainExtractor();
2424

2525
/**
2626
* Constructor
2727
*/
28-
private StackCallerChainExtractor() {}
28+
private StackCallerProtectionDomainChainExtractor() {}
2929

3030
/**
3131
* Folds the stack

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
import org.opensearch.javaagent.bootstrap.AgentPolicy;
1212

1313
import java.lang.StackWalker.Option;
14+
import java.security.Policy;
15+
import java.util.stream.Stream;
1416

1517
import net.bytebuddy.asm.Advice;
1618

@@ -29,11 +31,18 @@ public SystemExitInterceptor() {}
2931
* @throws Exception exceptions
3032
*/
3133
@Advice.OnMethodEnter()
34+
@SuppressWarnings("removal")
3235
public static void intercept(int code) throws Exception {
36+
final Policy policy = AgentPolicy.getPolicy();
37+
if (policy == null) {
38+
return; /* noop */
39+
}
40+
3341
final StackWalker walker = StackWalker.getInstance(Option.RETAIN_CLASS_REFERENCE);
3442
final Class<?> caller = walker.getCallerClass();
43+
final Stream<Class<?>> chain = walker.walk(StackCallerClassChainExtractor.INSTANCE);
3544

36-
if (!AgentPolicy.isClassThatCanExit(caller.getName())) {
45+
if (AgentPolicy.isChainThatCanExit(caller, chain) == false) {
3746
throw new SecurityException("The class " + caller + " is not allowed to call System::exit(" + code + ")");
3847
}
3948
}

‎libs/agent-sm/agent/src/test/java/org/opensearch/javaagent/SystemExitInterceptorTests.java renamed to ‎libs/agent-sm/agent/src/test/java/org/opensearch/javaagent/AgentTests.java

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,21 @@
1515
import java.security.Policy;
1616
import java.util.Set;
1717

18-
public class SystemExitInterceptorTests {
18+
public class AgentTests {
1919
@SuppressWarnings("removal")
2020
@BeforeClass
2121
public static void setUp() {
2222
AgentPolicy.setPolicy(new Policy() {
23-
}, Set.of(), new String[] { "worker.org.gradle.process.internal.worker.GradleWorkerMain" });
23+
}, Set.of(), (caller, chain) -> caller.getName().equalsIgnoreCase("worker.org.gradle.process.internal.worker.GradleWorkerMain"));
2424
}
2525

2626
@Test(expected = SecurityException.class)
2727
public void testSystemExitIsForbidden() {
2828
System.exit(0);
2929
}
30+
31+
@Test(expected = SecurityException.class)
32+
public void testRuntimeHaltIsForbidden() {
33+
Runtime.getRuntime().halt(0);
34+
}
3035
}

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

Lines changed: 102 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,12 +13,13 @@
1313
import java.security.Permission;
1414
import java.security.Policy;
1515
import java.security.ProtectionDomain;
16-
import java.util.Arrays;
1716
import java.util.Collections;
1817
import java.util.List;
1918
import java.util.Set;
19+
import java.util.function.BiFunction;
2020
import java.util.logging.Logger;
2121
import java.util.stream.Collectors;
22+
import java.util.stream.Stream;
2223

2324
/**
2425
* Agent Policy
@@ -28,7 +29,92 @@ public class AgentPolicy {
2829
private static final Logger LOGGER = Logger.getLogger(AgentPolicy.class.getName());
2930
private static volatile Policy policy;
3031
private static volatile Set<String> trustedHosts;
31-
private static volatile Set<String> classesThatCanExit;
32+
private static volatile BiFunction<Class<?>, Stream<Class<?>>, Boolean> classesThatCanExit;
33+
34+
/**
35+
* None of the classes is allowed to call {@link System#exit} or {@link Runtime#halt}
36+
*/
37+
public static final class NoneCanExit implements BiFunction<Class<?>, Stream<Class<?>>, Boolean> {
38+
/**
39+
* NoneCanExit
40+
*/
41+
public NoneCanExit() {}
42+
43+
/**
44+
* Check if class is allowed to call {@link System#exit}, {@link Runtime#halt}
45+
* @param caller caller class
46+
* @param chain chain of call classes
47+
* @return is class allowed to call {@link System#exit}, {@link Runtime#halt} or not
48+
*/
49+
@Override
50+
public Boolean apply(Class<?> caller, Stream<Class<?>> chain) {
51+
return true;
52+
}
53+
}
54+
55+
/**
56+
* Only caller is allowed to call {@link System#exit} or {@link Runtime#halt}
57+
*/
58+
public static final class CallerCanExit implements BiFunction<Class<?>, Stream<Class<?>>, Boolean> {
59+
private final String[] classesThatCanExit;
60+
61+
/**
62+
* CallerCanExit
63+
* @param classesThatCanExit classes that can exit
64+
*/
65+
public CallerCanExit(final String[] classesThatCanExit) {
66+
this.classesThatCanExit = classesThatCanExit;
67+
}
68+
69+
/**
70+
* Check if class is allowed to call {@link System#exit}, {@link Runtime#halt}
71+
* @param caller caller class
72+
* @param chain chain of call classes
73+
* @return is class allowed to call {@link System#exit}, {@link Runtime#halt} or not
74+
*/
75+
@Override
76+
public Boolean apply(Class<?> caller, Stream<Class<?>> chain) {
77+
for (final String classThatCanExit : classesThatCanExit) {
78+
if (caller.getName().equalsIgnoreCase(classThatCanExit)) {
79+
return true;
80+
}
81+
}
82+
return false;
83+
}
84+
}
85+
86+
/**
87+
* Any caller in the chain is allowed to call {@link System#exit} or {@link Runtime#halt}
88+
*/
89+
public static final class AnyCanExit implements BiFunction<Class<?>, Stream<Class<?>>, Boolean> {
90+
private final String[] classesThatCanExit;
91+
92+
/**
93+
* AnyCanExit
94+
* @param classesThatCanExit classes that can exit
95+
*/
96+
public AnyCanExit(final String[] classesThatCanExit) {
97+
this.classesThatCanExit = classesThatCanExit;
98+
}
99+
100+
/**
101+
* Check if class is allowed to call {@link System#exit}, {@link Runtime#halt}
102+
* @param caller caller class
103+
* @param chain chain of call classes
104+
* @return is class allowed to call {@link System#exit}, {@link Runtime#halt} or not
105+
*/
106+
@Override
107+
public Boolean apply(Class<?> caller, Stream<Class<?>> chain) {
108+
return chain.anyMatch(clazz -> {
109+
for (final String classThatCanExit : classesThatCanExit) {
110+
if (clazz.getName().matches(classThatCanExit)) {
111+
return true;
112+
}
113+
}
114+
return false;
115+
});
116+
}
117+
}
32118

33119
private AgentPolicy() {}
34120

@@ -37,20 +123,24 @@ private AgentPolicy() {}
37123
* @param policy policy
38124
*/
39125
public static void setPolicy(Policy policy) {
40-
setPolicy(policy, Set.of(), new String[0]);
126+
setPolicy(policy, Set.of(), new NoneCanExit());
41127
}
42128

43129
/**
44130
* Set Agent policy
45131
* @param policy policy
46132
* @param trustedHosts trusted hosts
47-
* @param classesThatCanExit classed that are allowed to call {@link System#exit}
133+
* @param classesThatCanExit classed that are allowed to call {@link System#exit}, {@link Runtime#halt}
48134
*/
49-
public static void setPolicy(Policy policy, final Set<String> trustedHosts, final String[] classesThatCanExit) {
135+
public static void setPolicy(
136+
Policy policy,
137+
final Set<String> trustedHosts,
138+
final BiFunction<Class<?>, Stream<Class<?>>, Boolean> classesThatCanExit
139+
) {
50140
if (AgentPolicy.policy == null) {
51141
AgentPolicy.policy = policy;
52142
AgentPolicy.trustedHosts = Collections.unmodifiableSet(trustedHosts);
53-
AgentPolicy.classesThatCanExit = Arrays.stream(classesThatCanExit).collect(Collectors.toSet());
143+
AgentPolicy.classesThatCanExit = classesThatCanExit;
54144
LOGGER.info("Policy attached successfully: " + policy);
55145
} else {
56146
throw new SecurityException("The Policy has been set already: " + AgentPolicy.policy);
@@ -92,11 +182,12 @@ public static boolean isTrustedHost(String hostname) {
92182
}
93183

94184
/**
95-
* Check if class is allowed to call {@link System#exit}
96-
* @param name class name
97-
* @return is class allowed to call {@link System#exit} or not
185+
* Check if class is allowed to call {@link System#exit}, {@link Runtime#halt}
186+
* @param caller caller class
187+
* @param chain chain of call classes
188+
* @return is class allowed to call {@link System#exit}, {@link Runtime#halt} or not
98189
*/
99-
public static boolean isClassThatCanExit(String name) {
100-
return AgentPolicy.classesThatCanExit.contains(name);
190+
public static boolean isChainThatCanExit(Class<?> caller, Stream<Class<?>> chain) {
191+
return classesThatCanExit.apply(caller, chain);
101192
}
102193
}

0 commit comments

Comments
 (0)