Skip to content

Commit 4e3746b

Browse files
committed
feat: add request callback config option
1 parent 1331ba7 commit 4e3746b

File tree

4 files changed

+22
-0
lines changed

4 files changed

+22
-0
lines changed

dht.go

+3
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,8 @@ type IpfsDHT struct {
163163
// addrFilter is used to filter the addresses we put into the peer store.
164164
// Mostly used to filter out localhost and local addresses.
165165
addrFilter func([]ma.Multiaddr) []ma.Multiaddr
166+
167+
onRequestHook func(ctx context.Context, s network.Stream, req pb.Message)
166168
}
167169

168170
// Assert that IPFS assumptions about interfaces aren't broken. These aren't a
@@ -306,6 +308,7 @@ func makeDHT(h host.Host, cfg dhtcfg.Config) (*IpfsDHT, error) {
306308
routingTablePeerFilter: cfg.RoutingTable.PeerFilter,
307309
rtPeerDiversityFilter: cfg.RoutingTable.DiversityFilter,
308310
addrFilter: cfg.AddressFilter,
311+
onRequestHook: cfg.OnRequestHook,
309312

310313
fixLowPeersChan: make(chan struct{}, 1),
311314

dht_net.go

+2
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,8 @@ func (dht *IpfsDHT) handleNewMessage(s network.Stream) bool {
100100
metrics.ReceivedBytes.M(int64(msgLen)),
101101
)
102102

103+
dht.onRequestHook(ctx, s, req)
104+
103105
handler := dht.handlerForMsgType(req.GetType())
104106
if handler == nil {
105107
stats.Record(ctx, metrics.ReceivedMessageErrors.M(1))

dht_options.go

+12
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package dht
22

33
import (
4+
"context"
45
"fmt"
56
"testing"
67
"time"
@@ -12,6 +13,7 @@ import (
1213
"github.com/libp2p/go-libp2p-kbucket/peerdiversity"
1314
record "github.com/libp2p/go-libp2p-record"
1415
"github.com/libp2p/go-libp2p/core/host"
16+
"github.com/libp2p/go-libp2p/core/network"
1517
"github.com/libp2p/go-libp2p/core/peer"
1618
"github.com/libp2p/go-libp2p/core/protocol"
1719

@@ -368,3 +370,13 @@ func WithCustomMessageSender(messageSenderBuilder func(h host.Host, protos []pro
368370
return nil
369371
}
370372
}
373+
374+
// OnRequestHook registers a callback function that will be invoked
375+
// for every incoming DHT protocol message. Note: Ensure that the
376+
// callback executes efficiently, as it can block the entire message handler.
377+
func OnRequestHook(f func(ctx context.Context, s network.Stream, req pb.Message)) Option {
378+
return func(c *dhtcfg.Config) error {
379+
c.OnRequestHook = f
380+
return nil
381+
}
382+
}

internal/config/config.go

+5
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package config
22

33
import (
4+
"context"
45
"fmt"
56
"time"
67

@@ -14,6 +15,7 @@ import (
1415
"github.com/libp2p/go-libp2p-kbucket/peerdiversity"
1516
record "github.com/libp2p/go-libp2p-record"
1617
"github.com/libp2p/go-libp2p/core/host"
18+
"github.com/libp2p/go-libp2p/core/network"
1719
"github.com/libp2p/go-libp2p/core/peer"
1820
"github.com/libp2p/go-libp2p/core/protocol"
1921
ma "github.com/multiformats/go-multiaddr"
@@ -63,6 +65,7 @@ type Config struct {
6365

6466
BootstrapPeers func() []peer.AddrInfo
6567
AddressFilter func([]ma.Multiaddr) []ma.Multiaddr
68+
OnRequestHook func(ctx context.Context, s network.Stream, req pb.Message)
6669

6770
// test specific Config options
6871
DisableFixLowPeers bool
@@ -134,6 +137,8 @@ var Defaults = func(o *Config) error {
134137
// MAGIC: It makes sense to set it to a multiple of OptProvReturnRatio * BucketSize. We chose a multiple of 4.
135138
o.OptimisticProvideJobsPoolSize = 60
136139

140+
o.OnRequestHook = func(ctx context.Context, s network.Stream, req pb.Message) {}
141+
137142
return nil
138143
}
139144

0 commit comments

Comments
 (0)