Move logic to StackCallerProtectionDomainChainExtractor

Signed-off-by: Craig Perkins <[email protected]>

Author: cwperks

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

@@ -21,10 +21,6 @@
 import java.security.Policy;
 import java.security.ProtectionDomain;
 import java.util.Collection;
-import java.util.ArrayList;
-import java.util.List;
-import java.lang.StackWalker.StackFrame;
-import java.util.stream.Collectors;
 
 import net.bytebuddy.asm.Advice;
 
@@ -70,23 +66,7 @@ public static void intercept(@Advice.AllArguments Object[] args, @Advice.Origin
         }
 
         final StackWalker walker = StackWalker.getInstance(StackWalker.Option.RETAIN_CLASS_REFERENCE);
-        final Collection<ProtectionDomain> callers = walker.walk(s -> {
-            List<ProtectionDomain> domains = new ArrayList<>();
-            boolean foundPrivileged = false;
-
-            for (StackFrame frame : s.toList()) {
-                if (frame.getClassName().equals("java.security.AccessController") &&
-                    frame.getMethodName().equals("doPrivileged")) {
-                    foundPrivileged = true;
-                    break;
-                }
-                Class<?> callerClass = frame.getDeclaringClass();
-                domains.add(callerClass.getProtectionDomain());
-            }
-
-            return foundPrivileged ? domains : s.map(f -> f.getDeclaringClass().getProtectionDomain())
-                .collect(Collectors.toList());
-        });
+        final Collection<ProtectionDomain> callers = walker.walk(StackCallerProtectionDomainChainExtractor.INSTANCE);
 
         final String name = method.getName();
         boolean isMutating = name.equals("move") || name.equals("write") || name.startsWith("create");

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

@@ -10,7 +10,6 @@
 
 import org.opensearch.javaagent.bootstrap.AgentPolicy;
 
-import java.lang.StackWalker.Option;
 import java.lang.reflect.Method;
 import java.net.InetSocketAddress;
 import java.net.NetPermission;
@@ -19,10 +18,6 @@
 import java.security.Policy;
 import java.security.ProtectionDomain;
 import java.util.Collection;
-import java.util.ArrayList;
-import java.util.List;
-import java.lang.StackWalker.StackFrame;
-import java.util.stream.Collectors;
 
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.asm.Advice.Origin;
@@ -51,23 +46,7 @@ public static void intercept(@Advice.AllArguments Object[] args, @Origin Method
         }
 
         final StackWalker walker = StackWalker.getInstance(StackWalker.Option.RETAIN_CLASS_REFERENCE);
-        final Collection<ProtectionDomain> callers = walker.walk(s -> {
-            List<ProtectionDomain> domains = new ArrayList<>();
-            boolean foundPrivileged = false;
-
-            for (StackFrame frame : s.toList()) {
-                if (frame.getClassName().equals("java.security.AccessController") &&
-                    frame.getMethodName().equals("doPrivileged")) {
-                    foundPrivileged = true;
-                    break;
-                }
-                Class<?> callerClass = frame.getDeclaringClass();
-                domains.add(callerClass.getProtectionDomain());
-            }
-
-            return foundPrivileged ? domains : s.map(f -> f.getDeclaringClass().getProtectionDomain())
-                .collect(Collectors.toList());
-        });
+        final Collection<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/StackCallerProtectionDomainChainExtractor.java

@@ -35,9 +35,12 @@ private StackCallerProtectionDomainChainExtractor() {}
      */
     @Override
     public Collection<ProtectionDomain> apply(Stream<StackFrame> frames) {
-        return frames.map(StackFrame::getDeclaringClass)
+        return frames.takeWhile(
+            frame -> !frame.getClassName().equals("java.security.AccessController") || !frame.getMethodName().equals("doPrivileged")
+        )
+            .map(StackFrame::getDeclaringClass)
             .map(Class::getProtectionDomain)
-            .filter(pd -> pd.getCodeSource() != null) /* JDK */
+            .filter(pd -> pd.getCodeSource() != null) // Filter out JDK classes
             .collect(Collectors.toSet());
     }
 }