add new option to set custom

msgSenderBuilder
libp2p · Oct 2, 2024 · 53ef5be · 53ef5be
1 parent 3989875
commit 53ef5be
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/dht_options.go b/dht_options.go
@@ -6,9 +6,11 @@ import (
 	"time"
 
 	dhtcfg "github.com/libp2p/go-libp2p-kad-dht/internal/config"
+	pb "github.com/libp2p/go-libp2p-kad-dht/pb"
 	"github.com/libp2p/go-libp2p-kad-dht/providers"
 	"github.com/libp2p/go-libp2p-kbucket/peerdiversity"
 	record "github.com/libp2p/go-libp2p-record"
+	"github.com/libp2p/go-libp2p/core/host"
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/libp2p/go-libp2p/core/protocol"
 
@@ -356,3 +358,13 @@ func AddressFilter(f func([]ma.Multiaddr) []ma.Multiaddr) Option {
 		return nil
 	}
 }
+
+// --- NEW FEATURE ---
+// WithCustomMessageSender configures the pb.MessageSender of the IpfsDHT to use the
+// custom implementation of the pb.MessageSender
+func WithCustomMessageSender(messageSenderBuilder func(h host.Host, protos []protocol.ID) pb.MessageSenderWithDisconnect) Option {
+	return func(c *dhtcfg.Config) error {
+		c.MsgSenderBuilder = messageSenderBuilder
+		return nil
+	}
+}
diff --git a/internal/config/config.go b/internal/config/config.go
@@ -7,6 +7,8 @@ import (
 	"github.com/ipfs/boxo/ipns"
 	ds "github.com/ipfs/go-datastore"
 	dssync "github.com/ipfs/go-datastore/sync"
+	"github.com/libp2p/go-libp2p-kad-dht/internal/net"
+	pb "github.com/libp2p/go-libp2p-kad-dht/pb"
 	"github.com/libp2p/go-libp2p-kad-dht/providers"
 	"github.com/libp2p/go-libp2p-kbucket/peerdiversity"
 	record "github.com/libp2p/go-libp2p-record"
@@ -48,6 +50,7 @@ type Config struct {
 	ProviderStore          providers.ProviderStore
 	QueryPeerFilter        QueryFilterFunc
 	LookupCheckConcurrency int
+	MsgSenderBuilder       func(h host.Host, protos []protocol.ID) pb.MessageSenderWithDisconnect
 
 	RoutingTable struct {
 		RefreshQueryTimeout time.Duration
@@ -114,6 +117,7 @@ var Defaults = func(o *Config) error {
 	o.EnableProviders = true
 	o.EnableValues = true
 	o.QueryPeerFilter = EmptyQueryFilter
+	o.MsgSenderBuilder = net.NewMessageSenderImpl
 
 	o.RoutingTable.LatencyTolerance = 10 * time.Second
 	o.RoutingTable.RefreshQueryTimeout = 10 * time.Second