From ebef1f83fe4f70ba2f5c683a69f08e7ba32c8ad3 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Thu, 27 Jan 2022 18:39:52 +0000
Subject: [PATCH] [Java] More carefully ensure sockets are closed in
 `NioPeerHandler`

There is no indication that the sockets are leaking in the handling
loop, but its good to be careful and ensure we always both close
the SocketChannel and cancel the key registration.

During connection, however, it appears we may leak a SocketChannel
if the connection times out, leaking a file descriptor at least
until the GC kicks in and cleans up after us.  Here we are more
careful.
---
 .../org/ldk/batteries/NioPeerHandler.java     | 31 ++++++++++++-------
 1 file changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/main/java/org/ldk/batteries/NioPeerHandler.java b/src/main/java/org/ldk/batteries/NioPeerHandler.java
index e0dd83e9..2ab48d62 100644
--- a/src/main/java/org/ldk/batteries/NioPeerHandler.java
+++ b/src/main/java/org/ldk/batteries/NioPeerHandler.java
@@ -87,7 +87,7 @@ public class NioPeerHandler {
 							(peer.key.interestOps() | SelectionKey.OP_READ) & (~SelectionKey.OP_WRITE)));
 					}
                     return written;
-                } catch (IOException e) {
+                } catch (IOException|CancelledKeyException ignored) {
                     // Most likely the socket is disconnected, let the background thread handle it.
                     return 0;
                 }
@@ -97,7 +97,7 @@ public class NioPeerHandler {
             public void disconnect_socket() {
                 try {
                     do_selector_action(() -> {
-                        peer.key.cancel();
+                        try { peer.key.cancel(); } catch (CancelledKeyException ignored) {}
                         peer.key.channel().close();
                     });
                 } catch (IOException ignored) { }
@@ -189,8 +189,8 @@ public class NioPeerHandler {
                             if (key.isValid() && (key.interestOps() & SelectionKey.OP_WRITE) != 0 && key.isWritable()) {
                                 Result_NonePeerHandleErrorZ res = this.peer_manager.write_buffer_space_avail(peer.descriptor);
                                 if (res instanceof Result_NonePeerHandleErrorZ.Result_NonePeerHandleErrorZ_Err) {
-                                    key.channel().close();
                                     key.cancel();
+                                    key.channel().close();
                                 }
                             }
                             if (key.isValid() && (key.interestOps() & SelectionKey.OP_READ) != 0 && key.isReadable()) {
@@ -199,6 +199,7 @@ public class NioPeerHandler {
                                 if (read == -1) {
                                     this.peer_manager.socket_disconnected(peer.descriptor);
                                     key.cancel();
+                                    key.channel().close(); // This may throw, we read -1 so the channel should already be closed, but do this to be safe
                                 } else if (read > 0) {
                                     ((Buffer)buf).flip();
                                     // This code is quite hot during initial network graph sync, so we go a ways out of
@@ -223,15 +224,15 @@ public class NioPeerHandler {
                                             key.interestOps(key.interestOps() & (~SelectionKey.OP_READ));
                                         }
                                     } else {
-                                        key.channel().close();
                                         key.cancel();
+                                        key.channel().close();
                                     }
                                     bindings.CResult_boolPeerHandleErrorZ_free(read_result_pointer);
                                 }
                             }
                         } catch (IOException ignored) {
-                            try { key.channel().close(); } catch (IOException ignored2) { }
                             key.cancel();
+                            try { key.channel().close(); } catch (IOException ignored2) { }
                             peer_manager.socket_disconnected(peer.descriptor);
                         }
                     } catch (CancelledKeyException e) {
@@ -265,13 +266,21 @@ public class NioPeerHandler {
      */
     public void connect(byte[] their_node_id, SocketAddress remote, int timeout_ms) throws IOException {
         SocketChannel chan = SocketChannel.open();
-        chan.configureBlocking(false);
-        Selector open_selector = Selector.open();
-        chan.register(open_selector, SelectionKey.OP_CONNECT);
-        if (!chan.connect(remote)) {
-            open_selector.select(timeout_ms);
+        boolean connected;
+        try {
+            chan.configureBlocking(false);
+            Selector open_selector = Selector.open();
+            chan.register(open_selector, SelectionKey.OP_CONNECT);
+            if (!chan.connect(remote)) {
+                open_selector.select(timeout_ms);
+            }
+            connected = chan.finishConnect();
+        } catch (IOException e) {
+            try { chan.close(); } catch (IOException _e) { }
+            throw e;
         }
-        if (!chan.finishConnect()) { // Note that this may throw its own IOException if we failed for another reason
+        if (!connected) {
+            try { chan.close(); } catch (IOException _e) { }
             throw new IOException("Timed out");
         }
         Peer peer = setup_socket(chan);
-- 
2.39.5