Fix long-standing race in net-tokio reading after a disconnect event

If rust-lightning tells us to disconnect a socket after we read
some bytes from the socket, but before we actually give those bytes
to rust-lightning, we may end up calling rust-lightning with a
Descriptor that isn't registered anymore.

Sadly, there really isn't a good way to solve this, and it should
be a pretty quick event, so we just busy-wait.

diff --git a/lightning-net-tokio/src/lib.rs b/lightning-net-tokio/src/lib.rs
index 08aa12571985a985f896929e2034a248ff77ca1a..e635d93d383162560ad28060de65dd5a361dd902 100644 (file)
--- a/lightning-net-tokio/src/lib.rs
+++ b/lightning-net-tokio/src/lib.rs
@@ -70,7 +70,7 @@ use lightning::ln::peer_handler;
 use lightning::ln::peer_handler::SocketDescriptor as LnSocketTrait;
 use lightning::ln::msgs::ChannelMessageHandler;
 
-use std::task;
+use std::{task, thread};
 use std::net::SocketAddr;
 use std::sync::{Arc, Mutex, MutexGuard};
 use std::sync::atomic::{AtomicU64, Ordering};
@@ -101,6 +101,11 @@ struct Connection {
        // socket. To wake it up (without otherwise changing its state, we can push a value into this
        // Sender.
        read_waker: mpsc::Sender<()>,
+       // When we are told by rust-lightning to disconnect, we can't return to rust-lightning until we
+       // are sure we won't call any more read/write PeerManager functions with the same connection.
+       // This is set to true if we're in such a condition (with disconnect checked before with the
+       // top-level mutex held) and false when we can return.
+       block_disconnect_socket: bool,
        read_paused: bool,
        rl_requested_disconnect: bool,
        id: u64,
@@ -143,28 +148,35 @@ impl Connection {
                                } }
                        }
 
+                       macro_rules! prepare_read_write_call {
+                               () => { {
+                                       let mut us_lock = us.lock().unwrap();
+                                       if us_lock.rl_requested_disconnect {
+                                               shutdown_socket!("disconnect_socket() call from RL", Disconnect::CloseConnection);
+                                       }
+                                       us_lock.block_disconnect_socket = true;
+                               } }
+                       }
+
                        let read_paused = us.lock().unwrap().read_paused;
                        tokio::select! {
                                v = write_avail_receiver.recv() => {
                                        assert!(v.is_some()); // We can't have dropped the sending end, its in the us Arc!
-                                       if us.lock().unwrap().rl_requested_disconnect {
-                                               shutdown_socket!("disconnect_socket() call from RL", Disconnect::CloseConnection);
-                                       }
+                                       prepare_read_write_call!();
                                        if let Err(e) = peer_manager.write_buffer_space_avail(&mut our_descriptor) {
                                                shutdown_socket!(e, Disconnect::CloseConnection);
                                        }
+                                       us.lock().unwrap().block_disconnect_socket = false;
                                },
                                _ = read_wake_receiver.recv() => {},
                                read = reader.read(&mut buf), if !read_paused => match read {
                                        Ok(0) => shutdown_socket!("Connection closed", Disconnect::PeerDisconnected),
                                        Ok(len) => {
-                                               if us.lock().unwrap().rl_requested_disconnect {
-                                                       shutdown_socket!("disconnect_socket() call from RL", Disconnect::CloseConnection);
-                                               }
+                                               prepare_read_write_call!();
                                                let read_res = peer_manager.read_event(&mut our_descriptor, &buf[0..len]);
+                                               let mut us_lock = us.lock().unwrap();
                                                match read_res {
                                                        Ok(pause_read) => {
-                                                               let mut us_lock = us.lock().unwrap();
                                                                if pause_read {
                                                                        us_lock.read_paused = true;
                                                                }
@@ -172,6 +184,7 @@ impl Connection {
                                                        },
                                                        Err(e) => shutdown_socket!(e, Disconnect::CloseConnection),
                                                }
+                                               us_lock.block_disconnect_socket = false;
                                        },
                                        Err(e) => shutdown_socket!(e, Disconnect::PeerDisconnected),
                                },
@@ -203,8 +216,8 @@ impl Connection {
 
                (reader, write_receiver, read_receiver,
                Arc::new(Mutex::new(Self {
-                       writer: Some(writer), event_notify, write_avail, read_waker,
-                       read_paused: false, rl_requested_disconnect: false,
+                       writer: Some(writer), event_notify, write_avail, read_waker, read_paused: false,
+                       block_disconnect_socket: false, rl_requested_disconnect: false,
                        id: ID_COUNTER.fetch_add(1, Ordering::AcqRel)
                })))
        }
@@ -411,15 +424,18 @@ impl peer_handler::SocketDescriptor for SocketDescriptor {
        }
 
        fn disconnect_socket(&mut self) {
-               let mut us = self.conn.lock().unwrap();
-               us.rl_requested_disconnect = true;
-               us.read_paused = true;
-               // Wake up the sending thread, assuming it is still alive
-               let _ = us.write_avail.try_send(());
-               // TODO: There's a race where we don't meet the requirements of disconnect_socket if the
-               // read task is about to call a PeerManager function (eg read_event or write_event).
-               // Ideally we need to release the us lock and block until we have confirmation from the
-               // read task that it has broken out of its main loop.
+               {
+                       let mut us = self.conn.lock().unwrap();
+                       us.rl_requested_disconnect = true;
+                       us.read_paused = true;
+                       // Wake up the sending thread, assuming it is still alive
+                       let _ = us.write_avail.try_send(());
+                       // Happy-path return:
+                       if !us.block_disconnect_socket { return; }
+               }
+               while self.conn.lock().unwrap().block_disconnect_socket {
+                       thread::yield_now();
+               }
        }
 }
 impl Clone for SocketDescriptor {