From 886f4905d378c4225cbf13d53f073022fbb5acbc Mon Sep 17 00:00:00 2001 From: Matt Corallo Date: Wed, 3 Nov 2021 21:10:12 +0000 Subject: [PATCH] Fix race condition in NioPeerHandler on `socket_disconnected` --- src/main/java/org/ldk/batteries/NioPeerHandler.java | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/main/java/org/ldk/batteries/NioPeerHandler.java b/src/main/java/org/ldk/batteries/NioPeerHandler.java index bf111729..afe728b0 100644 --- a/src/main/java/org/ldk/batteries/NioPeerHandler.java +++ b/src/main/java/org/ldk/batteries/NioPeerHandler.java @@ -174,10 +174,11 @@ public class NioPeerHandler { if (chan == null) continue; try { Peer peer = setup_socket(chan); + peer.key = chan.register(this.selector, SelectionKey.OP_READ, peer); Result_NonePeerHandleErrorZ res = this.peer_manager.new_inbound_connection(peer.descriptor); - if (res instanceof Result_NonePeerHandleErrorZ.Result_NonePeerHandleErrorZ_OK) { - peer.key = chan.register(this.selector, SelectionKey.OP_READ, peer); - } + if (res instanceof Result_NonePeerHandleErrorZ.Result_NonePeerHandleErrorZ_Err) { + peer.descriptor.disconnect_socket(); + } } catch (IOException ignored) { } } continue; // There is no attachment so the rest of the loop is useless @@ -273,14 +274,17 @@ public class NioPeerHandler { throw new IOException("Timed out"); } Peer peer = setup_socket(chan); + do_selector_action(() -> peer.key = chan.register(this.selector, SelectionKey.OP_READ, peer)); Result_CVec_u8ZPeerHandleErrorZ res = this.peer_manager.new_outbound_connection(their_node_id, peer.descriptor); if (res instanceof Result_CVec_u8ZPeerHandleErrorZ.Result_CVec_u8ZPeerHandleErrorZ_OK) { byte[] initial_bytes = ((Result_CVec_u8ZPeerHandleErrorZ.Result_CVec_u8ZPeerHandleErrorZ_OK) res).res; if (chan.write(ByteBuffer.wrap(initial_bytes)) != initial_bytes.length) { + peer.descriptor.disconnect_socket(); + this.peer_manager.socket_disconnected(peer.descriptor); throw new IOException("We assume TCP socket buffer is at least a single packet in length"); } - do_selector_action(() -> peer.key = chan.register(this.selector, SelectionKey.OP_READ, peer)); } else { + peer.descriptor.disconnect_socket(); throw new IOException("LDK rejected outbound connection. This likely shouldn't ever happen."); } } -- 2.30.2