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.

Author: TheBlueMatt

lightning-net-tokio/src/lib.rs

@@ -70,9 +70,9 @@ 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};
+use std::sync::{Arc, Mutex, MutexGuard};
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::Duration;
 use std::hash::Hash;
@@ -101,29 +101,31 @@ 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,
 }
 impl Connection {
+	fn event_trigger(us: &mut MutexGuard<Self>) {
+		match us.event_notify.try_send(()) {
+			Ok(_) => {},
+			Err(mpsc::error::TrySendError::Full(_)) => {
+				// Ignore full errors as we just need the user to poll after this point, so if they
+				// haven't received the last send yet, it doesn't matter.
+			},
+			_ => panic!()
+		}
+	}
 	async fn schedule_read<CMH: ChannelMessageHandler + 'static>(peer_manager: Arc<peer_handler::PeerManager<SocketDescriptor, Arc<CMH>>>, us: Arc<Mutex<Self>>, mut reader: io::ReadHalf<TcpStream>, mut read_wake_receiver: mpsc::Receiver<()>, mut write_avail_receiver: mpsc::Receiver<()>) {
 		let peer_manager_ref = peer_manager.clone();
 		// 8KB is nice and big but also should never cause any issues with stack overflowing.
 		let mut buf = [0; 8192];
 
-		macro_rules! event_occurred {
-			($lock: expr) => { {
-				match $lock.event_notify.try_send(()) {
-					Ok(_) => {},
-					Err(mpsc::error::TrySendError::Full(_)) => {
-						// Ignore full errors as we just need the user to poll after this point, so if they
-						// haven't received the last send yet, it doesn't matter.
-					},
-					_ => panic!()
-				}
-			} }
-		}
-
 		// If Rust-Lightning tells us to disconnect, we don't, then, call socket_disconnect() after
 		// closing the socket, but if the socket closes on its own (ie our counterparty does it), we do.
 		let need_disconnect_event;
@@ -136,48 +138,60 @@ 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", false);
+					}
+					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", false);
-					}
+					prepare_read_write_call!();
 					if let Err(e) = peer_manager.write_buffer_space_avail(&mut SocketDescriptor::new(us.clone())) {
 						shutdown_socket!(e, false);
 					}
+					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", true),
 					Ok(len) => {
-						if us.lock().unwrap().rl_requested_disconnect {
-							shutdown_socket!("disconnect_socket() call from RL", false);
-						}
+						prepare_read_write_call!();
 						let read_res = peer_manager.read_event(&mut SocketDescriptor::new(Arc::clone(&us)), &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;
 								}
-								event_occurred!(us_lock);
+								Self::event_trigger(&mut us_lock);
 							},
 							Err(e) => shutdown_socket!(e, false),
 						}
+						us_lock.block_disconnect_socket = false;
 					},
 					Err(e) => shutdown_socket!(e, true),
 				},
 			}
 		}
-		let writer_option = us.lock().unwrap().writer.take();
+		let writer_option = {
+			let mut us_lock = us.lock().unwrap();
+			us_lock.block_disconnect_socket = false;
+			us_lock.writer.take()
+		};
 		if let Some(mut writer) = writer_option {
 			// If the socket is already closed, shutdown() will fail, so just ignore it.
 			let _ = writer.shutdown().await;
 		}
 		if need_disconnect_event {
 			peer_manager_ref.socket_disconnected(&SocketDescriptor::new(Arc::clone(&us)));
-			event_occurred!(us.lock().unwrap());
+			Self::event_trigger(&mut us.lock().unwrap());
 		}
 	}
 
@@ -196,8 +210,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)
 		})))
 	}
@@ -384,15 +398,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 {