diff --git a/libp2p/protocols/pubsub/pubsubpeer.nim b/libp2p/protocols/pubsub/pubsubpeer.nim index 43fe8459a1..3a6ce79d49 100644 --- a/libp2p/protocols/pubsub/pubsubpeer.nim +++ b/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(): diff --git a/tests/pubsub/testpubsubpeer.nim b/tests/pubsub/testpubsubpeer.nim index 11a4a2c71b..a9930fbe2c 100644 --- a/tests/pubsub/testpubsubpeer.nim +++ b/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..