Fix StackCallerClassChainExtractor and respectively RuntimeHaltInterceptor, SystemExitInterceptor

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

@@ -12,7 +12,7 @@
 
 import java.lang.StackWalker.Option;
 import java.security.Policy;
-import java.util.stream.Stream;
+import java.util.Collection;
 
 import net.bytebuddy.asm.Advice;
 
@@ -40,7 +40,7 @@
 
         final StackWalker walker = StackWalker.getInstance(Option.RETAIN_CLASS_REFERENCE);
         final Class<?> caller = walker.getCallerClass();
-        final Stream<Class<?>> chain = walker.walk(StackCallerClassChainExtractor.INSTANCE);
+        final Collection<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/StackCallerClassChainExtractor.java

@@ -9,13 +9,16 @@
 package org.opensearch.javaagent;
 
 import java.lang.StackWalker.StackFrame;
+import java.util.Collection;
+import java.util.Set;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 /**
  * Stack Caller Class Chain Extractor
  */
-public final class StackCallerClassChainExtractor implements Function<Stream<StackFrame>, Stream<Class<?>>> {
+public final class StackCallerClassChainExtractor implements Function<Stream<StackFrame>, Collection<Class<?>>> {
     /**
      * Single instance of stateless class.
      */
@@ -31,12 +34,12 @@ private StackCallerClassChainExtractor() {}
      * @param frames stack frames
      */
     @Override
-    public Stream<Class<?>> apply(Stream<StackFrame> frames) {
+    public Collection<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();
+    private static <A> Set<A> cast(Stream<StackFrame> frames) {
+        return (Set<A>) frames.map(StackFrame::getDeclaringClass).filter(c -> !c.isHidden()).collect(Collectors.toSet());
     }
 }

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

@@ -12,7 +12,7 @@
 
 import java.lang.StackWalker.Option;
 import java.security.Policy;
-import java.util.stream.Stream;
+import java.util.Collection;
 
 import net.bytebuddy.asm.Advice;
 
@@ -40,7 +40,7 @@
 
         final StackWalker walker = StackWalker.getInstance(Option.RETAIN_CLASS_REFERENCE);
         final Class<?> caller = walker.getCallerClass();
-        final Stream<Class<?>> chain = walker.walk(StackCallerClassChainExtractor.INSTANCE);
+        final Collection<Class<?>> chain = walker.walk(StackCallerClassChainExtractor.INSTANCE);
 
         if (AgentPolicy.isChainThatCanExit(caller, chain) == false) {
             throw new SecurityException("The class " + caller + " is not allowed to call System::exit(" + code + ")");

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

@@ -13,6 +13,7 @@
 import java.security.Permission;
 import java.security.Policy;
 import java.security.ProtectionDomain;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
@@ -29,12 +30,12 @@
     private static final Logger LOGGER = Logger.getLogger(AgentPolicy.class.getName());
     private static volatile Policy policy;
     private static volatile Set<String> trustedHosts;
-    private static volatile BiFunction<Class<?>, Stream<Class<?>>, Boolean> classesThatCanExit;
+    private static volatile BiFunction<Class<?>, Collection<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> {
+    public static final class NoneCanExit implements BiFunction<Class<?>, Collection<Class<?>>, Boolean> {
         /**
          * NoneCanExit
          */
@@ -47,7 +48,7 @@
          * @return is class allowed to call {@link System#exit}, {@link Runtime#halt} or not
          */
         @Override
-        public Boolean apply(Class<?> caller, Stream<Class<?>> chain) {
+        public Boolean apply(Class<?> caller, Collection<Class<?>> chain) {
             return true;
         }
     }
@@ -86,7 +87,7 @@
     /**
      * 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> {
+    public static final class AnyCanExit implements BiFunction<Class<?>, Collection<Class<?>>, Boolean> {
         private final String[] classesThatCanExit;
 
         /**
@@ -104,15 +105,15 @@
          * @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 -> {
+        public Boolean apply(Class<?> caller, Collection<Class<?>> chain) {
+            for (final Class<?> clazz : chain) {
                 for (final String classThatCanExit : classesThatCanExit) {
                     if (clazz.getName().matches(classThatCanExit)) {
                         return true;
                     }
                 }
-                return false;
-            });
+            }
+            return false;
         }
     }
 
@@ -135,7 +136,7 @@
     public static void setPolicy(
         Policy policy,
         final Set<String> trustedHosts,
-        final BiFunction<Class<?>, Stream<Class<?>>, Boolean> classesThatCanExit
+        final BiFunction<Class<?>, Collection<Class<?>>, Boolean> classesThatCanExit
     ) {
         if (AgentPolicy.policy == null) {
             AgentPolicy.policy = policy;
@@ -187,7 +188,7 @@
      * @param chain chain of call classes
      * @return is class allowed to call {@link System#exit}, {@link Runtime#halt} or not
      */
-    public static boolean isChainThatCanExit(Class<?> caller, Stream<Class<?>> chain) {
+    public static boolean isChainThatCanExit(Class<?> caller, Collection<Class<?>> chain) {
         return classesThatCanExit.apply(caller, chain);
     }
 }