116 lines
2.4 KiB
Go
116 lines
2.4 KiB
Go
package storage
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
"os"
|
|
"path/filepath"
|
|
"sync"
|
|
|
|
"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
|
|
|
|
mu sync.Mutex
|
|
fdMap map[string]*os.File
|
|
}
|
|
|
|
// 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,
|
|
fdMap: make(map[string]*os.File),
|
|
}
|
|
}
|
|
|
|
// Close closes all open file descriptors
|
|
func (pr *PieceReader) Close() {
|
|
pr.mu.Lock()
|
|
defer pr.mu.Unlock()
|
|
for _, f := range pr.fdMap {
|
|
f.Close()
|
|
}
|
|
pr.fdMap = make(map[string]*os.File)
|
|
}
|
|
|
|
func (pr *PieceReader) getFile(path string) (*os.File, error) {
|
|
pr.mu.Lock()
|
|
defer pr.mu.Unlock()
|
|
|
|
if f, ok := pr.fdMap[path]; ok {
|
|
return f, nil
|
|
}
|
|
|
|
f, err := os.Open(path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
pr.fdMap[path] = f
|
|
return f, nil
|
|
}
|
|
|
|
// 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 := pr.getFile(filePath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
n, err := f.ReadAt(buf[bytesRead:bytesRead+readLength], int64(fileOffset))
|
|
if err != nil && err != io.EOF {
|
|
return nil, err
|
|
}
|
|
|
|
bytesRead += n
|
|
if n < readLength {
|
|
break // EOF reached prematurely
|
|
}
|
|
|
|
currentOffset += file.Length
|
|
}
|
|
|
|
if bytesRead < length {
|
|
return nil, errors.New("could not read full block")
|
|
}
|
|
|
|
return buf, nil
|
|
}
|