diff --git a/dht.go b/dht.go index 76e010b..1e89b58 100644 --- a/dht.go +++ b/dht.go @@ -17,8 +17,13 @@ const ( CrawlMode ) -// ErrNotReady is the error when DHT is not initialized. -var ErrNotReady = errors.New("dht is not ready") +var ( + // ErrNotReady is the error when DHT is not initialized. + ErrNotReady = errors.New("dht is not ready") + // ErrOnGetPeersResponseNotSet is the error that config + // OnGetPeersResponseNotSet is not set when call dht.GetPeers. + ErrOnGetPeersResponseNotSet = errors.New("OnGetPeersResponse is not set") +) // Config represents the configure of dht. type Config struct { @@ -47,6 +52,8 @@ type Config struct { MaxNodes int // callback when got get_peers request OnGetPeers func(string, string, int) + // callback when receive get_peers response + OnGetPeersResponse func(string, *Peer) // callback when got announce_peer request OnAnnouncePeer func(string, string, int) // blcoked ips @@ -229,48 +236,31 @@ func (dht *DHT) id(target string) string { } // GetPeers returns peers who have announced having infoHash. -func (dht *DHT) GetPeers(infoHash string) ([]*Peer, error) { +func (dht *DHT) GetPeers(infoHash string) error { if !dht.Ready { - return nil, ErrNotReady + return ErrNotReady + } + + if dht.OnGetPeersResponse == nil { + return ErrOnGetPeersResponseNotSet } if len(infoHash) == 40 { data, err := hex.DecodeString(infoHash) if err != nil { - return nil, err + return err } infoHash = string(data) } - peers := dht.peersManager.GetPeers(infoHash, dht.K) - if len(peers) != 0 { - return peers, nil + neighbors := dht.routingTable.GetNeighbors( + newBitmapFromString(infoHash), dht.routingTable.Len()) + + for _, no := range neighbors { + dht.transactionManager.getPeers(no, infoHash) } - ch := make(chan struct{}) - - go func() { - neighbors := dht.routingTable.GetNeighbors( - newBitmapFromString(infoHash), dht.K) - - for _, no := range neighbors { - dht.transactionManager.getPeers(no, infoHash) - } - - i := 0 - for _ = range time.Tick(time.Second * 1) { - i++ - peers = dht.peersManager.GetPeers(infoHash, dht.K) - if len(peers) != 0 || i == 30 { - break - } - } - - ch <- struct{}{} - }() - - <-ch - return peers, nil + return nil } // Run starts the dht. diff --git a/krpc.go b/krpc.go index 9a397ae..dc0bd77 100644 --- a/krpc.go +++ b/krpc.go @@ -702,6 +702,9 @@ func handleResponse(dht *DHT, addr *net.UDPAddr, continue } dht.peersManager.Insert(infoHash, p) + if dht.OnGetPeersResponse != nil { + dht.OnGetPeersResponse(infoHash, p) + } } } else if findOn( dht, r, newBitmapFromString(infoHash), getPeersType) != nil { diff --git a/sample/getpeers/getpeers.go b/sample/getpeers/getpeers.go index f9e6b47..45d0d36 100644 --- a/sample/getpeers/getpeers.go +++ b/sample/getpeers/getpeers.go @@ -2,22 +2,34 @@ package main import ( "fmt" - "github.com/shiyanhui/dht" + "log" "time" + + "github.com/shiyanhui/dht" ) func main() { d := dht.New(nil) - go d.Run() - - for { - // ubuntu-14.04.2-desktop-amd64.iso - peers, err := d.GetPeers("546cf15f724d19c4319cc17b179d7e035f89c1f4") - if err != nil { - time.Sleep(time.Second * 1) - continue - } - - fmt.Println("Found peers:", peers) + d.OnGetPeersResponse = func(infoHash string, peer *dht.Peer) { + fmt.Printf("GOT PEER: <%s:%d>\n", peer.IP, peer.Port) } + + go func() { + for { + // ubuntu-14.04.2-desktop-amd64.iso + err := d.GetPeers("546cf15f724d19c4319cc17b179d7e035f89c1f4") + if err != nil && err != dht.ErrNotReady { + log.Fatal(err) + } + + if err == dht.ErrNotReady { + time.Sleep(time.Second * 1) + continue + } + + break + } + }() + + d.Run() }