Add timeout mechanism

p2p/p2p.go

@@ -1,15 +1,11 @@
 package p2p
 
 import (
-	"crypto/sha1"
-	"encoding/hex"
-	"fmt"
-	"math"
 	"net"
 	"strconv"
+	"time"
 
 	"github.com/veggiedefender/torrent-client/handshake"
-	"github.com/veggiedefender/torrent-client/message"
 )
 
 // Peer encodes connection information for a peer
@@ -18,8 +14,8 @@ type Peer struct {
 	Port uint16
 }
 
-// Downloader holds data required to download a torrent from a list of peers
-type Downloader struct {
+// Download holds data required to download a torrent from a list of peers
+type Download struct {
 	Peers       []Peer
 	PeerID      [20]byte
 	InfoHash    [20]byte
@@ -27,79 +23,26 @@ type Downloader struct {
 	Length      int
 }
 
-// Download downloads a torrent
-func (d *Downloader) Download() error {
-	conn, err := d.Peers[0].connect(d.PeerID, d.InfoHash)
-	if err != nil {
-		return err
-	}
-	defer conn.Close()
-	h, err := d.handshake(conn)
-	if err != nil {
-		return err
-	}
-	fmt.Println(h)
+type peerState struct {
+	peer *Peer
+	conn net.Conn
+}
 
-	choked := false
-	pieceSize := d.Length / len(d.PieceHashes)
-	buf := make([]byte, pieceSize)
-	i := 0
-	for i < pieceSize {
-		msg, err := message.Read(conn)
-		if err != nil {
-			return err
-		}
-
-		if msg.ID != message.MsgPiece {
-			fmt.Println(msg.String())
-		} else {
-			fmt.Println("Received", len(msg.Payload), "bytes")
-		}
-
-		switch msg.ID {
-		case message.MsgChoke:
-			choked = true
-		case message.MsgUnchoke:
-			choked = false
-		case message.MsgPiece:
-			n, err := message.ParsePiece(0, buf, msg)
-			if err != nil {
-				return err
-			}
-			i += n
-		}
-
-		if !choked {
-			index := 0 // Piece number
-			begin := i // Offset
-			remain := pieceSize - i
-			length := int(math.Min(float64(16384), float64(pieceSize)))
-			length = int(math.Min(float64(remain), float64(length)))
-			_, err := conn.Write(message.FormatRequest(index, begin, length).Serialize())
-			if err != nil {
-				return err
-			}
-		}
-	}
-
-	s := sha1.Sum(buf)
-	fmt.Printf("Downloaded %d bytes.\n", len(buf))
-	fmt.Printf("Got SHA1\t%s\n", hex.EncodeToString(s[:]))
-	fmt.Printf("Expected\t%s\n", hex.EncodeToString(d.PieceHashes[0][:]))
-
-	return nil
+type swarm struct {
+	peerStates []*peerState
 }
 
 func (p *Peer) connect(peerID [20]byte, infoHash [20]byte) (net.Conn, error) {
 	hostPort := net.JoinHostPort(p.IP.String(), strconv.Itoa(int(p.Port)))
-	conn, err := net.Dial("tcp", hostPort)
+	conn, err := net.DialTimeout("tcp", hostPort, 3*time.Second)
 	if err != nil {
 		return nil, err
 	}
 	return conn, nil
 }
 
-func (d *Downloader) handshake(conn net.Conn) (*handshake.Handshake, error) {
+func (d *Download) handshake(conn net.Conn) (*handshake.Handshake, error) {
+	conn.SetDeadline(time.Now().Local().Add(3 * time.Second))
 	req := handshake.New(d.InfoHash, d.PeerID)
 	_, err := conn.Write(req.Serialize())
 	if err != nil {
@@ -110,5 +53,43 @@ func (d *Downloader) handshake(conn net.Conn) (*handshake.Handshake, error) {
 	if err != nil {
 		return nil, err
 	}
+	conn.SetDeadline(time.Time{}) // Disable the deadline
 	return res, nil
 }
+
+func (d *Download) initPeer(p *Peer, c chan *peerState) {
+	conn, err := p.connect(d.PeerID, d.InfoHash)
+	if err != nil {
+		c <- nil
+		return
+	}
+	_, err = d.handshake(conn)
+	if err != nil {
+		c <- nil
+		return
+	}
+	c <- &peerState{p, conn}
+}
+
+func (d *Download) startSwarm() *swarm {
+	c := make(chan *peerState)
+	for i := range d.Peers {
+		go d.initPeer(&d.Peers[i], c)
+	}
+
+	peerStates := make([]*peerState, 0)
+	for range d.Peers {
+		ps := <-c
+		if ps != nil {
+			peerStates = append(peerStates, ps)
+		}
+	}
+
+	return &swarm{peerStates}
+}
+
+// Download downloads a torrent
+func (d *Download) Download() error {
+	d.startSwarm()
+	return nil
+}