Send control msg only once and improve tests

libp2p/protocols/pubsub/pubsubpeer.nim

@@ -63,7 +63,7 @@ type
     heDontWants*: Deque[HashSet[MessageId]]
     iHaveBudget*: int
     pingBudget*: int
-    maxMessageSize: int
+    maxMessageSize*: int
     appScore*: float64 # application specific score
     behaviourPenalty*: float64 # the eventual penalty score
 
@@ -291,20 +291,27 @@ proc send*(p: PubSubPeer, msg: RPCMsg, anonymize: bool): seq[RPCMsg] {.raises: [
     # protobuf for every peer - this could easily be improved!
     sendMetrics(msg)
     encodeRpcMsg(msg, anonymize)
-  var res = newSeq[RPCMsg]()
+
+  var sentMessages: seq[RPCMsg]
   # Check if the encoded message size exceeds the maxMessageSize
   if encoded.len > p.maxMessageSize:
+    var controlSent = false
     # Split the RPCMsg into individual messages and send them separately
     for message in msg.messages:
-      var newMsg = RPCMsg(messages: @[message], control: msg.control)
+      var newMsg: RPCMsg
+      if not controlSent:
+        newMsg = RPCMsg(messages: @[message], control: msg.control)
+        controlSent = true
+      else:
+        newMsg = RPCMsg(messages: @[message])
       let newMsgEncoded = encodeRpcMsg(newMsg, anonymize)
       asyncSpawn p.sendEncoded(newMsgEncoded)
-      res.add(newMsg)
+      sentMessages.add(newMsg)
   else:
     # If the message size is within limits, send it as is
     asyncSpawn p.sendEncoded(encoded)
-    res.add(msg)
-  return res
+    sentMessages.add(msg)
+  return sentMessages
 
 proc canAskIWant*(p: PubSubPeer, msgId: MessageId): bool =
   for sentIHave in p.sentIHaves.mitems():

tests/pubsub/testpubsubpeer.nim

@@ -22,22 +22,47 @@ suite "GossipSub":
   test "Test oversized RPCMsg splitting":
     let
       gossipSub = TestGossipSub.init(newStandardSwitch())
-      maxMessageSize = 1048576  # 1MB for example, adjust as needed
       topic = "test_topic"
-      largeMessageSize = maxMessageSize div 2 + 100  # Just over half the max size
       peer = gossipSub.getPubSubPeer(randomPeerId())
+      largeMessageSize = peer.maxMessageSize div 2 + 100  # Just over half the max size
+
+    # Create a dummy ControlMessage
+    let
+      graft = @[ControlGraft(topicID: "topic1")]
+      prune = @[ControlPrune(topicID: "topic2")]
+      ihave = @[ControlIHave(topicID: "topic3", messageIDs: @[@[1'u8, 2, 3], @[4'u8, 5, 6]])]
+      iwant = @[ControlIWant(messageIDs: @[@[7'u8, 8, 9]])]
+
+    let control = some(ControlMessage(
+      graft: graft,
+      prune: prune,
+      ihave: ihave,
+      iwant: iwant
+    ))
 
     # Create a message that's just over half the max size
     let largeMessage = Message(data: newSeq[byte](largeMessageSize))
 
     # Create an RPCMsg with two such messages, making it oversized
-    var oversizedMsg = RPCMsg(messages: @[largeMessage, largeMessage])
+    var oversizedMsg = RPCMsg(messages: @[largeMessage, largeMessage], control: control)
 
     # Mock the sendEncoded function to capture the messages being sent
     var sentMessages = peer.send(oversizedMsg, false)
 
+    # The first split message should have the control data, others shouldn't
+    check: sentMessages[0].control == control
+    for i in 1..<sentMessages.len:
+      check: sentMessages[i].control.isNone
+
     # Check if the send function correctly split and sent the RPCMsg
     check:
       sentMessages.len == 2
       sentMessages[0].messages[0].data.len == largeMessageSize
       sentMessages[1].messages[0].data.len == largeMessageSize
+
+    # If the original message is below maxMessageSize, it's sent as-is
+    var smallMsg = RPCMsg(messages: @[Message(data: newSeqOfCap[byte](peer.maxMessageSize div 2))], control: control)
+
+    let sentSmallMessages = peer.send(smallMsg, false)
+    check: sentSmallMessages.len == 1
+    check: sentSmallMessages[0] == smallMsg