Fix concurrency issue from modifying the selectedKeysSet

kevink-sq · kohlschuetter · commit b8386a1d8391 · 2023-10-27T14:08:37.000+02:00
This bug may have been a regression introduced in optimizing the key removal: a7c3106 To reproduce refer to poc: https://github.com/kevink-sq/jetty-concurrency-issue-poc To verify fix, refer to poc branch (same as these changes): https://github.com/kevink-sq/jetty-concurrency-issue-poc/tree/kevink/attempted-junixsocket-fix-2 Jetty's [ManagedSelector](https://github.com/jetty/jetty.project/blob/7a7d69a69f4f51772e20813332291189a24e91b1/jetty-io/src/main/java/org/eclipse/jetty/io/ManagedSelector.java#L655) iterates through the selectedKeysSet and when separate processes cancel the AFSelectorKey, the iterator throws a `ConcurrentModificationException`. Fix is to introduce `cancelledKeysSet` to defer removal of cancelled keys until the select process. Use of the cancelled sets are referenced from java.nio.channels.spi's `AbstractSelector` and `AbstractSelectionKey`. Ideally we'd have `AFSelectionKey` extend from `AbstractSelectionKey` to leverage the base class's internal cancelled set but`AbstractSelectionKey` defines its own `isValid` and `cancel` and cannot override the base's final methods. Also added an atomic check to `AFSelectionKey#cancel` to prevent "Too many open files" socket exception.
diff --git a/junixsocket-common/src/main/java/org/newsclub/net/unix/AFSelectionKey.java b/junixsocket-common/src/main/java/org/newsclub/net/unix/AFSelectionKey.java
@@ -80,8 +80,10 @@ boolean isSelected() {
 
   @Override
   public void cancel() {
-    sel.remove(this);
-    cancelNoRemove();
+    if (!cancelled.compareAndSet(false, true) || !chann.isOpen()) {
+      return;
+    }
+    sel.prepareRemove(this);
   }
 
   void cancelNoRemove() {
diff --git a/junixsocket-common/src/main/java/org/newsclub/net/unix/AFSelector.java b/junixsocket-common/src/main/java/org/newsclub/net/unix/AFSelector.java
@@ -49,6 +49,7 @@ final class AFSelector extends AbstractSelector {
 
   private final Set<SelectionKey> selectedKeysSet = new HashSet<>();
   private final Set<SelectionKey> selectedKeysPublic = new UngrowableSet<>(selectedKeysSet);
+  private final Set<SelectionKey> cancelledKeysSet = new HashSet<>();
 
   private PollFd pollFd = null;
 
@@ -112,6 +113,7 @@ private int select0(int timeout) throws IOException {
         throw new ClosedSelectorException();
       }
       pfd = pollFd = initPollFd(pollFd);
+      performRemove();
       selectedKeysSet.clear();
     }
     int num;
@@ -299,6 +301,23 @@ synchronized void remove(AFSelectionKey key) {
     pollFd = null;
   }
 
+  void prepareRemove(AFSelectionKey key) {
+    synchronized (cancelledKeysSet) {
+      cancelledKeysSet.add(key);
+    }
+  }
+
+  void performRemove() {
+    synchronized (cancelledKeysSet) {
+      for (SelectionKey key : cancelledKeysSet) {
+        selectedKeysSet.remove(key);
+        deregister((AFSelectionKey) key);
+        pollFd = null;
+      }
+      cancelledKeysSet.clear();
+    }
+  }
+
   private void deregister(AFSelectionKey key) {
     // super.deregister unnecessarily casts SelectionKey to AbstractSelectionKey, and
     // ((AbstractSelectableChannel)key.channel()).removeKey(key); is not visible.
