download.go

Documentation: github.com/concurrency-8/torrent

     1  package torrent
     2  
     3  import (
     4  	"bytes"
     5  	"crypto/sha1"
     6  	"encoding/binary"
     7  	"encoding/gob"
     8  	"io"
     9  	"log"
    10  	"net"
    11  	"net/url"
    12  	"os"
    13  	"path/filepath"
    14  	"sync"
    15  	"time"
    16  
    17  	"github.com/concurrency-8/args"
    18  	"github.com/concurrency-8/parser"
    19  	"github.com/concurrency-8/piece"
    20  	"github.com/concurrency-8/queue"
    21  	"github.com/concurrency-8/tracker"
    22  )
    23  
    24  type handler func(tracker.Peer, []byte, net.Conn, *piece.PieceTracker, *queue.Queue, *tracker.ClientStatusReport) error
    25  
    26  // MaxTry is the maximum number of times we should try to connect to a tracker
    27  var MaxTry int = 1
    28  
    29  // TCPTimeout is the maximum time for which one must wait for connection to a peer
    30  var TCPTimeout time.Duration = 15
    31  
    32  // ReadTimeout is the maximum time for which one must wait for the nect message from the peer. If no message arrives till this point, handshake again
    33  var ReadTimeout time.Duration = 150
    34  
    35  var wg sync.WaitGroup
    36  
    37  // Info is logger for information
    38  var Info *log.Logger
    39  
    40  // Error is logger for errors
    41  var Error *log.Logger
    42  
    43  // DownloadFromFile downloads torrent from path using port
    44  func DownloadFromFile(path string, port int) {
    45  
    46  	// Set up logs
    47  	logFolder := filepath.Join("Logs", path)
    48  	os.MkdirAll(logFolder, os.ModePerm)
    49  	logFile, _ := os.Create(filepath.Join(logFolder, "Download.log"))
    50  
    51  	Info = log.New(logFile, "INFO ", log.Ldate|log.Ltime|log.Lshortfile)
    52  	Error = log.New(logFile, "ERROR ", log.Ldate|log.Ltime|log.Lshortfile)
    53  
    54  	torrentFile, err := parser.ParseFromFile(path)
    55  	if err != nil {
    56  		Error.Println("Unable to open torrentfile", err)
    57  		panic(err)
    58  	}
    59  	Info.Println("TorrentFile parsed")
    60  
    61  	// Generate client status report
    62  	clientReport := tracker.GetClientStatusReport(torrentFile, uint16(port))
    63  
    64  	// Getting peer list from one announce url only for now.
    65  	var announceResp *tracker.AnnounceResponse
    66  	for _, announceURL := range torrentFile.Announce {
    67  		u, err := url.Parse(announceURL)
    68  		if err != nil {
    69  			panic(err)
    70  		}
    71  		Info.Println("Contacting tracker[", announceURL, "] for peer list...")
    72  		count := 0
    73  		for count < MaxTry {
    74  			count++
    75  			announceResp, err = tracker.GetPeers(u, clientReport)
    76  			if err == nil {
    77  				break
    78  			}
    79  			Info.Println("Failed(", err, "). Trying again...")
    80  		}
    81  		if err == nil {
    82  			break
    83  		}
    84  	}
    85  
    86  	if announceResp == nil {
    87  		panic("Unable to receive peers! Problem with the torrent or internet")
    88  	}
    89  
    90  	pieceTracker := piece.NewPieceTracker(torrentFile)
    91  	if args.ARGS.Resume {
    92  		readGob(torrentFile.Name+"/resume.gob", pieceTracker)
    93  	}
    94  	// DownloadFromPeer(announceResp.Peers[0], clientReport, pieceTracker)
    95  	wg.Add(len(announceResp.Peers))
    96  	for _, peer := range announceResp.Peers {
    97  		Info.Println("Spawning peer thread: peer<", peer, ">")
    98  		go DownloadFromPeer(peer, clientReport, pieceTracker)
    99  	}
   100  
   101  	// DownloadFromPeer(announceResp.Peers[0], clientReport, pieceTracker)
   102  
   103  	wg.Wait()
   104  	pieceTracker.PrintPercentageDone()
   105  	Info.Println("All peer threads finished!")
   106  }
   107  
   108  // DownloadFromPeer is a function that handshakes with a peer specified by peer object.
   109  // Concurrently call this function to establish parallel connections to many peers.
   110  func DownloadFromPeer(peer tracker.Peer, report *tracker.ClientStatusReport, pieces *piece.PieceTracker) error {
   111  	defer wg.Done()
   112  
   113  	//safely handle reading using onWholeMessage
   114  
   115  	queue := queue.NewQueue(report.TorrentFile)
   116  
   117  	exitStatus := 1
   118  	var err error
   119  	for exitStatus == 1 && err == nil {
   120  		queue.Choked = true
   121  		conn, err := sendHandshake(peer, report)
   122  		if err != nil {
   123  			break
   124  		}
   125  		exitStatus, err = onWholeMessage(peer, conn, msgHandler, pieces, queue, report)
   126  		if err != nil {
   127  			break
   128  		}
   129  	}
   130  
   131  	Info.Println("peer: <", peer, ">: ends!")
   132  	return err
   133  }
   134  
   135  func sendHandshake(peer tracker.Peer, report *tracker.ClientStatusReport) (conn net.Conn, err error) {
   136  	buffer, err := BuildHandshake(*report)
   137  	if err != nil {
   138  		return nil, err
   139  	}
   140  	peerip := make([]byte, 4)
   141  	binary.BigEndian.PutUint32(peerip, peer.IPAdress)
   142  	service := net.TCPAddr{
   143  		IP:   peerip,
   144  		Port: int(peer.Port),
   145  		Zone: "",
   146  	}
   147  	Info.Println("peer: <", peer, ">: Dialing TCP connection")
   148  	d := net.Dialer{Timeout: TCPTimeout * time.Second}
   149  	err = nil
   150  	count := 0
   151  	for count < MaxTry {
   152  		count++
   153  		conn, err = d.Dial("tcp", service.String())
   154  		if err != nil {
   155  			Info.Println("peer: <", peer, ">: Unable to set up TCP connection: ", count)
   156  		} else {
   157  			Info.Println("peer: <", peer, ">: Successfully connected to Peer")
   158  			break
   159  		}
   160  	}
   161  
   162  	if err != nil {
   163  		Error.Println("peer: <", peer, ">: Could not connect to peer!")
   164  		return nil, err
   165  	}
   166  	Info.Println("peer: <", peer, ">: Handshaking")
   167  
   168  	//write the handshake content into the connection.
   169  	_, err = conn.Write(buffer.Bytes())
   170  	if err != nil {
   171  		return nil, err
   172  	}
   173  
   174  	return conn, nil
   175  }
   176  
   177  func msgHandler(peer tracker.Peer, msg []byte, conn net.Conn, pieces *piece.PieceTracker, queue *queue.Queue, report *tracker.ClientStatusReport) error {
   178  	// Info.Println("peer: <", peer, ">: Message:", msg)
   179  
   180  	if (len(msg) == int(uint8(msg[0]))+49) && (bytes.Equal(msg[1:20], []byte("BitTorrent protocol"))) {
   181  		Info.Println("peer: <", peer, ">: Handshake successful")
   182  		message, err := BuildInterested()
   183  		if err != nil {
   184  			Info.Println("peer: <", peer, ">: Error", err.Error())
   185  			return err
   186  		}
   187  		conn.Write(message.Bytes())
   188  		// Info.Println("peer: <", peer, ">: Request(", len(message.Bytes()), "): ", message.Bytes())
   189  	} else {
   190  
   191  		_, id, payload := ParseMsg(bytes.NewBuffer(msg))
   192  
   193  		if id == 0 {
   194  			Info.Println("peer: <", peer, ">: Choke")
   195  			ChokeHandler(peer, conn, pieces, report)
   196  		}
   197  		if id == 1 {
   198  			Info.Println("peer: <", peer, ">: Unchoke")
   199  			UnchokeHandler(peer, conn, pieces, queue)
   200  		}
   201  		if id == 4 {
   202  			Info.Println("peer: <", peer, ">: Have")
   203  			HaveHandler(peer, conn, pieces, queue, payload)
   204  		}
   205  		if id == 5 {
   206  			Info.Println("peer: <", peer, ">: BitField")
   207  			BitFieldHandler(peer, conn, pieces, queue, payload)
   208  		}
   209  		if id == 7 {
   210  			Info.Println("peer: <", peer, ">: Piece")
   211  			PieceHandler(peer, conn, pieces, queue, report, parser.PieceBlock{
   212  				Index: payload["index"].(uint32),
   213  				Begin: payload["begin"].(uint32),
   214  				Bytes: payload["block"].(*bytes.Buffer).Bytes(),
   215  			})
   216  		}
   217  	}
   218  
   219  	return nil
   220  
   221  }
   222  
   223  // onWholeMessage sends complete messages to callback function
   224  func onWholeMessage(peer tracker.Peer, conn net.Conn, msgHandler handler, pieces *piece.PieceTracker, queue *queue.Queue, report *tracker.ClientStatusReport) (status int, err error) {
   225  	buffer := new(bytes.Buffer)
   226  	handshake := true
   227  	resp := make([]byte, 1000)
   228  	msgLen := -1
   229  	count := 0
   230  	for pieces != nil && !pieces.IsDone() {
   231  		conn.SetReadDeadline(time.Now().Add(ReadTimeout * time.Second)) // Setting Read deadline from a connection
   232  		respLen, err := conn.Read(resp)
   233  		//Please look for a better connection handling in the future.
   234  		//Maybe use defer?
   235  
   236  		if err != nil {
   237  			if netErr, ok := err.(net.Error); ok && netErr.Timeout() && !handshake {
   238  				count++
   239  				if count < MaxTry {
   240  					Info.Println("Timeout error - Try again")
   241  					continue
   242  				} else {
   243  					Info.Println("peer: <", peer, ">: Many timeout errors - Peer not responding. Should try to reconnect")
   244  					conn.Close()
   245  					return 1, err
   246  				}
   247  			} else if err != io.EOF {
   248  				Error.Println("peer: <", peer, ">: Error while reading from connection: ", err)
   249  				Info.Println("peer: <", peer, ">: Restarting connection")
   250  				conn.Close()
   251  				return 1, err
   252  			} else {
   253  				Error.Println("peer: <", peer, ">: Peer does not respond. Stop contacting this peer.")
   254  				return 0, err
   255  			}
   256  		}
   257  
   258  		binary.Write(buffer, binary.BigEndian, resp[:respLen])
   259  
   260  		if handshake {
   261  			Info.Println("peer: <", peer, ">: First message from peer afte connection starts - Must be handshake")
   262  			length := uint8((buffer.Bytes())[0])
   263  			msgLen = int(length) + 49
   264  		} else if msgLen == -1 {
   265  			length := binary.BigEndian.Uint32(buffer.Bytes()[0:4])
   266  			// length := uint32((buffer.Bytes())[0:4])
   267  			msgLen = int(length) + 4
   268  			Info.Println("peer: <", peer, ">: New message reception started, len =", msgLen)
   269  			// Info.Println("peer: <", peer, ">: Setting msgLen to", msgLen)
   270  		}
   271  
   272  		for len(buffer.Bytes()) >= 4 && msgLen != -1 && len(buffer.Bytes()) >= msgLen {
   273  			Info.Println("peer: <", peer, ">: Message received, msgLen =", msgLen)
   274  			messageBytes := make([]byte, msgLen)
   275  			binary.Read(buffer, binary.BigEndian, messageBytes)
   276  			// Info.Println("peer: <", peer, ">: msgLen:", msgLen)
   277  			msgHandler(peer, messageBytes, conn, pieces, queue, report)
   278  			Info.Println("peer: <", peer, ">: Message handled - setting msgLen = -1")
   279  			msgLen = -1
   280  			handshake = false
   281  			if len(buffer.Bytes()) > 4 {
   282  				length := binary.BigEndian.Uint32(buffer.Bytes()[0:4])
   283  				msgLen = int(length) + 4
   284  				Info.Println("peer: <", peer, ">: New message was in previous one - msgLen =", msgLen)
   285  				// Info.Println("peer: <", peer, ">: Setting msgLen to", msgLen)
   286  			}
   287  		}
   288  	}
   289  	return 0, nil
   290  }
   291  
   292  // ChokeHandler handles choking protocol
   293  func ChokeHandler(peer tracker.Peer, conn net.Conn, pieces *piece.PieceTracker, report *tracker.ClientStatusReport) {
   294  	Info.Println("peer:<", peer, ">: Choke")
   295  	if pieces != nil && pieces.IsDone() {
   296  		Info.Println("All pieces done. Closing connection.")
   297  		conn.Close()
   298  	} else if report != nil {
   299  		Info.Println("peer: <", peer, ">: Handshaking again")
   300  		// time.Sleep(2 * time.Second) // Sleep for 2 seconds and try handshaking again
   301  		handshake, err := BuildHandshake(*report)
   302  		if err != nil {
   303  			panic("Problem with the torrentFile")
   304  		} else {
   305  			conn.Write(handshake.Bytes())
   306  		}
   307  	}
   308  }
   309  
   310  // UnchokeHandler handles unchoking protocol
   311  func UnchokeHandler(peer tracker.Peer, conn net.Conn, pieces *piece.PieceTracker, queue *queue.Queue) {
   312  	if queue.Choked && queue.Length() != 0 {
   313  		Info.Println("peer:<", peer, "> Unchoke: queue was choked, but queue was non-empty")
   314  		queue.Choked = false
   315  		Info.Println("peer:<", peer, ">: Requesting next piece")
   316  		RequestPiece(peer, conn, pieces, queue)
   317  	} else if queue.Choked {
   318  		Info.Println("peer:<", peer, ">: Unchoke - Queue empty and choked - Sending interested")
   319  		queue.Choked = false
   320  		message, _ := BuildInterested()
   321  		// if err != nil {
   322  		// 	Info.Println("peer: <", peer, ">: Error", err.Error())
   323  		// 	return err
   324  		// }
   325  		if conn != nil {
   326  			conn.Write(message.Bytes())
   327  		}
   328  	}
   329  	// Info.Println("peer: <", peer, ">: RequestPiece : Called from Unchokehandler")
   330  }
   331  
   332  // HaveHandler handles Have protocol
   333  func HaveHandler(peer tracker.Peer, conn net.Conn, pieces *piece.PieceTracker, queue *queue.Queue, payload Payload) (pieceIndex uint32, err error) {
   334  	binary.Read(payload["payload"].(*bytes.Buffer), binary.BigEndian, &pieceIndex)
   335  	queueempty := (queue.Length() == 0)
   336  	err = queue.Enqueue(pieceIndex)
   337  	if err != nil {
   338  		return
   339  	}
   340  	if queueempty {
   341  		Info.Println("peer: <", peer, ">: HaveHandler: Queue was empty. Requesting pieces.")
   342  		err = RequestPiece(peer, conn, pieces, queue)
   343  	}
   344  	return
   345  }
   346  
   347  // BitFieldHandler handles bitfield protocol
   348  func BitFieldHandler(peer tracker.Peer, conn net.Conn, pieces *piece.PieceTracker, queue *queue.Queue, payload Payload) (err error) {
   349  	queueempty := (queue.Length() == 0)
   350  	msg := payload["payload"]
   351  	for i, bytevalue := range msg.(*bytes.Buffer).Bytes() {
   352  		for j := 7; j >= 0; j-- {
   353  			if 1 == bytevalue&1 {
   354  				err = queue.Enqueue(uint32(i*8 + j))
   355  			}
   356  			bytevalue = bytevalue >> 1
   357  		}
   358  	}
   359  	if queueempty {
   360  		Info.Println("peer: <", peer, ">: BitFieldHandler: Queue was empty. Requesting pieces")
   361  		err = RequestPiece(peer, conn, pieces, queue)
   362  	}
   363  
   364  	return
   365  }
   366  
   367  // PieceHandler - TODO Write comment
   368  func PieceHandler(peer tracker.Peer, conn net.Conn, pieces *piece.PieceTracker, queue *queue.Queue, report *tracker.ClientStatusReport, pieceResp parser.PieceBlock) {
   369  	pieces.AddReceived(pieceResp)
   370  
   371  	Info.Println("peer: <", peer, ">: Received piece[", pieceResp.Index, "] [", pieceResp.Begin/parser.BLOCK_LEN, "]")
   372  	report.Data[pieceResp.Index].Blocks[pieceResp.Begin/parser.BLOCK_LEN] = pieceResp
   373  
   374  	toSHA1 := func(data []byte) []byte {
   375  		hash := sha1.New()
   376  		hash.Write(data)
   377  		return hash.Sum(nil)
   378  	}
   379  	var piece []byte
   380  	if pieces.PieceIsDone(pieceResp.Index) {
   381  		for _, i := range report.Data[pieceResp.Index].Blocks {
   382  			if len(i.Bytes) != 0 {
   383  				piece = append(piece, i.Bytes...)
   384  			}
   385  		}
   386  		// numZeros := int(report.TorrentFile.PieceLength) - len(piece)
   387  		// for i := 0; i < numZeros; i++ {
   388  		// 	piece = append(piece, byte(0))
   389  		// }
   390  		same := true
   391  		expected := report.TorrentFile.Piece[pieceResp.Index*20 : (pieceResp.Index+1)*20]
   392  		actual := toSHA1(piece)
   393  		for i := range expected {
   394  			same = same && expected[i] == actual[i]
   395  		}
   396  		if !same {
   397  			Error.Println("peer: <", peer, ">: SHA do not match for piece:", pieceResp.Index)
   398  			Error.Println("peer: <", peer, ">: Expected:\t", report.TorrentFile.Piece[pieceResp.Index*20:(pieceResp.Index+1)*20])
   399  			Error.Println("peer: <", peer, ">: Actual:\t", toSHA1(piece))
   400  			report.Data[pieceResp.Index].Blocks[pieceResp.Begin/parser.BLOCK_LEN] = parser.PieceBlock{}
   401  
   402  			pieces.Reset(pieceResp.Index)
   403  			queue.Enqueue(pieceResp.Index)
   404  			Info.Println("peer: <", peer, ">: Reset queue and pieceTracker for", pieceResp.Index)
   405  			RequestPiece(peer, conn, pieces, queue)
   406  			return
   407  		}
   408  		Info.Println("peer: <", peer, ">: Piece[", pieceResp.Index, "] downloaded SUCCESSFULLY!")
   409  	}
   410  
   411  	// Info.Println("Bytes Received : ", pieceResp.Bytes)
   412  
   413  	offsetInFile := uint64(pieceResp.Index)*uint64(report.TorrentFile.PieceLength) + uint64(pieceResp.Begin)
   414  	file := report.TorrentFile.Files[0].FilePointer
   415  	for key, value := range report.TorrentFile.Files {
   416  		if offsetInFile > value.Length {
   417  			offsetInFile -= value.Length
   418  			file = report.TorrentFile.Files[key+1].FilePointer
   419  		} else {
   420  			break
   421  		}
   422  	}
   423  	Info.Println("peer: <", peer, ">: Writing block to file ", file.Name())
   424  	file.WriteAt(pieceResp.Bytes, int64(offsetInFile))
   425  	if args.ARGS.ResumeCapability {
   426  		writeGob(report.TorrentFile.Name+"/resume.gob", pieces)
   427  	}
   428  	// file.Sync()
   429  	pieces.PrintPercentageDone()
   430  
   431  	if pieces.IsDone() {
   432  		for _, file := range report.TorrentFile.Files {
   433  			defer file.FilePointer.Close()
   434  		}
   435  		Info.Println("peer: <", peer, ">: Done")
   436  		conn.Close()
   437  	} else {
   438  		Info.Println("peer<", peer, " >: Called from piecehandler")
   439  		RequestPiece(peer, conn, pieces, queue)
   440  	}
   441  }
   442  
   443  var pieceTrackerLock sync.Mutex
   444  
   445  // RequestPiece requests a piece
   446  func RequestPiece(peer tracker.Peer, conn net.Conn, pieces *piece.PieceTracker, queue *queue.Queue) (err error) {
   447  	if queue.Choked {
   448  		Error.Println("peer: <", peer, ">: Queue is choked")
   449  		return
   450  	}
   451  
   452  	for queue.Length() > 0 {
   453  		pieceBlock, err := queue.Peek()
   454  
   455  		if err != nil {
   456  			break
   457  		}
   458  
   459  		err = queue.Dequeue()
   460  
   461  		if err != nil {
   462  			break
   463  		}
   464  
   465  		pieceTrackerLock.Lock()
   466  		if pieces.Needed(pieceBlock) {
   467  			pieces.AddRequested(pieceBlock)
   468  			pieceTrackerLock.Unlock()
   469  			Info.Println("peer: <", peer, ">: Requesting piece[", pieceBlock.Index, "][", pieceBlock.Begin/parser.BLOCK_LEN, "]")
   470  			message, err := BuildRequest(pieceBlock)
   471  
   472  			if err != nil {
   473  				break
   474  			}
   475  			_, err = conn.Write(message.Bytes())
   476  
   477  			if err != nil {
   478  				Info.Println(err.Error())
   479  				break
   480  			}
   481  			break
   482  		} else {
   483  			pieceTrackerLock.Unlock()
   484  		}
   485  	}
   486  	return
   487  }
   488  
   489  func writeGob(filePath string, object interface{}) error {
   490  	file, err := os.Create(filePath)
   491  	if err == nil {
   492  		encoder := gob.NewEncoder(file)
   493  		encoder.Encode(object)
   494  	}
   495  	file.Close()
   496  	return err
   497  }
   498  
   499  func readGob(filePath string, object interface{}) error {
   500  	file, err := os.Open(filePath)
   501  	if err == nil {
   502  		decoder := gob.NewDecoder(file)
   503  		err = decoder.Decode(object)
   504  	}
   505  	file.Close()
   506  	return err
   507  }
   508
View as plain text