From 5ae183b49bb8e3724f8bbdb51859f67769cafc90 Mon Sep 17 00:00:00 2001
From: Rusty Russell <rusty@rustcorp.com.au>
Date: Sat, 25 Jan 2025 11:10:11 +1030
Subject: [PATCH] connectd: attach input filtering for incoming dev_disconnect.

Signed-off-by: Rusty Russell <rusty@rustcorp.com.au>
---
 connectd/multiplex.c             | 28 ++++++++++++++++++++++------
 connectd/peer_exchange_initmsg.c |  7 +++++++
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/connectd/multiplex.c b/connectd/multiplex.c
index e49bb768d..227dbe6a9 100644
--- a/connectd/multiplex.c
+++ b/connectd/multiplex.c
@@ -1136,6 +1136,13 @@ static struct subd *new_subd(struct peer *peer,
 	return subd;
 }
 
+static struct io_plan *close_peer_dev_disconnect(struct io_conn *peer_conn,
+						 struct peer *peer)
+{
+	assert(peer->to_peer == peer_conn);
+	return io_close_cb(peer_conn, peer);
+}
+
 static struct io_plan *read_hdr_from_peer(struct io_conn *peer_conn,
 					  struct peer *peer);
 static struct io_plan *read_body_from_peer_done(struct io_conn *peer_conn,
@@ -1145,7 +1152,8 @@ static struct io_plan *read_body_from_peer_done(struct io_conn *peer_conn,
        struct channel_id channel_id;
        struct subd *subd;
        enum peer_wire type;
-
+       struct io_plan *(*next_read)(struct io_conn *peer_conn,
+				    struct peer *peer) = read_hdr_from_peer;
 
        decrypted = cryptomsg_decrypt_body(tmpctx, &peer->cs,
 					  peer->peer_in);
@@ -1162,16 +1170,24 @@ static struct io_plan *read_body_from_peer_done(struct io_conn *peer_conn,
        if (!peer->dev_read_enabled)
 	       return read_hdr_from_peer(peer_conn, peer);
 
+       switch (dev_disconnect_in(&peer->id, type)) {
+       case DEV_DISCONNECT_IN_NORMAL:
+	       break;
+       case DEV_DISCONNECT_IN_AFTER_RECV:
+	       next_read = close_peer_dev_disconnect;
+	       break;
+       }
+
        /* We got something! */
        peer->last_recv_time = time_now();
 
        /* Don't process packets while we're closing */
        if (peer->draining)
-	       return read_hdr_from_peer(peer_conn, peer);
+	       return next_read(peer_conn, peer);
 
        /* If we swallow this, just try again. */
        if (handle_message_locally(peer, decrypted))
-	       return read_hdr_from_peer(peer_conn, peer);
+	       return next_read(peer_conn, peer);
 
        /* After this we should be able to match to subd by channel_id */
        if (!extract_channel_id(decrypted, &channel_id)) {
@@ -1186,7 +1202,7 @@ static struct io_plan *read_body_from_peer_done(struct io_conn *peer_conn,
 					"Received %s: %s",
 					peer_wire_name(type), desc);
 		       if (type == WIRE_WARNING)
-			       return read_hdr_from_peer(peer_conn, peer);
+			       return next_read(peer_conn, peer);
 		       return io_close(peer_conn);
 	       }
 
@@ -1194,7 +1210,7 @@ static struct io_plan *read_body_from_peer_done(struct io_conn *peer_conn,
 	       send_warning(peer, "Unexpected message %s: %s",
 			    peer_wire_name(type),
 			    tal_hex(tmpctx, decrypted));
-	       return read_hdr_from_peer(peer_conn, peer);
+	       return next_read(peer_conn, peer);
        }
 
        /* If we don't find a subdaemon for this, create a new one. */
@@ -1237,7 +1253,7 @@ static struct io_plan *read_body_from_peer_done(struct io_conn *peer_conn,
        }
 
        /* Wait for them to wake us */
-       return io_wait(peer_conn, &peer->peer_in, read_hdr_from_peer, peer);
+       return io_wait(peer_conn, &peer->peer_in, next_read, peer);
 }
 
 static struct io_plan *read_body_from_peer(struct io_conn *peer_conn,
diff --git a/connectd/peer_exchange_initmsg.c b/connectd/peer_exchange_initmsg.c
index a7df0d6eb..b0a29549c 100644
--- a/connectd/peer_exchange_initmsg.c
+++ b/connectd/peer_exchange_initmsg.c
@@ -61,6 +61,13 @@ static struct io_plan *peer_init_received(struct io_conn *conn,
 
 	status_peer_io(LOG_IO_IN, &peer->id, msg);
 
+	switch (dev_disconnect_in(&peer->id, fromwire_peektype(msg))) {
+	case DEV_DISCONNECT_IN_NORMAL:
+		break;
+	case DEV_DISCONNECT_IN_AFTER_RECV:
+		return io_close(conn);
+	}
+
 	/* BOLT #1:
 	 *
 	 * A receiving node: