Add support for batches of other types

While the spec only includes commitment_signed messages in batches,
there may be other types of batches in the future. Generalize the
message batching code to allow for other types in the future.

Author: jkczyz

lightning/src/ln/msgs.rs

@@ -696,6 +696,8 @@ pub struct StartBatch {
 	pub channel_id: ChannelId,
 	/// The number of messages to follow.
 	pub batch_size: u16,
+	/// The type of all messages expected in the batch.
+	pub message_type: Option<u16>,
 }
 
 /// An [`update_add_htlc`] message to be sent to or received from a peer.
@@ -3099,7 +3101,9 @@ impl_writeable_msg!(PeerStorageRetrieval, { data }, {});
 impl_writeable_msg!(StartBatch, {
 	channel_id,
 	batch_size
-}, {});
+}, {
+	(1, message_type, option)
+});
 
 // Note that this is written as a part of ChannelManager objects, and thus cannot change its
 // serialization format in a way which assumes we know the total serialized length/message end

lightning/src/ln/peer_handler.rs

@@ -539,6 +539,24 @@ enum InitSyncTracker{
 	NodesSyncing(NodeId),
 }
 
+/// A batch of messages initiated when receiving a `start_batch` message.
+struct MessageBatch {
+	/// The channel associated with all the messages in the batch.
+	channel_id: ChannelId,
+
+	/// The number of messages expected to be in the batch.
+	batch_size: usize,
+
+	/// The batch of messages, which should all be of the same type.
+	messages: MessageBatchImpl,
+}
+
+/// The representation of the message batch, which may different for each message type.
+enum MessageBatchImpl {
+	/// A batch of `commitment_signed` messages, where each has a unique `funding_txid`.
+	CommitmentSigned(BTreeMap<Txid, msgs::CommitmentSigned>),
+}
+
 /// The ratio between buffer sizes at which we stop sending initial sync messages vs when we stop
 /// forwarding gossip messages to peers altogether.
 const FORWARD_INIT_SYNC_BUFFER_LIMIT_RATIO: usize = 2;
@@ -620,7 +638,7 @@ struct Peer {
 
 	inbound_connection: bool,
 
-	commitment_signed_batch: Option<(ChannelId, usize, BTreeMap<Txid, msgs::CommitmentSigned>)>,
+	message_batch: Option<MessageBatch>,
 }
 
 impl Peer {
@@ -1159,7 +1177,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 					received_channel_announce_since_backlogged: false,
 					inbound_connection: false,
 
-					commitment_signed_batch: None,
+					message_batch: None,
 				}));
 				Ok(res)
 			}
@@ -1217,7 +1235,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 					received_channel_announce_since_backlogged: false,
 					inbound_connection: true,
 
-					commitment_signed_batch: None,
+					message_batch: None,
 				}));
 				Ok(())
 			}
@@ -1773,7 +1791,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 		// During splicing, commitment_signed messages need to be collected into a single batch
 		// before they are handled.
 		if let wire::Message::StartBatch(msg) = message {
-			if peer_lock.commitment_signed_batch.is_some() {
+			if peer_lock.message_batch.is_some() {
 				let error = format!("Peer {} sent start_batch for channel {} before previous batch completed", log_pubkey!(their_node_id), &msg.channel_id);
 				log_debug!(logger, "{}", error);
 				return Err(LightningError {
@@ -1818,15 +1836,41 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 				}.into());
 			}
 
-			peer_lock.commitment_signed_batch = Some((msg.channel_id, batch_size, BTreeMap::new()));
+			let messages = match msg.message_type {
+				Some(message_type) if message_type == msgs::CommitmentSigned::TYPE => {
+					MessageBatchImpl::CommitmentSigned(BTreeMap::new())
+				},
+				_ => {
+					let error = format!("Peer {} sent start_batch for channel {} without a known message type", log_pubkey!(their_node_id), &msg.channel_id);
+					log_debug!(logger, "{}", error);
+					return Err(LightningError {
+						err: error.clone(),
+						action: msgs::ErrorAction::DisconnectPeerWithWarning {
+							msg: msgs::WarningMessage {
+								channel_id: msg.channel_id,
+								data: error,
+							},
+						},
+					}.into());
+				},
+			};
+
+			let message_batch = MessageBatch {
+				channel_id: msg.channel_id,
+				batch_size,
+				messages,
+			};
+			peer_lock.message_batch = Some(message_batch);
 
 			return Ok(None);
 		}
 
 		if let wire::Message::CommitmentSigned(msg) = message {
-			if let Some((channel_id, batch_size, buffer)) = &mut peer_lock.commitment_signed_batch {
-				if msg.channel_id != *channel_id {
-					let error = format!("Peer {} sent batched commitment_signed for the wrong channel (expected: {}, actual: {})", log_pubkey!(their_node_id), channel_id, &msg.channel_id);
+			if let Some(message_batch) = &mut peer_lock.message_batch {
+				let MessageBatchImpl::CommitmentSigned(ref mut buffer) = &mut message_batch.messages;
+
+				if msg.channel_id != message_batch.channel_id {
+					let error = format!("Peer {} sent batched commitment_signed for the wrong channel (expected: {}, actual: {})", log_pubkey!(their_node_id), message_batch.channel_id, &msg.channel_id);
 					log_debug!(logger, "{}", error);
 					return Err(LightningError {
 						err: error.clone(),
@@ -1842,30 +1886,37 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 				let funding_txid = match msg.funding_txid {
 					Some(funding_txid) => funding_txid,
 					None => {
-						log_debug!(logger, "Peer {} sent batched commitment_signed without a funding_txid for channel {}", log_pubkey!(their_node_id), channel_id);
+						log_debug!(logger, "Peer {} sent batched commitment_signed without a funding_txid for channel {}", log_pubkey!(their_node_id), message_batch.channel_id);
 						return Err(PeerHandleError { }.into());
 					},
 				};
 
 				match buffer.entry(funding_txid) {
 					btree_map::Entry::Vacant(entry) => { entry.insert(msg); },
 					btree_map::Entry::Occupied(_) => {
-						log_debug!(logger, "Peer {} sent batched commitment_signed with duplicate funding_txid {} for channel {}", log_pubkey!(their_node_id), funding_txid, channel_id);
+						log_debug!(logger, "Peer {} sent batched commitment_signed with duplicate funding_txid {} for channel {}", log_pubkey!(their_node_id), funding_txid, message_batch.channel_id);
 						return Err(PeerHandleError { }.into());
 					}
 				}
 
-				if buffer.len() == *batch_size {
-					let (channel_id, _, batch) = peer_lock.commitment_signed_batch.take().expect("batch should have been inserted");
+				if buffer.len() == message_batch.batch_size {
+					let MessageBatch { channel_id, batch_size: _, messages } = peer_lock.message_batch.take().expect("batch should have been inserted");
+					let MessageBatchImpl::CommitmentSigned(batch) = messages;
+
 					return Ok(Some(LogicalMessage::CommitmentSignedBatch(channel_id, batch)));
 				} else {
 					return Ok(None);
 				}
 			} else {
 				return Ok(Some(LogicalMessage::FromWire(wire::Message::CommitmentSigned(msg))));
 			}
-		} else if peer_lock.commitment_signed_batch.is_some() {
-			log_debug!(logger, "Peer {} sent non-commitment_signed message when expecting batched commitment_signed", log_pubkey!(their_node_id));
+		} else if let Some(message_batch) = &peer_lock.message_batch {
+			match message_batch.messages {
+				MessageBatchImpl::CommitmentSigned(_) => {
+					log_debug!(logger, "Peer {} sent an unexpected message for a commitment_signed batch", log_pubkey!(their_node_id));
+				},
+			}
+
 			return Err(PeerHandleError { }.into());
 		}
 
@@ -2465,6 +2516,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 								let msg = msgs::StartBatch {
 									channel_id: *channel_id,
 									batch_size: commitment_signed.len() as u16,
+									message_type: Some(msgs::CommitmentSigned::TYPE),
 								};
 								self.enqueue_message(&mut *peer, &msg);
 							}