Skip to content

Commit d44d15a

Browse files
committed
Keep the hashes of IDONTWANT message ids instead
Rather than keeping the raw message ids, keep their hashes instead to save memory and protect again memory DoS attacks.
1 parent afceb70 commit d44d15a

File tree

1 file changed

+10
-5
lines changed

1 file changed

+10
-5
lines changed

gossipsub.go

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package pubsub
22

33
import (
44
"context"
5+
"crypto/sha256"
56
"fmt"
67
"math/rand"
78
"sort"
@@ -64,6 +65,8 @@ var (
6465
GossipSubIDontWantMessageTTL = 3 // 3 heartbeats
6566
)
6667

68+
type checksum [32]byte
69+
6770
// GossipSubParams defines all the gossipsub specific parameters.
6871
type GossipSubParams struct {
6972
// overlay parameters.
@@ -243,7 +246,7 @@ func DefaultGossipSubRouter(h host.Host) *GossipSubRouter {
243246
backoff: make(map[string]map[peer.ID]time.Time),
244247
peerhave: make(map[peer.ID]int),
245248
peerdontwant: make(map[peer.ID]int),
246-
unwanted: make(map[peer.ID]map[string]int),
249+
unwanted: make(map[peer.ID]map[checksum]int),
247250
iasked: make(map[peer.ID]int),
248251
outbound: make(map[peer.ID]bool),
249252
connect: make(chan connectInfo, params.MaxPendingConnections),
@@ -448,7 +451,7 @@ type GossipSubRouter struct {
448451
control map[peer.ID]*pb.ControlMessage // pending control messages
449452
peerhave map[peer.ID]int // number of IHAVEs received from peer in the last heartbeat
450453
peerdontwant map[peer.ID]int // number of IDONTWANTs received from peer in the last heartbeat
451-
unwanted map[peer.ID]map[string]int // TTL of the message ids peers don't want
454+
unwanted map[peer.ID]map[checksum]int // TTL of the message ids peers don't want
452455
iasked map[peer.ID]int // number of messages we have asked from peer in the last heartbeat
453456
outbound map[peer.ID]bool // connection direction cache, marks peers with outbound connections
454457
backoff map[string]map[peer.ID]time.Time // prune backoff
@@ -966,7 +969,7 @@ func (gs *GossipSubRouter) handlePrune(p peer.ID, ctl *pb.ControlMessage) {
966969

967970
func (gs *GossipSubRouter) handleIDontWant(p peer.ID, ctl *pb.ControlMessage) {
968971
if gs.unwanted[p] == nil {
969-
gs.unwanted[p] = make(map[string]int)
972+
gs.unwanted[p] = make(map[checksum]int)
970973
}
971974

972975
// IDONTWANT flood protection
@@ -979,7 +982,8 @@ func (gs *GossipSubRouter) handleIDontWant(p peer.ID, ctl *pb.ControlMessage) {
979982
// Remember all the unwanted message ids
980983
for _, idontwant := range ctl.GetIdontwant() {
981984
for _, mid := range idontwant.GetMessageIDs() {
982-
gs.unwanted[p][mid] = gs.params.IDontWantMessageTTL
985+
hashed := sha256.Sum256([]byte(mid))
986+
gs.unwanted[p][hashed] = gs.params.IDontWantMessageTTL
983987
}
984988
}
985989
}
@@ -1145,9 +1149,10 @@ func (gs *GossipSubRouter) Publish(msg *Message) {
11451149

11461150
for p := range gmap {
11471151
mid := gs.p.idGen.ID(msg)
1152+
hashed := sha256.Sum256([]byte(mid))
11481153
// Check if it has already received an IDONTWANT for the message.
11491154
// If so, don't send it to the peer
1150-
if _, ok := gs.unwanted[p][mid]; ok {
1155+
if _, ok := gs.unwanted[p][hashed]; ok {
11511156
continue
11521157
}
11531158
tosend[p] = struct{}{}

0 commit comments

Comments
 (0)