Ztorrent/internal/storage/reader.go

82 lines
1.8 KiB
Go

package storage
import (
"errors"
"io"
"os"
"path/filepath"
"github.com/veggiedefender/torrent-client/internal/torrentfile"
)
// PieceReader reads pieces from the finalized torrent files on disk.
type PieceReader struct {
files []torrentfile.File
outputRoot string
pieceLength int
totalLength int
}
// NewPieceReader creates a new PieceReader.
func NewPieceReader(files []torrentfile.File, pieceLength, totalLength int, outputRoot string) *PieceReader {
return &PieceReader{
files: files,
outputRoot: outputRoot,
pieceLength: pieceLength,
totalLength: totalLength,
}
}
// ReadBlock reads a specific block of data from the files.
func (pr *PieceReader) ReadBlock(pieceIndex, begin, length int) ([]byte, error) {
absoluteOffset := pieceIndex*pr.pieceLength + begin
if absoluteOffset+length > pr.totalLength {
return nil, errors.New("read block exceeds total length")
}
buf := make([]byte, length)
bytesRead := 0
var currentOffset int
for _, file := range pr.files {
if currentOffset+file.Length <= absoluteOffset {
currentOffset += file.Length
continue
}
if currentOffset >= absoluteOffset+length {
break
}
fileOffset := 0
if absoluteOffset > currentOffset {
fileOffset = absoluteOffset - currentOffset
}
readLength := file.Length - fileOffset
if readLength > length-bytesRead {
readLength = length - bytesRead
}
filePath := filepath.Join(pr.outputRoot, file.Path)
f, err := os.Open(filePath)
if err != nil {
return nil, err
}
_, err = f.ReadAt(buf[bytesRead:bytesRead+readLength], int64(fileOffset))
f.Close()
if err != nil && err != io.EOF {
return nil, err
}
bytesRead += readLength
currentOffset += file.Length
}
if bytesRead < length {
return nil, errors.New("could not read full block")
}
return buf, nil
}