From 06b849a44d3de5c2233dadc673f8e84188792905 Mon Sep 17 00:00:00 2001
From: Jesse Li <jessejesse123@gmail.com>
Date: Thu, 2 Jan 2020 11:49:08 -0500
Subject: [PATCH] Fix bug which produced incorrect piece size

---
 p2p/p2p.go                 | 20 ++++++++++++++------
 torrentfile/torrentfile.go |  1 +
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/p2p/p2p.go b/p2p/p2p.go
index 52c356e..9276253 100644
--- a/p2p/p2p.go
+++ b/p2p/p2p.go
@@ -24,6 +24,7 @@ type Torrent struct {
 	PeerID      [20]byte
 	InfoHash    [20]byte
 	PieceHashes [][20]byte
+	PieceLength int
 	Length      int
 	Name        string
 }
@@ -181,13 +182,20 @@ func (t *Torrent) startDownloadWorker(peer peers.Peer, workQueue chan *pieceWork
 	}
 }
 
-func calculateBoundsForPiece(index, numPieces, length int) (begin int, end int) {
-	pieceLength := length / numPieces
-	begin = index * pieceLength
-	end = begin + pieceLength
+func (t *Torrent) calculateBoundsForPiece(index int) (begin int, end int) {
+	begin = index * t.PieceLength
+	end = begin + t.PieceLength
+	if end > t.Length {
+		end = t.Length
+	}
 	return begin, end
 }
 
+func (t *Torrent) calculatePieceSize(index int) int {
+	begin, end := t.calculateBoundsForPiece(index)
+	return end - begin
+}
+
 // Download downloads the torrent. This stores the entire file in memory.
 func (t *Torrent) Download() ([]byte, error) {
 	log.Println("Starting download for", t.Name)
@@ -195,7 +203,7 @@ func (t *Torrent) Download() ([]byte, error) {
 	workQueue := make(chan *pieceWork, len(t.PieceHashes))
 	results := make(chan *pieceResult, len(t.PieceHashes))
 	for index, hash := range t.PieceHashes {
-		length := t.Length / len(t.PieceHashes)
+		length := t.calculatePieceSize(index)
 		workQueue <- &pieceWork{index, hash, length}
 	}
 
@@ -209,7 +217,7 @@ func (t *Torrent) Download() ([]byte, error) {
 	donePieces := 0
 	for donePieces < len(t.PieceHashes) {
 		res := <-results
-		begin, end := calculateBoundsForPiece(res.index, len(t.PieceHashes), t.Length)
+		begin, end := t.calculateBoundsForPiece(res.index)
 		copy(buf[begin:end], res.buf)
 		donePieces++
 
diff --git a/torrentfile/torrentfile.go b/torrentfile/torrentfile.go
index b7953e4..8d6cafc 100644
--- a/torrentfile/torrentfile.go
+++ b/torrentfile/torrentfile.go
@@ -54,6 +54,7 @@ func (t *TorrentFile) Download() ([]byte, error) {
 		PeerID:      peerID,
 		InfoHash:    t.InfoHash,
 		PieceHashes: t.PieceHashes,
+		PieceLength: t.PieceLength,
 		Length:      t.Length,
 		Name:        t.Name,
 	}