diff --git a/dht.go b/dht.go index 45aec84..ff479e8 100644 --- a/dht.go +++ b/dht.go @@ -25,9 +25,14 @@ const ( // Config represents the configure of dht. type Config struct { - K int - KBucketSize int - Network string + // in mainline dht, k = 8 + K int + // for crawling mode, we put all nodes in one bucket, so KBucketSize may + // not be K + KBucketSize int + // candidates are udp, udp4, udp6 + Network string + // format is `ip:port` Address string PrimeNodes []string KBucketExpiredAfter time.Duration diff --git a/krpc.go b/krpc.go index a87a25e..04af4ca 100644 --- a/krpc.go +++ b/krpc.go @@ -573,7 +573,9 @@ func handleRequest(dht *DHT, addr *net.UDPAddr, port = addr.Port } - dht.peersManager.Insert(infoHash, newPeer(addr.IP, port, token)) + if dht.IsStandardMode() { + dht.peersManager.Insert(infoHash, newPeer(addr.IP, port, token)) + } send(dht.conn, addr, makeResponse(t, map[string]interface{}{ "id": dht.id(id), diff --git a/peerwire.go b/peerwire.go index 8ad3095..112c625 100644 --- a/peerwire.go +++ b/peerwire.go @@ -38,17 +38,12 @@ var handshakePrefix = []byte{ } // read reads size-length bytes from conn to data. -func read(conn *net.TCPConn, size int, data *bytes.Buffer) (err error) { - buffer := make([]byte, size) - if _, err = io.ReadFull(conn, buffer); err != nil { - return - } - - n, err := data.Write(buffer) - if err != nil || n != size { +func read(conn *net.TCPConn, size int, data *bytes.Buffer) error { + n, err := io.CopyN(data, conn, int64(size)) + if err != nil || n != int64(size) { return errors.New("read error") } - return + return nil } // readMessage gets a message from the tcp connection. @@ -215,11 +210,22 @@ func (wire *Wire) requestPieces( sendMessage(conn, buffer[:length]) } + buffer = nil } // fetchMetadata fetchs medata info accroding to infohash from dht. func (wire *Wire) fetchMetadata(r Request) { + var ( + length int + msgType byte + piecesNum int + pieces [][]byte + utMetadata int + metadataSize int + ) + defer func() { + pieces = nil recover() }() @@ -245,15 +251,6 @@ func (wire *Wire) fetchMetadata(r Request) { return } - var ( - length int - msgType byte - pieces [][]byte - piecesNum int - utMetadata int - metadataSize int - ) - for { length, err = readMessage(conn, data) if err != nil { diff --git a/routingtable.go b/routingtable.go index d582f3e..95a43de 100644 --- a/routingtable.go +++ b/routingtable.go @@ -119,7 +119,12 @@ func (pm *peersManager) Insert(infoHash string, peer *Peer) { pm.Unlock() v, _ := pm.table.Get(infoHash) - v.(*keyedDeque).Push(peer.CompactIPPortInfo(), peer) + queue := v.(*keyedDeque) + + queue.Push(peer.CompactIPPortInfo(), peer) + if queue.Len() > pm.dht.K { + queue.Remove(queue.Front()) + } } // GetPeers returns size-length peers who announces having infoHash. @@ -326,6 +331,7 @@ type routingTable struct { cachedNodes *syncedMap cachedKBuckets *keyedDeque dht *DHT + clearQueue *syncedList } // newRoutingTable returns a new routingTable pointer. @@ -339,6 +345,7 @@ func newRoutingTable(k int, dht *DHT) *routingTable { cachedNodes: newSyncedMap(), cachedKBuckets: newKeyedDeque(), dht: dht, + clearQueue: newSyncedList(), } rt.cachedKBuckets.Push(root.bucket.prefix.String(), root.bucket) @@ -495,7 +502,6 @@ func (rt *routingTable) RemoveByAddr(address string) { // Fresh sends findNode to all nodes in the expired nodes. func (rt *routingTable) Fresh() { now := time.Now() - queue := newSyncedList() for e := range rt.cachedKBuckets.Iter() { bucket := e.Value.(*kbucket) @@ -507,13 +513,17 @@ func (rt *routingTable) Fresh() { for e := range bucket.nodes.Iter() { no := e.Value.(*node) rt.dht.transactionManager.findNode(no, bucket.RandomChildID()) - queue.PushBack(no) + rt.clearQueue.PushBack(no) } } - for e := range queue.Iter() { - rt.Remove(e.Value.(*node).id) + if rt.dht.IsCrawlMode() { + for e := range rt.clearQueue.Iter() { + rt.Remove(e.Value.(*node).id) + } } + + rt.clearQueue.Clear() } // Len returns the number of nodes in table.