From 9674127e9970fd818382bcc20cc24cfd1be28a91 Mon Sep 17 00:00:00 2001
From: ucwong <ucwong@126.com>
Date: Fri, 16 Aug 2024 22:20:13 +0800
Subject: [PATCH] dial nodes from discv5 dht

---
 ctxc/backend.go   | 29 ++++++++++++++++++++++++++---
 ctxc/discovery.go | 18 +++++++++++-------
 nightly.sh        |  2 +-
 3 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/ctxc/backend.go b/ctxc/backend.go
index 3ecadc0101..efb32bcf0a 100644
--- a/ctxc/backend.go
+++ b/ctxc/backend.go
@@ -48,6 +48,7 @@ import (
 	"github.com/CortexFoundation/CortexTheseus/miner"
 	"github.com/CortexFoundation/CortexTheseus/node"
 	"github.com/CortexFoundation/CortexTheseus/p2p"
+	"github.com/CortexFoundation/CortexTheseus/p2p/dnsdisc"
 	"github.com/CortexFoundation/CortexTheseus/p2p/enode"
 	"github.com/CortexFoundation/CortexTheseus/p2p/enr"
 	"github.com/CortexFoundation/CortexTheseus/params"
@@ -70,7 +71,7 @@ type Cortex struct {
 	blockchain      *core.BlockChain
 	protocolManager *ProtocolManager
 
-	dialCandidates enode.Iterator
+	discmix *enode.FairMix
 
 	// DB interfaces
 	chainDb ctxcdb.Database // Block chain database
@@ -152,6 +153,7 @@ func New(stack *node.Node, config *Config) (*Cortex, error) {
 		bloomRequests:     make(chan chan *bloombits.Retrieval),
 		bloomIndexer:      NewBloomIndexer(chainDb, params.BloomBitsBlocks, params.BloomConfirms),
 		p2pServer:         stack.Server(),
+		discmix:           enode.NewFairMix(0),
 		shutdownTracker:   shutdowncheck.NewShutdownTracker(chainDb),
 	}
 
@@ -252,7 +254,6 @@ func New(stack *node.Node, config *Config) (*Cortex, error) {
 		gpoParams.Default = config.Miner.GasPrice
 	}
 	ctxc.APIBackend.gpo = gasprice.NewOracle(ctxc.APIBackend, gpoParams)
-	ctxc.dialCandidates, err = ctxc.setupDiscovery()
 	if err != nil {
 		return nil, err
 	}
@@ -569,7 +570,7 @@ func (s *Cortex) Protocols() []p2p.Protocol {
 	for i, vsn := range ProtocolVersions {
 		protos[i] = s.protocolManager.makeProtocol(vsn)
 		protos[i].Attributes = []enr.Entry{s.currentCtxcEntry(s.BlockChain())}
-		protos[i].DialCandidates = s.dialCandidates
+		protos[i].DialCandidates = s.discmix
 	}
 	return protos
 }
@@ -578,6 +579,7 @@ func (s *Cortex) Protocols() []p2p.Protocol {
 // Cortex protocol implementation.
 func (s *Cortex) Start(srvr *p2p.Server) error {
 	s.startCtxcEntryUpdate(srvr.LocalNode())
+	s.setupDiscovery()
 	// Start the bloom bits servicing goroutines
 	s.startBloomHandlers(params.BloomBitsBlocks)
 
@@ -594,12 +596,33 @@ func (s *Cortex) Start(srvr *p2p.Server) error {
 	return nil
 }
 
+func (s *Cortex) setupDiscovery() error {
+	dnsclient := dnsdisc.NewClient(dnsdisc.Config{})
+	if len(s.config.DiscoveryURLs) > 0 {
+		iter, err := dnsclient.NewIterator(s.config.DiscoveryURLs...)
+		if err != nil {
+			return err
+		}
+		s.discmix.AddSource(iter)
+	}
+
+	// Add DHT nodes from discv5.
+	if s.p2pServer.DiscoveryV5() != nil {
+		filter := NewNodeFilter(s.blockchain)
+		iter := enode.Filter(s.p2pServer.DiscoveryV5().RandomNodes(), filter)
+		s.discmix.AddSource(iter)
+	}
+
+	return nil
+}
+
 // Stop implements node.Service, terminating all internal goroutines used by the
 // Cortex protocol.
 func (s *Cortex) Stop() error {
 	if s.synapse != nil {
 		s.synapse.Close()
 	}
+	s.discmix.Close()
 	s.protocolManager.Stop()
 	// Then stop everything else.
 	s.bloomIndexer.Close()
diff --git a/ctxc/discovery.go b/ctxc/discovery.go
index 948c9231cd..7a7cd790d5 100644
--- a/ctxc/discovery.go
+++ b/ctxc/discovery.go
@@ -19,7 +19,6 @@ package ctxc
 import (
 	"github.com/CortexFoundation/CortexTheseus/core"
 	"github.com/CortexFoundation/CortexTheseus/core/forkid"
-	"github.com/CortexFoundation/CortexTheseus/p2p/dnsdisc"
 	"github.com/CortexFoundation/CortexTheseus/p2p/enode"
 	"github.com/CortexFoundation/CortexTheseus/rlp"
 )
@@ -65,11 +64,16 @@ func (ctxc *Cortex) currentCtxcEntry(chain *core.BlockChain) *ctxcEntry {
 	}
 }
 
-// setupDiscovery creates the node discovery source for the ctxc protocol.
-func (ctxc *Cortex) setupDiscovery() (enode.Iterator, error) {
-	if len(ctxc.config.DiscoveryURLs) == 0 {
-		return nil, nil
+// NewNodeFilter returns a filtering function that returns whether the provided
+// enode advertises a forkid compatible with the current chain.
+func NewNodeFilter(chain *core.BlockChain) func(*enode.Node) bool {
+	filter := forkid.NewFilter(chain)
+	return func(n *enode.Node) bool {
+		var entry ctxcEntry
+		if err := n.Load(entry); err != nil {
+			return false
+		}
+		err := filter(entry.ForkID)
+		return err == nil
 	}
-	client := dnsdisc.NewClient(dnsdisc.Config{})
-	return client.NewIterator(ctxc.config.DiscoveryURLs...)
 }
diff --git a/nightly.sh b/nightly.sh
index d725e90593..b93e6f1ca4 100755
--- a/nightly.sh
+++ b/nightly.sh
@@ -10,7 +10,7 @@ GOFLAGS=-modcacherw make
 echo "running cortex..."
 #./test.sh > nightly.log 2>&1 &
 rm -rf /tmp/.cortex_test
-./build/bin/cortex --datadir=/tmp/.cortex_test/ --port=0 --storage.mode=lazy --storage.dht > nightly.log 2>&1 &
+./build/bin/cortex --datadir=/tmp/.cortex_test/ --port=0 --storage.mode=lazy --storage.dht --v5disc > nightly.log 2>&1 &
 
 CORTEX_PID=$!