Intercept java.net.Socket::connect

Signed-off-by: Andrew Ross <andrross@amazon.com>

Author: andrross

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

@@ -11,6 +11,7 @@
 import org.opensearch.javaagent.bootstrap.AgentPolicy;
 
 import java.lang.instrument.Instrumentation;
+import java.net.Socket;
 import java.nio.channels.FileChannel;
 import java.nio.channels.SocketChannel;
 import java.nio.file.Files;
@@ -71,8 +72,9 @@ public static void agentmain(String agentArguments, Instrumentation instrumentat
         initAgent(instrumentation);
     }
 
-    private static AgentBuilder createAgentBuilder(Instrumentation inst) throws Exception {
-        final Junction<TypeDescription> systemType = ElementMatchers.isSubTypeOf(SocketChannel.class);
+    private static AgentBuilder createAgentBuilder() throws Exception {
+        final Junction<TypeDescription> socketType = ElementMatchers.isSubTypeOf(SocketChannel.class)
+            .or(ElementMatchers.isSubTypeOf(Socket.class));
         final Junction<TypeDescription> pathType = ElementMatchers.isSubTypeOf(Files.class);
         final Junction<TypeDescription> fileChannelType = ElementMatchers.isSubTypeOf(FileChannel.class);
 
@@ -98,11 +100,11 @@ private static AgentBuilder createAgentBuilder(Instrumentation inst) throws Exce
             );
 
         final ByteBuddy byteBuddy = new ByteBuddy().with(Implementation.Context.Disabled.Factory.INSTANCE);
-        final AgentBuilder agentBuilder = new AgentBuilder.Default(byteBuddy).with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE)
+        return new AgentBuilder.Default(byteBuddy).with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE)
             .with(AgentBuilder.RedefinitionStrategy.REDEFINITION)
             .with(AgentBuilder.TypeStrategy.Default.REDEFINE)
             .ignore(ElementMatchers.nameContains("$MockitoMock$")) /* ingore all Mockito mocks */
-            .type(systemType)
+            .type(socketType)
             .transform(socketTransformer)
             .type(pathType.or(fileChannelType))
             .transform(fileTransformer)
@@ -118,12 +120,10 @@ private static AgentBuilder createAgentBuilder(Instrumentation inst) throws Exce
                     Advice.to(RuntimeHaltInterceptor.class).on(ElementMatchers.named("halt"))
                 )
             );
-
-        return agentBuilder;
     }
 
     private static void initAgent(Instrumentation instrumentation) throws Exception {
-        AgentBuilder agentBuilder = createAgentBuilder(instrumentation);
+        AgentBuilder agentBuilder = createAgentBuilder();
         agentBuilder.installOn(instrumentation);
     }
 }

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

@@ -13,6 +13,7 @@
 import java.io.IOException;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
+import java.net.Socket;
 import java.net.UnixDomainSocketAddress;
 import java.nio.channels.SocketChannel;
 
@@ -28,6 +29,8 @@ public void testConnections() throws IOException {
 
             assertThrows(SecurityException.class, () -> channel.connect(new InetSocketAddress("opensearch.org", 80)));
         }
+
+        assertThrows(SecurityException.class, () -> new Socket("localhost", 9200));
     }
 
     @Test