refactor: add option OnGetPeersResponse to make it possible to deal with got peers

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
-	}
-
-	ch := make(chan struct{})
-
-	go func() {
 	neighbors := dht.routingTable.GetNeighbors(
-			newBitmapFromString(infoHash), dht.K)
+		newBitmapFromString(infoHash), dht.routingTable.Len())
 
 	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.

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 {

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()
+	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
-		peers, err := d.GetPeers("546cf15f724d19c4319cc17b179d7e035f89c1f4")
-		if err != nil {
+			err := d.GetPeers("546cf15f724d19c4319cc17b179d7e035f89c1f4")
+			if err != nil && err != dht.ErrNotReady {
+				log.Fatal(err)
+			}
+
+			if err == dht.ErrNotReady {
 				time.Sleep(time.Second * 1)
 				continue
 			}
 
-		fmt.Println("Found peers:", peers)
+			break
 		}
+	}()
+
+	d.Run()
 }