From 60e215aa817cc66691873ec3bad4b767f807de30 Mon Sep 17 00:00:00 2001
From: BruceWang <bruce.wang@dbappsecurity.com.cn>
Date: Thu, 23 May 2019 10:13:48 +0800
Subject: [PATCH] Add OnQuery Hook

---
 client.go | 1 +
 config.go | 4 ++++
 2 files changed, 5 insertions(+)

diff --git a/client.go b/client.go
index 2d2a3fa7..004ae686 100644
--- a/client.go
+++ b/client.go
@@ -313,6 +313,7 @@ func (cl *Client) newDhtServer(conn net.PacketConn) (s *dht.Server, err error) {
 		}(),
 		StartingNodes:      cl.config.DhtStartingNodes,
 		ConnectionTracking: cl.config.ConnTracker,
+		OnQuery:            cl.config.DHTOnQuery,
 	}
 	s, err = dht.NewServer(&cfg)
 	if err == nil {
diff --git a/config.go b/config.go
index 4ac212d6..f8057b23 100644
--- a/config.go
+++ b/config.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/anacrolix/dht"
+	"github.com/anacrolix/dht/krpc"
 	"github.com/anacrolix/missinggo"
 	"github.com/anacrolix/missinggo/conntrack"
 	"github.com/anacrolix/missinggo/expect"
@@ -124,6 +125,9 @@ type ClientConfig struct {
 	dropDuplicatePeerIds bool
 
 	ConnTracker *conntrack.Instance
+
+	// OnQuery hook func
+	DHTOnQuery func(query *krpc.Msg, source net.Addr) (propagate bool)
 }
 
 func (cfg *ClientConfig) SetListenAddr(addr string) *ClientConfig {