From c9235820da1d6ffa5dcfebfa4c6d9697140f1db3 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Fri, 1 Mar 2024 23:54:56 +1100
Subject: [PATCH] Fix file storage segments for v2 torrents

---
 common/upverted_files.go |  9 +++++++++
 segments/index.go        |  6 ++++++
 storage/file-piece.go    | 16 ++++++++++++----
 storage/file.go          |  2 +-
 torrent.go               |  6 +++++-
 webseed/client.go        |  2 +-
 6 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/common/upverted_files.go b/common/upverted_files.go
index 1933e16a..8984eb25 100644
--- a/common/upverted_files.go
+++ b/common/upverted_files.go
@@ -16,3 +16,12 @@ func LengthIterFromUpvertedFiles(fis []metainfo.FileInfo) segments.LengthIter {
 		return l, true
 	}
 }
+
+// Returns file segments, BitTorrent v2 aware.
+func TorrentOffsetFileSegments(info *metainfo.Info) (ret []segments.Extent) {
+	files := info.UpvertedFiles()
+	for _, fi := range files {
+		ret = append(ret, segments.Extent{fi.TorrentOffset, fi.Length})
+	}
+	return
+}
diff --git a/segments/index.go b/segments/index.go
index 6717dcba..888e90a8 100644
--- a/segments/index.go
+++ b/segments/index.go
@@ -17,6 +17,10 @@ type Index struct {
 	segments []Extent
 }
 
+func NewIndexFromSegments(segments []Extent) Index {
+	return Index{segments}
+}
+
 func (me Index) iterSegments() func() (Length, bool) {
 	return func() (Length, bool) {
 		if len(me.segments) == 0 {
@@ -29,6 +33,8 @@ func (me Index) iterSegments() func() (Length, bool) {
 	}
 }
 
+// Returns true if the callback returns false early, or extents are found in the index for all parts
+// of the given extent.
 func (me Index) Locate(e Extent, output Callback) bool {
 	first := sort.Search(len(me.segments), func(i int) bool {
 		_e := me.segments[i]
diff --git a/storage/file-piece.go b/storage/file-piece.go
index 47772017..99ff5fc8 100644
--- a/storage/file-piece.go
+++ b/storage/file-piece.go
@@ -1,6 +1,7 @@
 package storage
 
 import (
+	"github.com/anacrolix/torrent/segments"
 	"io"
 	"log"
 	"os"
@@ -32,12 +33,19 @@ func (fs *filePieceImpl) Completion() Completion {
 	verified := true
 	if c.Complete {
 		// If it's allegedly complete, check that its constituent files have the necessary length.
-		for _, fi := range extentCompleteRequiredLengths(fs.p.Info, fs.p.Offset(), fs.p.Length()) {
-			s, err := os.Stat(fs.files[fi.fileIndex].path)
-			if err != nil || s.Size() < fi.length {
+		if !fs.segmentLocater.Locate(segments.Extent{
+			Start:  fs.p.Offset(),
+			Length: fs.p.Length(),
+		}, func(i int, extent segments.Extent) bool {
+			file := fs.files[i]
+			s, err := os.Stat(file.path)
+			if err != nil || s.Size() < extent.Start+extent.Length {
 				verified = false
-				break
+				return false
 			}
+			return true
+		}) {
+			panic("files do not cover piece extent")
 		}
 	}
 
diff --git a/storage/file.go b/storage/file.go
index b8739647..231825fd 100644
--- a/storage/file.go
+++ b/storage/file.go
@@ -83,7 +83,7 @@ func (fs fileClientImpl) OpenTorrent(info *metainfo.Info, infoHash metainfo.Hash
 	}
 	t := &fileTorrentImpl{
 		files,
-		segments.NewIndex(common.LengthIterFromUpvertedFiles(upvertedFiles)),
+		segments.NewIndexFromSegments(common.TorrentOffsetFileSegments(info)),
 		infoHash,
 		fs.opts.PieceCompletion,
 	}
diff --git a/torrent.go b/torrent.go
index 64e80660..a361c4ed 100644
--- a/torrent.go
+++ b/torrent.go
@@ -1133,7 +1133,11 @@ func (t *Torrent) hashPiece(piece pieceIndex) (
 	if logPieceContents {
 		writers = append(writers, &examineBuf)
 	}
-	_, err = storagePiece.WriteTo(io.MultiWriter(writers...))
+	var written int64
+	written, err = storagePiece.WriteTo(io.MultiWriter(writers...))
+	if err == nil && written != int64(p.length()) {
+		err = io.ErrShortWrite
+	}
 	if logPieceContents {
 		t.logger.WithDefaultLevel(log.Debug).Printf("hashed %q with copy err %v", examineBuf.Bytes(), err)
 	}
diff --git a/webseed/client.go b/webseed/client.go
index d5ae3ac1..4614a3e4 100644
--- a/webseed/client.go
+++ b/webseed/client.go
@@ -64,7 +64,7 @@ func (me *Client) SetInfo(info *metainfo.Info) {
 		// http://ia600500.us.archive.org/1/items URLs in archive.org torrents.
 		return
 	}
-	me.fileIndex = segments.NewIndex(common.LengthIterFromUpvertedFiles(info.UpvertedFiles()))
+	me.fileIndex = segments.NewIndexFromSegments(common.TorrentOffsetFileSegments(info))
 	me.info = info
 	me.Pieces.AddRange(0, uint64(info.NumPieces()))
 }
-- 
2.44.0