From d6234aa328624d2ca07dd27b071b1eff12611d47 Mon Sep 17 00:00:00 2001
From: petricadaipegsp <155911522+petricadaipegsp@users.noreply.github.com>
Date: Tue, 19 Nov 2024 11:42:29 +0100
Subject: [PATCH] Avoid BlossomSubRouter race condition (#364)

---
 node/p2p/blossomsub.go         | 15 +++-----------
 node/p2p/internal/discovery.go | 37 ++++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 12 deletions(-)
 create mode 100644 node/p2p/internal/discovery.go

diff --git a/node/p2p/blossomsub.go b/node/p2p/blossomsub.go
index aa637fd..b41f6c3 100644
--- a/node/p2p/blossomsub.go
+++ b/node/p2p/blossomsub.go
@@ -448,6 +448,9 @@ func NewBlossomSub(
 		internal.PeerAddrInfosToPeerIDSlice(bootstrappers),
 		true,
 	)))
+	blossomOpts = append(blossomOpts, blossomsub.WithDiscovery(
+		internal.NewPeerConnectorDiscovery(discovery),
+	))
 
 	params := toBlossomSubParams(p2pConfig)
 	rt := blossomsub.NewBlossomSubRouter(h, params, bs.network)
@@ -463,18 +466,6 @@ func NewBlossomSub(
 	bs.h = h
 	bs.signKey = privKey
 
-	go func() {
-		for {
-			time.Sleep(30 * time.Second)
-			for _, mask := range pubsub.GetBitmasks() {
-				if !rt.EnoughPeers([]byte(mask), 0) {
-					_ = discovery.Connect(ctx)
-					break
-				}
-			}
-		}
-	}()
-
 	return bs
 }
 
diff --git a/node/p2p/internal/discovery.go b/node/p2p/internal/discovery.go
new file mode 100644
index 0000000..1d89e8a
--- /dev/null
+++ b/node/p2p/internal/discovery.go
@@ -0,0 +1,37 @@
+package internal
+
+import (
+	"context"
+	"time"
+
+	"github.com/libp2p/go-libp2p/core/discovery"
+	"github.com/libp2p/go-libp2p/core/peer"
+)
+
+type peerConnectorDiscovery struct {
+	connector PeerConnector
+}
+
+var _ discovery.Discovery = (*peerConnectorDiscovery)(nil)
+
+// Advertise implements discovery.Advertiser.
+func (d *peerConnectorDiscovery) Advertise(ctx context.Context, ns string, opts ...discovery.Option) (time.Duration, error) {
+	return time.Duration(1<<63 - 1), nil
+}
+
+// FindPeers implements discovery.Discoverer.
+func (d *peerConnectorDiscovery) FindPeers(ctx context.Context, ns string, opts ...discovery.Option) (<-chan peer.AddrInfo, error) {
+	if err := d.connector.Connect(ctx); err != nil {
+		return nil, err
+	}
+	ch := make(chan peer.AddrInfo)
+	close(ch)
+	return ch, nil
+}
+
+// NewPeerConnectorDiscovery creates a new peer connector discovery.
+// The discovery instance does not do any advertisements and just triggers
+// the peer connector once FindPeers is called.
+func NewPeerConnectorDiscovery(connector PeerConnector) discovery.Discovery {
+	return &peerConnectorDiscovery{connector: connector}
+}