Support custom DHT servers

[btrtrc.git] / client.go

diff --git a/client.go b/client.go
index 8bb1ca38a6f8defdbad2afdab0955d82a424cc49..e5dfa045c2f9f8b63d542c68db17d22de614b13a 100644 (file)
--- a/client.go
+++ b/client.go
@@ -58,7 +58,7 @@ type Client struct {
        onClose        []func()
        dialers        []Dialer
        listeners      []Listener
-       dhtServers     []*dht.Server
+       dhtServers     []DhtServer
        ipBlockList    iplist.Ranger
        // Our BitTorrent protocol extension bytes, sent in our BT handshakes.
        extensionBytes pp.PeerExtensionBits
@@ -101,12 +101,10 @@ func (cl *Client) LocalPort() (port int) {
        return
 }
 
-func writeDhtServerStatus(w io.Writer, s *dht.Server) {
+func writeDhtServerStatus(w io.Writer, s DhtServer) {
        dhtStats := s.Stats()
-       fmt.Fprintf(w, "\t# Nodes: %d (%d good, %d banned)\n", dhtStats.Nodes, dhtStats.GoodNodes, dhtStats.BadNodes)
        fmt.Fprintf(w, "\tServer ID: %x\n", s.ID())
-       fmt.Fprintf(w, "\tAnnounces: %d\n", dhtStats.SuccessfulOutboundAnnouncePeerQueries)
-       fmt.Fprintf(w, "\tOutstanding transactions: %d\n", dhtStats.OutstandingTransactions)
+       spew.Fdump(w, dhtStats)
 }
 
 // Writes out a human readable status of the client, such as for writing to a
@@ -120,7 +118,7 @@ func (cl *Client) WriteStatus(_w io.Writer) {
        fmt.Fprintf(w, "Peer ID: %+q\n", cl.PeerID())
        fmt.Fprintf(w, "Announce key: %x\n", cl.announceKey())
        fmt.Fprintf(w, "Banned IPs: %d\n", len(cl.badPeerIPsLocked()))
-       cl.eachDhtServer(func(s *dht.Server) {
+       cl.eachDhtServer(func(s DhtServer) {
                fmt.Fprintf(w, "%s DHT server at %s:\n", s.Addr().Network(), s.Addr().String())
                writeDhtServerStatus(w, s)
        })
@@ -237,11 +235,11 @@ func NewClient(cfg *ClientConfig) (cl *Client, err error) {
        if !cfg.NoDHT {
                for _, s := range sockets {
                        if pc, ok := s.(net.PacketConn); ok {
-                               ds, err := cl.newDhtServer(pc)
+                               ds, err := cl.newAnacrolixDhtServer(pc)
                                if err != nil {
                                        panic(err)
                                }
-                               cl.dhtServers = append(cl.dhtServers, ds)
+                               cl.dhtServers = append(cl.dhtServers, anacrolixDhtServerWrapper{ds})
                                cl.onClose = append(cl.onClose, func() { ds.Close() })
                        }
                }
@@ -250,6 +248,10 @@ func NewClient(cfg *ClientConfig) (cl *Client, err error) {
        return
 }
 
+func (cl *Client) AddDhtServer(d DhtServer) {
+       cl.dhtServers = append(cl.dhtServers, d)
+}
+
 // Adds a Dialer for outgoing connections. All Dialers are used when attempting to connect to a
 // given address for any Torrent.
 func (cl *Client) AddDialer(d Dialer) {
@@ -300,7 +302,7 @@ func (cl *Client) listenNetworks() (ns []network) {
        return
 }
 
-func (cl *Client) newDhtServer(conn net.PacketConn) (s *dht.Server, err error) {
+func (cl *Client) newAnacrolixDhtServer(conn net.PacketConn) (s *dht.Server, err error) {
        cfg := dht.ServerConfig{
                IPBlocklist:    cl.ipBlockList,
                Conn:           conn,
@@ -335,7 +337,7 @@ func (cl *Client) Closed() <-chan struct{} {
        return cl.closed.C()
 }
 
-func (cl *Client) eachDhtServer(f func(*dht.Server)) {
+func (cl *Client) eachDhtServer(f func(DhtServer)) {
        for _, ds := range cl.dhtServers {
                f(ds)
        }
@@ -929,14 +931,14 @@ func (cl *Client) sendInitialMessages(conn *connection, torrent *Torrent) {
 }
 
 func (cl *Client) dhtPort() (ret uint16) {
-       cl.eachDhtServer(func(s *dht.Server) {
+       cl.eachDhtServer(func(s DhtServer) {
                ret = uint16(missinggo.AddrPort(s.Addr()))
        })
        return
 }
 
 func (cl *Client) haveDhtServer() (ret bool) {
-       cl.eachDhtServer(func(_ *dht.Server) {
+       cl.eachDhtServer(func(_ DhtServer) {
                ret = true
        })
        return
@@ -1071,7 +1073,7 @@ func (cl *Client) AddTorrentInfoHashWithStorage(infoHash metainfo.Hash, specStor
        new = true
 
        t = cl.newTorrent(infoHash, specStorage)
-       cl.eachDhtServer(func(s *dht.Server) {
+       cl.eachDhtServer(func(s DhtServer) {
                go t.dhtAnnouncer(s)
        })
        cl.torrents[infoHash] = t
@@ -1188,7 +1190,7 @@ func (cl *Client) AddTorrentFromFile(filename string) (T *Torrent, err error) {
        return cl.AddTorrent(mi)
 }
 
-func (cl *Client) DhtServers() []*dht.Server {
+func (cl *Client) DhtServers() []DhtServer {
        return cl.dhtServers
 }
 
@@ -1206,7 +1208,7 @@ func (cl *Client) AddDHTNodes(nodes []string) {
                                Port: hmp.Port,
                        },
                }
-               cl.eachDhtServer(func(s *dht.Server) {
+               cl.eachDhtServer(func(s DhtServer) {
                        s.AddNode(ni)
                })
        }