cancell enqueued rpcs on receiving IDontWant

Author: sukunrt

comm.go

@@ -165,7 +165,7 @@ func (p *PubSub) handlePeerDead(s network.Stream) {
 }
 
 func (p *PubSub) handleSendingMessages(ctx context.Context, s network.Stream, outgoing *rpcQueue) {
-	writeRpc := func(rpc *RPC) error {
+	writeRpc := func(rpc *pb.RPC) error {
 		size := uint64(rpc.Size())
 
 		buf := pool.Get(varint.UvarintSize(size) + int(size))
@@ -193,8 +193,11 @@ func (p *PubSub) handleSendingMessages(ctx context.Context, s network.Stream, ou
 			p.logger.Debug("error popping message from the queue to send to peer", "peer", s.Conn().RemotePeer(), "err", err)
 			return
 		}
+		if rpc.Size() == 0 {
+			continue
+		}
 
-		err = writeRpc(rpc)
+		err = writeRpc(&rpc.RPC)
 		if err != nil {
 			s.Reset()
 			p.logger.Debug("error writing message to peer", "peer", s.Conn().RemotePeer(), "err", err)
@@ -215,6 +218,10 @@ func rpcWithMessages(msgs ...*pb.Message) *RPC {
 	return &RPC{RPC: pb.RPC{Publish: msgs}}
 }
 
+func rpcWithMessageAndMsgID(msg *pb.Message, msgID string) *RPC {
+	return &RPC{RPC: pb.RPC{Publish: []*pb.Message{msg}}, MsgIDs: []string{msgID}}
+}
+
 func rpcWithControl(msgs []*pb.Message,
 	ihave []*pb.ControlIHave,
 	iwant []*pb.ControlIWant,

gossipsub.go

@@ -1163,6 +1163,9 @@ func (gs *GossipSubRouter) handleIDontWant(p peer.ID, ctl *pb.ControlMessage) {
 	gs.peerdontwant[p]++
 
 	totalUnwantedIds := 0
+	// Collect message IDs for cancellation
+	var msgIDsToCancel []string
+
 	// Remember all the unwanted message ids
 mainIDWLoop:
 	for _, idontwant := range ctl.GetIdontwant() {
@@ -1175,8 +1178,14 @@ mainIDWLoop:
 
 			totalUnwantedIds++
 			gs.unwanted[p][computeChecksum(mid)] = gs.params.IDontWantMessageTTL
+			msgIDsToCancel = append(msgIDsToCancel, mid)
 		}
 	}
+
+	// Cancel these messages in the RPC queue if it exists
+	if queue, ok := gs.p.peers[p]; ok && len(msgIDsToCancel) > 0 {
+		queue.CancelMessages(msgIDsToCancel)
+	}
 }
 
 func (gs *GossipSubRouter) addBackoff(p peer.ID, topic string, isUnsubscribe bool) {
@@ -1370,7 +1379,7 @@ func (gs *GossipSubRouter) rpcs(msg *Message) iter.Seq2[peer.ID, *RPC] {
 			}
 		}
 
-		out := rpcWithMessages(msg.Message)
+		out := rpcWithMessageAndMsgID(msg.Message, gs.p.idGen.ID(msg))
 		for pid := range tosend {
 			if pid == from || pid == peer.ID(msg.GetFrom()) {
 				continue

gossipsub_spam_test.go

@@ -930,20 +930,25 @@ func TestGossipsubHandleIDontwantSpam(t *testing.T) {
 	rPid := hosts[1].ID()
 	ctrlMessage := &pb.ControlMessage{Idontwant: []*pb.ControlIDontWant{{MessageIDs: idwIds}}}
 	grt := psubs[0].rt.(*GossipSubRouter)
-	grt.handleIDontWant(rPid, ctrlMessage)
+	completed := make(chan struct{})
+	psubs[0].eval <- func() {
+		grt.handleIDontWant(rPid, ctrlMessage)
 
-	if grt.peerdontwant[rPid] != 1 {
-		t.Errorf("Wanted message count of %d but received %d", 1, grt.peerdontwant[rPid])
-	}
-	mid := fmt.Sprintf("idontwant-%d", GossipSubMaxIDontWantLength-1)
-	if _, ok := grt.unwanted[rPid][computeChecksum(mid)]; !ok {
-		t.Errorf("Desired message id was not stored in the unwanted map: %s", mid)
-	}
+		if grt.peerdontwant[rPid] != 1 {
+			t.Errorf("Wanted message count of %d but received %d", 1, grt.peerdontwant[rPid])
+		}
+		mid := fmt.Sprintf("idontwant-%d", GossipSubMaxIDontWantLength-1)
+		if _, ok := grt.unwanted[rPid][computeChecksum(mid)]; !ok {
+			t.Errorf("Desired message id was not stored in the unwanted map: %s", mid)
+		}
 
-	mid = fmt.Sprintf("idontwant-%d", GossipSubMaxIDontWantLength)
-	if _, ok := grt.unwanted[rPid][computeChecksum(mid)]; ok {
-		t.Errorf("Unwanted message id was stored in the unwanted map: %s", mid)
+		mid = fmt.Sprintf("idontwant-%d", GossipSubMaxIDontWantLength)
+		if _, ok := grt.unwanted[rPid][computeChecksum(mid)]; ok {
+			t.Errorf("Unwanted message id was stored in the unwanted map: %s", mid)
+		}
+		close(completed)
 	}
+	<-completed
 }
 
 type mockGSOnRead func(writeMsg func(*pb.RPC), irpc *pb.RPC)