[Entitlements] Add missing NIO async network instrumentation

libs/entitlement/bridge/src/main/java/org/elasticsearch/entitlement/bridge/EntitlementChecker.java

@@ -53,6 +53,8 @@
 import java.nio.channels.AsynchronousSocketChannel;
 import java.nio.channels.CompletionHandler;
 import java.nio.channels.DatagramChannel;
+import java.nio.channels.SelectableChannel;
+import java.nio.channels.Selector;
 import java.nio.channels.ServerSocketChannel;
 import java.nio.channels.SocketChannel;
 import java.nio.channels.spi.SelectorProvider;
@@ -588,6 +590,16 @@ public interface EntitlementChecker {
      * (not instrumentable).
      */
 
+    void check$java_nio_channels_spi_AbstractSelectableChannel$register(
+        Class<?> callerClass,
+        SelectableChannel that,
+        Selector sel,
+        int ops,
+        Object att
+    );
+
+    void check$java_nio_channels_SelectableChannel$register(Class<?> callerClass, SelectableChannel that, Selector sel, int ops);
+
     // bind
 
     void check$java_nio_channels_AsynchronousServerSocketChannel$bind(
@@ -611,6 +623,12 @@ public interface EntitlementChecker {
 
     void check$sun_nio_ch_ServerSocketChannelImpl$bind(Class<?> callerClass, ServerSocketChannel that, SocketAddress local, int backlog);
 
+    void check$java_nio_channels_SocketChannel$$open(Class<?> callerClass);
+
+    void check$java_nio_channels_SocketChannel$$open(Class<?> callerClass, java.net.ProtocolFamily family);
+
+    void check$java_nio_channels_SocketChannel$$open(Class<?> callerClass, SocketAddress remote);
+
     void check$sun_nio_ch_SocketChannelImpl$bind(Class<?> callerClass, SocketChannel that, SocketAddress local);
 
     // connect
@@ -658,6 +676,18 @@ public interface EntitlementChecker {
     // provider methods (dynamic)
     void checkSelectorProviderInheritedChannel(Class<?> callerClass, SelectorProvider that);
 
+    void checkSelectorProviderOpenDatagramChannel(Class<?> callerClass, SelectorProvider that);
+
+    void checkSelectorProviderOpenDatagramChannel(Class<?> callerClass, SelectorProvider that, java.net.ProtocolFamily family);
+
+    void checkSelectorProviderOpenServerSocketChannel(Class<?> callerClass, SelectorProvider that);
+
+    void checkSelectorProviderOpenServerSocketChannel(Class<?> callerClass, SelectorProvider that, java.net.ProtocolFamily family);
+
+    void checkSelectorProviderOpenSocketChannel(Class<?> callerClass, SelectorProvider that);
+
+    void checkSelectorProviderOpenSocketChannel(Class<?> callerClass, SelectorProvider that, java.net.ProtocolFamily family);
+
     /// /////////////////
     //
     // Load native libraries

libs/entitlement/qa/entitlement-test-plugin/src/main/java/org/elasticsearch/entitlement/qa/test/DummyImplementations.java

@@ -13,6 +13,7 @@
 
 import org.elasticsearch.core.SuppressForbidden;
 
+import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.net.DatagramPacket;
@@ -41,9 +42,12 @@
 import java.nio.channels.ReadableByteChannel;
 import java.nio.channels.SeekableByteChannel;
 import java.nio.channels.SelectableChannel;
+import java.nio.channels.SelectionKey;
+import java.nio.channels.Selector;
 import java.nio.channels.ServerSocketChannel;
 import java.nio.channels.SocketChannel;
 import java.nio.channels.WritableByteChannel;
+import java.nio.channels.spi.AbstractSelectableChannel;
 import java.nio.channels.spi.AbstractSelector;
 import java.nio.channels.spi.AsynchronousChannelProvider;
 import java.nio.channels.spi.SelectorProvider;
@@ -846,4 +850,71 @@ public void implCloseChannel(SelectableChannel sc) {}
         @Override
         public void implReleaseChannel(SelectableChannel sc) {}
     }
+
+    static class DummySelectableChannel extends AbstractSelectableChannel {
+        protected DummySelectableChannel(SelectorProvider provider) {
+            super(provider);
+        }
+
+        @Override
+        protected void implCloseSelectableChannel() throws IOException {
+
+        }
+
+        @Override
+        protected void implConfigureBlocking(boolean block) throws IOException {
+
+        }
+
+        @Override
+        public int validOps() {
+            return SelectionKey.OP_ACCEPT | SelectionKey.OP_CONNECT;
+        }
+    }
+
+    static class DummySelector extends AbstractSelector {
+        protected DummySelector(SelectorProvider provider) {
+            super(provider);
+        }
+
+        @Override
+        protected void implCloseSelector() throws IOException {
+
+        }
+
+        @Override
+        protected SelectionKey register(AbstractSelectableChannel ch, int ops, Object att) {
+            return null;
+        }
+
+        @Override
+        public Set<SelectionKey> keys() {
+            return Set.of();
+        }
+
+        @Override
+        public Set<SelectionKey> selectedKeys() {
+            return Set.of();
+        }
+
+        @Override
+        public int selectNow() throws IOException {
+            return 0;
+        }
+
+        @Override
+        public int select(long timeout) throws IOException {
+            return 0;
+        }
+
+        @Override
+        public int select() throws IOException {
+            return 0;
+        }
+
+        @Override
+        public Selector wakeup() {
+            return null;
+        }
+    }
 }

libs/entitlement/qa/entitlement-test-plugin/src/main/java/org/elasticsearch/entitlement/qa/test/NetworkAccessCheckActions.java

@@ -25,6 +25,7 @@
 import java.net.ServerSocket;
 import java.net.Socket;
 import java.net.SocketException;
+import java.net.StandardProtocolFamily;
 import java.net.URL;
 import java.net.URLConnection;
 import java.net.URLStreamHandler;
@@ -203,6 +204,20 @@ static void socketChannelConnect() throws IOException {
         }
     }
 
+    @EntitlementTest(expectedAccess = PLUGINS)
+    static void socketChannelOpenProtocol() throws IOException {
+        SocketChannel.open(StandardProtocolFamily.INET).close();
+    }
+
+    @EntitlementTest(expectedAccess = PLUGINS)
+    static void socketChannelOpenAddress() throws IOException {
+        try {
+            SocketChannel.open(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0)).close();
+        } catch (SocketException ex) {
+            // Some sort of SocketException is expected, we are trying to connect to port 0
+        }
+    }
+
     @EntitlementTest(expectedAccess = PLUGINS)
     static void asynchronousSocketChannelBind() throws IOException {
         try (var socketChannel = AsynchronousSocketChannel.open()) {

libs/entitlement/qa/entitlement-test-plugin/src/main/java/org/elasticsearch/entitlement/qa/test/NioChannelsActions.java

@@ -15,8 +15,11 @@
 
 import java.io.FileDescriptor;
 import java.io.IOException;
+import java.net.StandardProtocolFamily;
 import java.nio.channels.AsynchronousFileChannel;
 import java.nio.channels.FileChannel;
+import java.nio.channels.SelectionKey;
+import java.nio.channels.spi.SelectorProvider;
 import java.nio.file.StandardOpenOption;
 import java.util.Set;
 
@@ -85,4 +88,50 @@ static void asynchronousFileChannelOpenForReadWithOptions() throws IOException {
     static void channelsReadWriteSelectableChannel() throws IOException {
         jdk.nio.Channels.readWriteSelectableChannel(new FileDescriptor(), new DummyImplementations.DummySelectableChannelCloser()).close();
     }
+
+    @EntitlementTest(expectedAccess = PLUGINS)
+    static void selectableChannelRegisterConnect() throws IOException {
+        try (var selectableChannel = new DummyImplementations.DummySelectableChannel(SelectorProvider.provider())) {
+            selectableChannel.configureBlocking(false);
+            selectableChannel.register(new DummyImplementations.DummySelector(SelectorProvider.provider()), SelectionKey.OP_CONNECT);
+        }
+    }
+
+    @EntitlementTest(expectedAccess = PLUGINS)
+    static void selectableChannelRegisterAccept() throws IOException {
+        try (var selectableChannel = new DummyImplementations.DummySelectableChannel(SelectorProvider.provider())) {
+            selectableChannel.configureBlocking(false);
+            selectableChannel.register(new DummyImplementations.DummySelector(SelectorProvider.provider()), SelectionKey.OP_ACCEPT);
+        }
+    }
+
+    @EntitlementTest(expectedAccess = PLUGINS)
+    static void selectorProviderOpenSocketChannel() throws IOException {
+        SelectorProvider.provider().openSocketChannel().close();
+    }
+
+    @EntitlementTest(expectedAccess = PLUGINS)
+    static void selectorProviderOpenDatagramChannel() throws IOException {
+        SelectorProvider.provider().openDatagramChannel().close();
+    }
+
+    @EntitlementTest(expectedAccess = PLUGINS)
+    static void selectorProviderOpenServerSocketChannel() throws IOException {
+        SelectorProvider.provider().openServerSocketChannel().close();
+    }
+
+    @EntitlementTest(expectedAccess = PLUGINS)
+    static void selectorProviderOpenSocketChannelWithProtocol() throws IOException {
+        SelectorProvider.provider().openSocketChannel(StandardProtocolFamily.INET).close();
+    }
+
+    @EntitlementTest(expectedAccess = PLUGINS)
+    static void selectorProviderOpenDatagramChannelWithProtocol() throws IOException {
+        SelectorProvider.provider().openDatagramChannel(StandardProtocolFamily.INET).close();
+    }
+
+    @EntitlementTest(expectedAccess = PLUGINS)
+    static void selectorProviderOpenServerSocketChannelWithProtocol() throws IOException {
+        SelectorProvider.provider().openServerSocketChannel(StandardProtocolFamily.INET).close();
+    }
 }

libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/DynamicInstrumentation.java

@@ -122,20 +122,7 @@ static void initialize(Instrumentation inst, Class<?> checkerInterface, boolean
     private static Map<MethodKey, CheckMethod> getMethodsToInstrument(Class<?> checkerInterface) throws ClassNotFoundException,
         NoSuchMethodException {
         Map<MethodKey, CheckMethod> checkMethods = new HashMap<>(INSTRUMENTATION_SERVICE.lookupMethods(checkerInterface));
-        Stream.of(
-            fileSystemProviderChecks(),
-            fileStoreChecks(),
-            pathChecks(),
-            Stream.of(
-                INSTRUMENTATION_SERVICE.lookupImplementationMethod(
-                    SelectorProvider.class,
-                    "inheritedChannel",
-                    SelectorProvider.provider().getClass(),
-                    EntitlementChecker.class,
-                    "checkSelectorProviderInheritedChannel"
-                )
-            )
-        )
+        Stream.of(fileSystemProviderChecks(), fileStoreChecks(), pathChecks(), selectorProviderChecks())
             .flatMap(Function.identity())
             .forEach(instrumentation -> checkMethods.put(instrumentation.targetMethod(), instrumentation.checkMethod()));
 
@@ -257,6 +244,39 @@ private static Stream<InstrumentationService.InstrumentationInfo> pathChecks() {
         });
     }
 
+    private static Stream<InstrumentationService.InstrumentationInfo> selectorProviderChecks() {
+        var selectorProviderClass = SelectorProvider.provider().getClass();
+
+        var instrumentation = new InstrumentationInfoFactory() {
+            @Override
+            public InstrumentationService.InstrumentationInfo of(String methodName, Class<?>... parameterTypes)
+                throws ClassNotFoundException, NoSuchMethodException {
+                return INSTRUMENTATION_SERVICE.lookupImplementationMethod(
+                    SelectorProvider.class,
+                    methodName,
+                    selectorProviderClass,
+                    EntitlementChecker.class,
+                    "checkSelectorProvider" + Character.toUpperCase(methodName.charAt(0)) + methodName.substring(1),
+                    parameterTypes
+                );
+            }
+        };
+
+        try {
+            return Stream.of(
+                instrumentation.of("inheritedChannel"),
+                instrumentation.of("openDatagramChannel"),
+                instrumentation.of("openDatagramChannel", java.net.ProtocolFamily.class),
+                instrumentation.of("openServerSocketChannel"),
+                instrumentation.of("openServerSocketChannel", java.net.ProtocolFamily.class),
+                instrumentation.of("openSocketChannel"),
+                instrumentation.of("openSocketChannel", java.net.ProtocolFamily.class)
+            );
+        } catch (NoSuchMethodException | ClassNotFoundException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     private static Class<?>[] findClassesToRetransform(Class<?>[] loadedClasses, Set<String> classesToTransform) {
         List<Class<?>> retransform = new ArrayList<>();
         for (Class<?> loadedClass : loadedClasses) {

libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java

@@ -85,7 +85,11 @@ public static void initialize(Instrumentation inst) throws Exception {
      * transformed and undergo verification. In order to avoid circularity errors as much as possible, we force a partial order.
      */
     private static void ensureClassesSensitiveToVerificationAreInitialized() {
-        var classesToInitialize = Set.of("sun.net.www.protocol.http.HttpURLConnection");
+        var classesToInitialize = Set.of(
+            "sun.net.www.protocol.http.HttpURLConnection",
+            "sun.nio.ch.DatagramChannelImpl",
+            "sun.nio.ch.ServerSocketChannelImpl"
+        );
         for (String className : classesToInitialize) {
             try {
                 Class.forName(className);