Fix socket interceptor and add unit test (#17784)

The optimization of using a Stream instead of a Collection caused problems
with class not found and/or illegal access errors when using the lambda function
in the `Stream::forEach` call in the intercept method.

Signed-off-by: Andrew Ross <[email protected]>

Author: andrross

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

@@ -18,7 +18,7 @@
 import java.net.UnixDomainSocketAddress;
 import java.security.Policy;
 import java.security.ProtectionDomain;
-import java.util.stream.Stream;
+import java.util.Collection;
 
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.asm.Advice.Origin;
@@ -47,26 +47,26 @@ public static void intercept(@Advice.AllArguments Object[] args, @Origin Method
         }
 
         final StackWalker walker = StackWalker.getInstance(Option.RETAIN_CLASS_REFERENCE);
-        final Stream<ProtectionDomain> callers = walker.walk(StackCallerProtectionDomainChainExtractor.INSTANCE);
+        final Collection<ProtectionDomain> callers = walker.walk(StackCallerProtectionDomainChainExtractor.INSTANCE);
 
         if (args[0] instanceof InetSocketAddress address) {
             if (!AgentPolicy.isTrustedHost(address.getHostString())) {
                 final String host = address.getHostString() + ":" + address.getPort();
 
                 final SocketPermission permission = new SocketPermission(host, "connect,resolve");
-                callers.forEach(domain -> {
+                for (ProtectionDomain domain : callers) {
                     if (!policy.implies(domain, permission)) {
                         throw new SecurityException("Denied access to: " + host + ", domain " + domain);
                     }
-                });
+                }
             }
         } else if (args[0] instanceof UnixDomainSocketAddress address) {
             final NetPermission permission = new NetPermission("accessUnixDomainSocket");
-            callers.forEach(domain -> {
+            for (ProtectionDomain domain : callers) {
                 if (!policy.implies(domain, permission)) {
                     throw new SecurityException("Denied access to: " + address + ", domain " + domain);
                 }
-            });
+            }
         }
     }
 }

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

@@ -10,13 +10,15 @@
 
 import java.lang.StackWalker.StackFrame;
 import java.security.ProtectionDomain;
+import java.util.Collection;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 /**
  * Stack Caller Chain Extractor
  */
-public final class StackCallerProtectionDomainChainExtractor implements Function<Stream<StackFrame>, Stream<ProtectionDomain>> {
+public final class StackCallerProtectionDomainChainExtractor implements Function<Stream<StackFrame>, Collection<ProtectionDomain>> {
     /**
      * Single instance of stateless class.
      */
@@ -32,10 +34,10 @@ private StackCallerProtectionDomainChainExtractor() {}
      * @param frames stack frames
      */
     @Override
-    public Stream<ProtectionDomain> apply(Stream<StackFrame> frames) {
+    public Collection<ProtectionDomain> apply(Stream<StackFrame> frames) {
         return frames.map(StackFrame::getDeclaringClass)
             .map(Class::getProtectionDomain)
             .filter(pd -> pd.getCodeSource() != null) /* JDK */
-            .distinct();
+            .collect(Collectors.toSet());
     }
 }

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

@@ -0,0 +1,24 @@
+/*
+ * 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 java.security.Policy;
+import java.util.Set;
+
+public abstract class AgentTestCase {
+    @SuppressWarnings("removal")
+    @BeforeClass
+    public static void setUp() {
+        AgentPolicy.setPolicy(new Policy() {
+        }, Set.of(), (caller, chain) -> caller.getName().equalsIgnoreCase("worker.org.gradle.process.internal.worker.GradleWorkerMain"));
+    }
+}

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

@@ -8,21 +8,9 @@
 
 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 AgentTests {
-    @SuppressWarnings("removal")
-    @BeforeClass
-    public static void setUp() {
-        AgentPolicy.setPolicy(new Policy() {
-        }, Set.of(), (caller, chain) -> caller.getName().equalsIgnoreCase("worker.org.gradle.process.internal.worker.GradleWorkerMain"));
-    }
-
+public class AgentTests extends AgentTestCase {
     @Test(expected = SecurityException.class)
     public void testSystemExitIsForbidden() {
         System.exit(0);

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

@@ -0,0 +1,29 @@
+/*
+ * 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.junit.Test;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.net.UnixDomainSocketAddress;
+import java.nio.channels.SocketChannel;
+
+import static org.junit.Assert.assertThrows;
+
+public class SocketChannelInterceptorTests extends AgentTestCase {
+    @Test
+    public void test() throws IOException {
+        try (SocketChannel channel = SocketChannel.open()) {
+            assertThrows(SecurityException.class, () -> channel.connect(new InetSocketAddress("localhost", 9200)));
+
+            assertThrows(SecurityException.class, () -> channel.connect(UnixDomainSocketAddress.of("fake-path")));
+        }
+    }
+}