2446 lines
60 KiB
Go
2446 lines
60 KiB
Go
package torrent
|
||
|
||
import (
|
||
"context"
|
||
"crypto/rand"
|
||
"crypto/sha1"
|
||
"encoding/binary"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
mathrand "math/rand"
|
||
"net"
|
||
"os"
|
||
"path/filepath"
|
||
"sort"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
"unicode"
|
||
|
||
"golang.org/x/time/rate"
|
||
|
||
"github.com/veggiedefender/torrent-client/internal/dht"
|
||
"github.com/veggiedefender/torrent-client/internal/logger"
|
||
"github.com/veggiedefender/torrent-client/internal/magnet"
|
||
"github.com/veggiedefender/torrent-client/internal/portforward"
|
||
"github.com/veggiedefender/torrent-client/internal/sorter"
|
||
"github.com/veggiedefender/torrent-client/internal/storage"
|
||
"github.com/veggiedefender/torrent-client/internal/torrentfile"
|
||
"github.com/veggiedefender/torrent-client/internal/tracker"
|
||
)
|
||
|
||
type Engine struct {
|
||
mu sync.RWMutex
|
||
|
||
torrent *torrentfile.TorrentFile
|
||
peers []PeerStatus
|
||
peerIdx map[string]int // key → index in peers slice, O(1) lookup
|
||
trackers []TrackerStatus
|
||
peerID [20]byte
|
||
|
||
dhtServer *dht.Server
|
||
portManager *portforward.PortManager
|
||
|
||
cancel context.CancelFunc
|
||
downloadCtx context.Context
|
||
|
||
incomingConns chan net.Conn
|
||
|
||
lastErr error
|
||
phase string
|
||
|
||
downloadedBytes int64
|
||
uploadedBytes int64
|
||
completedPieces int
|
||
totalPieces int
|
||
outputPath string
|
||
outputRoot string
|
||
downloadSpeed float64
|
||
uploadSpeed float64
|
||
speedAt time.Time
|
||
downloadAtBytes int64
|
||
uploadAtBytes int64
|
||
pieceStates []PieceState
|
||
|
||
// rate limiting
|
||
downloadLimitBps atomic.Int64
|
||
uploadLimitBps atomic.Int64
|
||
|
||
scheduler *pieceScheduler // For toggling sequential mode
|
||
}
|
||
|
||
type PieceState uint8
|
||
|
||
const (
|
||
PieceMissing PieceState = iota
|
||
PieceDownloading
|
||
PieceCompleted
|
||
)
|
||
|
||
type Status struct {
|
||
Loaded bool
|
||
InfoHash [20]byte
|
||
Name string
|
||
Length int
|
||
PieceLength int
|
||
PieceCount int
|
||
|
||
CompletedPieces int
|
||
DownloadedBytes int64
|
||
UploadedBytes int64
|
||
DownloadSpeed float64
|
||
UploadSpeed float64
|
||
TotalBytes int64
|
||
OutputPath string
|
||
|
||
Files []torrentfile.File
|
||
Announce string
|
||
|
||
PeerCount int
|
||
ActivePeers int
|
||
Peers []PeerStatus
|
||
Trackers []TrackerStatus
|
||
|
||
PeerID string
|
||
Progress float64
|
||
Phase string
|
||
LastError string
|
||
PieceStates []PieceState
|
||
|
||
DownloadLimitBps int64
|
||
UploadLimitBps int64
|
||
}
|
||
|
||
type PeerStatus struct {
|
||
Address string
|
||
Port uint16
|
||
Source string
|
||
State string
|
||
Error string
|
||
DownloadedPieces int
|
||
DownloadedBytes int64
|
||
UploadedBytes int64
|
||
DownloadSpeed float64
|
||
LatencyMS int
|
||
Score float64
|
||
ErrorCount int
|
||
}
|
||
|
||
type TrackerStatus struct {
|
||
URL string
|
||
State string
|
||
PeerCount int
|
||
Error string
|
||
}
|
||
|
||
const (
|
||
trackerPort = 6881
|
||
trackerNumWant = 320
|
||
trackerAnnounceTimeout = 5 * time.Second
|
||
downloadStallTimeout = 90 * time.Second
|
||
reannounceInterval = 15 * time.Second
|
||
fastReannounceInterval = 6 * time.Second
|
||
idleRetryDelay = 900 * time.Millisecond
|
||
maxPeerConnections = 60
|
||
)
|
||
|
||
type pieceWriteRequest struct {
|
||
pieceIndex int
|
||
pieceData []byte
|
||
resultCh chan error
|
||
}
|
||
|
||
type bufferedPieceWriter struct {
|
||
requests chan pieceWriteRequest
|
||
done chan struct{}
|
||
closeOnce sync.Once
|
||
|
||
errMu sync.RWMutex
|
||
closeErr error
|
||
}
|
||
|
||
func newBufferedPieceWriter(partFile *os.File, pieceLength int, batchSize int, flushInterval time.Duration) *bufferedPieceWriter {
|
||
if batchSize < 1 {
|
||
batchSize = 1
|
||
}
|
||
if flushInterval <= 0 {
|
||
flushInterval = 120 * time.Millisecond
|
||
}
|
||
|
||
writer := &bufferedPieceWriter{
|
||
requests: make(chan pieceWriteRequest, batchSize*4),
|
||
done: make(chan struct{}),
|
||
}
|
||
|
||
go writer.loop(partFile, pieceLength, batchSize, flushInterval)
|
||
return writer
|
||
}
|
||
|
||
func (w *bufferedPieceWriter) WritePiece(ctx context.Context, pieceIndex int, pieceData []byte) error {
|
||
req := pieceWriteRequest{
|
||
pieceIndex: pieceIndex,
|
||
pieceData: pieceData,
|
||
resultCh: make(chan error, 1),
|
||
}
|
||
|
||
select {
|
||
case <-ctx.Done():
|
||
return ctx.Err()
|
||
case <-w.done:
|
||
if err := w.err(); err != nil {
|
||
return err
|
||
}
|
||
return errors.New("piece writer stopped")
|
||
case w.requests <- req:
|
||
}
|
||
|
||
select {
|
||
case <-ctx.Done():
|
||
return ctx.Err()
|
||
case <-w.done:
|
||
if err := w.err(); err != nil {
|
||
return err
|
||
}
|
||
return errors.New("piece writer stopped")
|
||
case err := <-req.resultCh:
|
||
return err
|
||
}
|
||
}
|
||
|
||
func (w *bufferedPieceWriter) Close() error {
|
||
w.closeOnce.Do(func() {
|
||
close(w.requests)
|
||
<-w.done
|
||
})
|
||
return w.err()
|
||
}
|
||
|
||
func (w *bufferedPieceWriter) loop(partFile *os.File, pieceLength int, batchSize int, flushInterval time.Duration) {
|
||
defer close(w.done)
|
||
flushTicker := time.NewTicker(flushInterval)
|
||
defer flushTicker.Stop()
|
||
|
||
batch := make([]pieceWriteRequest, 0, batchSize)
|
||
var terminalErr error
|
||
|
||
failRequest := func(req pieceWriteRequest, err error) {
|
||
req.resultCh <- err
|
||
close(req.resultCh)
|
||
}
|
||
|
||
flushBatch := func() {
|
||
if len(batch) == 0 {
|
||
return
|
||
}
|
||
var writeErr error
|
||
for i, req := range batch {
|
||
if writeErr == nil {
|
||
offset := int64(req.pieceIndex * pieceLength)
|
||
_, writeErr = partFile.WriteAt(req.pieceData, offset)
|
||
}
|
||
failRequest(req, writeErr)
|
||
if writeErr != nil {
|
||
for j := i + 1; j < len(batch); j++ {
|
||
failRequest(batch[j], writeErr)
|
||
}
|
||
break
|
||
}
|
||
}
|
||
batch = batch[:0]
|
||
if writeErr != nil {
|
||
terminalErr = writeErr
|
||
}
|
||
}
|
||
|
||
for {
|
||
select {
|
||
case req, ok := <-w.requests:
|
||
if !ok {
|
||
flushBatch()
|
||
w.setErr(terminalErr)
|
||
return
|
||
}
|
||
|
||
if terminalErr != nil {
|
||
failRequest(req, terminalErr)
|
||
continue
|
||
}
|
||
|
||
batch = append(batch, req)
|
||
if len(batch) >= batchSize {
|
||
flushBatch()
|
||
}
|
||
case <-flushTicker.C:
|
||
if terminalErr != nil {
|
||
continue
|
||
}
|
||
flushBatch()
|
||
}
|
||
}
|
||
}
|
||
|
||
func (w *bufferedPieceWriter) setErr(err error) {
|
||
w.errMu.Lock()
|
||
defer w.errMu.Unlock()
|
||
w.closeErr = err
|
||
}
|
||
|
||
func (w *bufferedPieceWriter) err() error {
|
||
w.errMu.RLock()
|
||
defer w.errMu.RUnlock()
|
||
return w.closeErr
|
||
}
|
||
|
||
func NewEngine() *Engine {
|
||
return &Engine{
|
||
peerID: generatePeerID(),
|
||
phase: "idle",
|
||
incomingConns: make(chan net.Conn, 128),
|
||
portManager: portforward.NewPortManager(),
|
||
}
|
||
}
|
||
|
||
// SetDownloadLimit устанавливает лимит скачивания в байт/с (0 = без лимита).
|
||
func (e *Engine) SetDownloadLimit(bps int64) {
|
||
if bps < 0 {
|
||
bps = 0
|
||
}
|
||
e.downloadLimitBps.Store(bps)
|
||
}
|
||
|
||
func (e *Engine) addUploadedBytes(n int64) {
|
||
e.mu.Lock()
|
||
e.uploadedBytes += n
|
||
e.mu.Unlock()
|
||
}
|
||
|
||
// SetUploadLimit устанавливает лимит отдачи в байт/с (0 = без лимита).
|
||
func (e *Engine) SetUploadLimit(bps int64) {
|
||
if bps < 0 {
|
||
bps = 0
|
||
}
|
||
e.uploadLimitBps.Store(bps)
|
||
}
|
||
|
||
// downloadRateLimiter возвращает rate.Limiter для текущего лимита скачивания.
|
||
func (e *Engine) downloadRateLimiter() *rate.Limiter {
|
||
bps := e.downloadLimitBps.Load()
|
||
if bps <= 0 {
|
||
return rate.NewLimiter(rate.Inf, 0)
|
||
}
|
||
return rate.NewLimiter(rate.Limit(bps), int(bps))
|
||
}
|
||
|
||
func (e *Engine) LoadTorrent(path, outputRoot string) error {
|
||
e.Stop()
|
||
e.resetStateForNewLoad()
|
||
e.setPhase("loading_metadata")
|
||
outputRoot = normalizeOutputRoot(outputRoot)
|
||
logger.Info("ENGINE", "loading torrent metadata from %s", path)
|
||
debugf("download root selected: %s", outputRoot)
|
||
|
||
downloadCtx, cancel := context.WithCancel(context.Background())
|
||
e.swapCancel(cancel)
|
||
e.mu.Lock()
|
||
e.downloadCtx = downloadCtx
|
||
e.mu.Unlock()
|
||
|
||
if strings.HasPrefix(path, "magnet:") {
|
||
go e.loadMagnet(downloadCtx, path, outputRoot)
|
||
return nil
|
||
}
|
||
|
||
tf, err := torrentfile.Open(path)
|
||
if err != nil {
|
||
logger.Info("ENGINE", "failed to parse torrent %s: %v", path, err)
|
||
e.setError(err)
|
||
e.setPhase("failed")
|
||
cancel()
|
||
return err
|
||
}
|
||
|
||
e.mu.Lock()
|
||
e.torrent = tf
|
||
if tf.InfoHash != [20]byte{} && len(tf.PieceHashes) > 0 {
|
||
e.totalPieces = len(tf.PieceHashes)
|
||
e.pieceStates = make([]PieceState, len(tf.PieceHashes))
|
||
e.outputPath = buildOutputPath(tf, outputRoot)
|
||
}
|
||
e.phase = "stopped"
|
||
e.mu.Unlock()
|
||
|
||
return nil
|
||
}
|
||
|
||
func (e *Engine) StartDownload() {
|
||
e.mu.Lock()
|
||
tf := e.torrent
|
||
out := e.outputRoot
|
||
phase := e.phase
|
||
|
||
if tf == nil || phase != "stopped" && phase != "failed" {
|
||
e.mu.Unlock()
|
||
return
|
||
}
|
||
|
||
downloadCtx, cancel := context.WithCancel(context.Background())
|
||
e.downloadCtx = downloadCtx
|
||
e.cancel = cancel
|
||
e.mu.Unlock()
|
||
|
||
go e.startTorrent(downloadCtx, tf, out)
|
||
}
|
||
|
||
func (e *Engine) loadMagnet(ctx context.Context, uri, outputRoot string) {
|
||
ml, err := magnet.Parse(uri)
|
||
if err != nil {
|
||
e.setError(err)
|
||
e.setPhase("failed")
|
||
return
|
||
}
|
||
|
||
e.setPhase("downloading_metadata")
|
||
|
||
queryCtx, queryCancel := context.WithTimeout(ctx, 30*time.Second)
|
||
defer queryCancel()
|
||
|
||
peers, _ := e.queryTrackers(queryCtx, ml.Trackers, ml.InfoHash, 0, tracker.AnnounceOptions{
|
||
PeerID: e.peerID,
|
||
Port: trackerPort,
|
||
NumWant: trackerNumWant,
|
||
Timeout: trackerAnnounceTimeout,
|
||
})
|
||
|
||
e.ensureDHTServer(ctx)
|
||
peersCh := make(chan []PeerStatus, 10)
|
||
if e.dhtServer != nil {
|
||
go func() {
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
case dhtPeers, ok := <-e.dhtServer.PeersFound:
|
||
if !ok {
|
||
return
|
||
}
|
||
var ps []PeerStatus
|
||
for _, p := range dhtPeers {
|
||
ps = append(ps, PeerStatus{Address: p.IP.String(), Port: p.Port})
|
||
}
|
||
select {
|
||
case peersCh <- ps:
|
||
default:
|
||
}
|
||
}
|
||
}
|
||
}()
|
||
go e.dhtServer.SearchForPeers(ctx, ml.InfoHash)
|
||
}
|
||
|
||
tf, err := e.downloadMetadataFromPeers(ctx, ml.InfoHash, ml.Name, peersCh, peers)
|
||
if err != nil {
|
||
e.setError(err)
|
||
e.setPhase("failed")
|
||
return
|
||
}
|
||
|
||
tf.Trackers = ml.Trackers
|
||
tf.Name = ml.Name
|
||
if tf.Name == "" {
|
||
tf.Name = hex.EncodeToString(ml.InfoHash[:])
|
||
}
|
||
|
||
e.mu.Lock()
|
||
e.torrent = tf
|
||
e.totalPieces = len(tf.PieceHashes)
|
||
e.pieceStates = make([]PieceState, len(tf.PieceHashes))
|
||
e.outputPath = buildOutputPath(tf, outputRoot)
|
||
e.phase = "stopped"
|
||
e.mu.Unlock()
|
||
}
|
||
|
||
func (e *Engine) downloadMetadataFromPeers(ctx context.Context, infoHash [20]byte, name string, peersCh <-chan []PeerStatus, initialPeers []PeerStatus) (*torrentfile.TorrentFile, error) {
|
||
resultCh := make(chan *torrentfile.TorrentFile, 1)
|
||
|
||
workCtx, workCancel := context.WithCancel(ctx)
|
||
defer workCancel()
|
||
|
||
tryPeer := func(p PeerStatus) {
|
||
addr := fmt.Sprintf("%s:%d", p.Address, p.Port)
|
||
|
||
dialCtx, cancel := context.WithTimeout(workCtx, 10*time.Second)
|
||
pc, err := newPeerClient(dialCtx, addr, infoHash, e.peerID, 0)
|
||
cancel()
|
||
if err != nil {
|
||
logger.Info("ENGINE", "peer dial failed %s: %v", addr, err)
|
||
return
|
||
}
|
||
defer pc.Close()
|
||
|
||
if pc.PeerUtMetadataID() == 0 {
|
||
logger.Info("ENGINE", "peer %s: no ut_metadata support", addr)
|
||
return
|
||
}
|
||
|
||
size := pc.MetadataSize()
|
||
logger.Info("ENGINE", "peer %s: metadata size=%d, ut_metadata=%d", addr, size, pc.PeerUtMetadataID())
|
||
if size <= 0 || size > 10*1024*1024 {
|
||
return
|
||
}
|
||
|
||
numPieces := (size + 16383) / 16384
|
||
metadata := make([]byte, size)
|
||
|
||
for i := 0; i < numPieces; i++ {
|
||
if err := pc.SendMetadataRequest(workCtx, i); err != nil {
|
||
return
|
||
}
|
||
pieceIdx, data, reject, err := pc.ReadMetadataMessage(workCtx)
|
||
if err != nil || reject || pieceIdx != i {
|
||
return
|
||
}
|
||
offset := i * 16384
|
||
copy(metadata[offset:], data)
|
||
}
|
||
|
||
if sha1.Sum(metadata) == infoHash {
|
||
if tf, err := torrentfile.FromMetadata(metadata, nil); err == nil {
|
||
select {
|
||
case resultCh <- tf:
|
||
default:
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
for _, p := range initialPeers {
|
||
go tryPeer(p)
|
||
}
|
||
|
||
timeout := time.After(60 * time.Second) // 60 seconds should be enough for DHT bootstrapping + metadata download
|
||
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
return nil, ctx.Err()
|
||
case <-timeout:
|
||
return nil, errors.New("timeout waiting for valid metadata from peers")
|
||
case tf := <-resultCh:
|
||
return tf, nil
|
||
case peers, ok := <-peersCh:
|
||
if !ok {
|
||
peersCh = nil
|
||
} else {
|
||
for _, p := range peers {
|
||
go tryPeer(p)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
func (e *Engine) startTorrent(downloadCtx context.Context, tf *torrentfile.TorrentFile, outputRoot string) {
|
||
queryCtx, queryCancel := context.WithTimeout(downloadCtx, 30*time.Second)
|
||
defer queryCancel()
|
||
|
||
e.setPhase("querying_trackers")
|
||
peers, trackerStatuses := e.queryTrackers(queryCtx, collectTrackersToTry(tf), tf.InfoHash, tf.Length, tracker.AnnounceOptions{
|
||
PeerID: e.peerID,
|
||
Port: trackerPort,
|
||
NumWant: trackerNumWant,
|
||
Timeout: trackerAnnounceTimeout,
|
||
})
|
||
|
||
e.mu.Lock()
|
||
e.torrent = tf
|
||
e.trackers = trackerStatuses
|
||
e.peers = peers
|
||
e.peerIdx = buildPeerIdx(peers)
|
||
e.totalPieces = len(tf.PieceHashes)
|
||
e.downloadedBytes = 0
|
||
e.uploadedBytes = 0
|
||
e.completedPieces = 0
|
||
e.outputRoot = outputRoot
|
||
e.outputPath = buildOutputPath(tf, outputRoot)
|
||
e.downloadSpeed = 0
|
||
e.uploadSpeed = 0
|
||
e.speedAt = time.Now()
|
||
e.downloadAtBytes = 0
|
||
e.uploadAtBytes = 0
|
||
e.pieceStates = make([]PieceState, len(tf.PieceHashes))
|
||
if len(peers) == 0 {
|
||
warnErr := fmt.Errorf("no peers yet via %d trackers", len(trackerStatuses))
|
||
if trErr := firstTrackerError(trackerStatuses); trErr != nil {
|
||
warnErr = fmt.Errorf("no peers yet: %w", trErr)
|
||
}
|
||
e.lastErr = warnErr
|
||
} else {
|
||
e.lastErr = nil
|
||
}
|
||
e.phase = "ready"
|
||
e.mu.Unlock()
|
||
logger.Info("ENGINE", "torrent loaded: name=%s size=%d peers=%d output=%s", tf.Name, tf.Length, len(peers), e.outputPath)
|
||
|
||
go e.runDownload(downloadCtx, tf)
|
||
}
|
||
|
||
func (e *Engine) Stop() {
|
||
e.mu.Lock()
|
||
cancel := e.cancel
|
||
e.cancel = nil
|
||
if cancel != nil {
|
||
e.phase = "stopped"
|
||
}
|
||
e.mu.Unlock()
|
||
|
||
if cancel != nil {
|
||
cancel()
|
||
}
|
||
}
|
||
|
||
func (e *Engine) Progress() float64 {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
return e.progressUnsafe()
|
||
}
|
||
|
||
func (e *Engine) Status() Status {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
|
||
status := Status{
|
||
Loaded: e.torrent != nil,
|
||
PeerCount: len(e.peers),
|
||
ActivePeers: countActivePeers(e.peers),
|
||
Peers: append([]PeerStatus(nil), e.peers...),
|
||
Trackers: append([]TrackerStatus(nil), e.trackers...),
|
||
PeerID: string(e.peerID[:]),
|
||
Progress: e.progressUnsafe(),
|
||
Phase: e.phase,
|
||
CompletedPieces: e.completedPieces,
|
||
DownloadedBytes: e.downloadedBytes,
|
||
UploadedBytes: e.uploadedBytes,
|
||
DownloadSpeed: e.downloadSpeed,
|
||
UploadSpeed: e.uploadSpeed,
|
||
OutputPath: e.outputPath,
|
||
PieceStates: append([]PieceState(nil), e.pieceStates...),
|
||
DownloadLimitBps: e.downloadLimitBps.Load(),
|
||
UploadLimitBps: e.uploadLimitBps.Load(),
|
||
}
|
||
|
||
if e.torrent != nil {
|
||
status.InfoHash = e.torrent.InfoHash
|
||
status.Loaded = true
|
||
status.Name = e.torrent.Name
|
||
status.Length = e.torrent.Length
|
||
status.TotalBytes = int64(e.torrent.Length)
|
||
status.PieceLength = e.torrent.PieceLength
|
||
status.PieceCount = len(e.torrent.PieceHashes)
|
||
|
||
status.Files = make([]torrentfile.File, len(e.torrent.Files))
|
||
copy(status.Files, e.torrent.Files)
|
||
|
||
status.Announce = e.torrent.Announce
|
||
}
|
||
if e.lastErr != nil {
|
||
status.LastError = e.lastErr.Error()
|
||
}
|
||
|
||
return status
|
||
}
|
||
|
||
func (e *Engine) queryTrackers(ctx context.Context, trackersToTry []string, infoHash [20]byte, length int, opts tracker.AnnounceOptions) ([]PeerStatus, []TrackerStatus) {
|
||
peerMap := make(map[string]PeerStatus, 512)
|
||
statuses := make([]TrackerStatus, 0, len(trackersToTry))
|
||
|
||
type trackerResult struct {
|
||
announce string
|
||
peers []tracker.Peer
|
||
err error
|
||
elapsed time.Duration
|
||
}
|
||
|
||
resultsCh := make(chan trackerResult, len(trackersToTry))
|
||
var wg sync.WaitGroup
|
||
for _, announce := range trackersToTry {
|
||
announce := announce
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
started := time.Now()
|
||
logger.Info("ENGINE", "tracker request started: %s", announce)
|
||
|
||
perTrackerCtx, cancel := context.WithTimeout(ctx, opts.Timeout)
|
||
peers, err := tracker.GetPeersFromURL(perTrackerCtx, announce, infoHash, length, opts)
|
||
cancel()
|
||
|
||
resultsCh <- trackerResult{
|
||
announce: announce,
|
||
peers: peers,
|
||
err: err,
|
||
elapsed: time.Since(started),
|
||
}
|
||
}()
|
||
}
|
||
go func() {
|
||
wg.Wait()
|
||
close(resultsCh)
|
||
}()
|
||
|
||
for res := range resultsCh {
|
||
st := TrackerStatus{URL: res.announce}
|
||
if res.err != nil {
|
||
st.State = "error"
|
||
st.Error = res.err.Error()
|
||
logger.Info("ENGINE", "tracker response received: %s state=error elapsed=%s err=%v", res.announce, res.elapsed.Round(time.Millisecond), res.err)
|
||
statuses = append(statuses, st)
|
||
continue
|
||
}
|
||
|
||
if len(res.peers) == 0 {
|
||
st.State = "empty"
|
||
} else {
|
||
st.State = "ok"
|
||
st.PeerCount = len(res.peers)
|
||
}
|
||
statuses = append(statuses, st)
|
||
logger.Info("ENGINE", "tracker response received: %s state=%s peers=%d elapsed=%s", res.announce, st.State, len(res.peers), res.elapsed.Round(time.Millisecond))
|
||
|
||
for _, peer := range res.peers {
|
||
addPeer(peerMap, peer, res.announce)
|
||
}
|
||
}
|
||
sort.Slice(statuses, func(i, j int) bool {
|
||
return statuses[i].URL < statuses[j].URL
|
||
})
|
||
|
||
peerList := make([]PeerStatus, 0, len(peerMap))
|
||
for _, peer := range peerMap {
|
||
peerList = append(peerList, peer)
|
||
}
|
||
sort.Slice(peerList, func(i, j int) bool {
|
||
if peerList[i].Address == peerList[j].Address {
|
||
return peerList[i].Port < peerList[j].Port
|
||
}
|
||
return peerList[i].Address < peerList[j].Address
|
||
})
|
||
|
||
return peerList, statuses
|
||
}
|
||
|
||
func (e *Engine) ensureDHTServer(ctx context.Context) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
if e.dhtServer != nil {
|
||
return
|
||
}
|
||
srv := dht.NewServer()
|
||
if err := srv.Start(ctx, dht.Port); err == nil {
|
||
e.dhtServer = srv
|
||
} else {
|
||
logger.Info("ENGINE", "Failed to start DHT server: %v", err)
|
||
}
|
||
}
|
||
|
||
func (e *Engine) runDownload(ctx context.Context, tf *torrentfile.TorrentFile) {
|
||
if e.dhtServer == nil {
|
||
e.ensureDHTServer(ctx)
|
||
if e.dhtServer != nil {
|
||
go func() {
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
case peers := <-e.dhtServer.PeersFound:
|
||
added := e.AddDiscoveredPeers(peers, "dht")
|
||
if added > 0 {
|
||
debugf("DHT: received %d new peers", added)
|
||
}
|
||
}
|
||
}
|
||
}()
|
||
go e.dhtServer.SearchForPeers(ctx, tf.InfoHash)
|
||
}
|
||
} else {
|
||
go e.dhtServer.SearchForPeers(ctx, tf.InfoHash)
|
||
}
|
||
|
||
defer e.clearCancel()
|
||
|
||
go e.listenIncoming(ctx, tf)
|
||
|
||
e.setPhase("preparing_download")
|
||
logger.Info("ENGINE", "preparing download to %s", e.outputRoot)
|
||
|
||
partPath, partFile, resumedPieces, err := createOrOpenPartFile(tf, e.outputRoot)
|
||
if err != nil {
|
||
logger.Info("ENGINE", "failed to create part file: %v", err)
|
||
e.setTerminalError("failed", fmt.Errorf("create temp file: %w", err))
|
||
return
|
||
}
|
||
e.mu.Lock()
|
||
e.outputPath = partPath
|
||
e.mu.Unlock()
|
||
defer partFile.Close()
|
||
|
||
// Определяем куски, которые нужно пропустить (состоят только из файлов с Priority == 0)
|
||
skippedPieces := make(map[int]bool)
|
||
{
|
||
fileOffset := int64(0)
|
||
fileSkippedRanges := make([][2]int64, 0)
|
||
for _, f := range tf.Files {
|
||
if f.Priority == 0 {
|
||
fileSkippedRanges = append(fileSkippedRanges, [2]int64{fileOffset, fileOffset + int64(f.Length)})
|
||
}
|
||
fileOffset += int64(f.Length)
|
||
}
|
||
|
||
for i := 0; i < len(tf.PieceHashes); i++ {
|
||
pStart := int64(i * tf.PieceLength)
|
||
pEnd := pStart + int64(tf.PieceLength)
|
||
if pEnd > int64(tf.Length) {
|
||
pEnd = int64(tf.Length)
|
||
}
|
||
|
||
// Проверяем, покрыт ли кусок полностью skipped файлами
|
||
covered := int64(0)
|
||
for _, r := range fileSkippedRanges {
|
||
// пересечение [pStart, pEnd) и [r[0], r[1])
|
||
start := pStart
|
||
if r[0] > start {
|
||
start = r[0]
|
||
}
|
||
end := pEnd
|
||
if r[1] < end {
|
||
end = r[1]
|
||
}
|
||
if start < end {
|
||
covered += (end - start)
|
||
}
|
||
}
|
||
|
||
if covered >= (pEnd - pStart) {
|
||
skippedPieces[i] = true
|
||
resumedPieces = append(resumedPieces, i)
|
||
}
|
||
}
|
||
}
|
||
|
||
// Восстанавливаем состояние уже скачанных кусков (включая пропущенные)
|
||
if len(resumedPieces) > 0 {
|
||
e.mu.Lock()
|
||
for _, idx := range resumedPieces {
|
||
if idx >= 0 && idx < len(e.pieceStates) {
|
||
e.pieceStates[idx] = PieceCompleted
|
||
}
|
||
}
|
||
e.completedPieces = len(resumedPieces)
|
||
e.downloadedBytes = int64(len(resumedPieces)) * int64(tf.PieceLength)
|
||
if e.downloadedBytes > int64(tf.Length) {
|
||
e.downloadedBytes = int64(tf.Length)
|
||
}
|
||
e.mu.Unlock()
|
||
logger.Info("ENGINE", "resumed download: %d/%d pieces already done", len(resumedPieces), len(tf.PieceHashes))
|
||
}
|
||
|
||
pieceWriter := newBufferedPieceWriter(partFile, tf.PieceLength, 8, 140*time.Millisecond)
|
||
defer pieceWriter.Close()
|
||
|
||
scheduler := newPieceSchedulerWithResume(len(tf.PieceHashes), resumedPieces)
|
||
e.mu.Lock()
|
||
e.scheduler = scheduler
|
||
e.mu.Unlock()
|
||
defer func() {
|
||
e.mu.Lock()
|
||
e.scheduler = nil
|
||
e.mu.Unlock()
|
||
scheduler.Stop()
|
||
}()
|
||
workerCtx, workerCancel := context.WithCancel(ctx)
|
||
var workersWG sync.WaitGroup
|
||
|
||
startedWorkers := make(map[string]struct{})
|
||
workerDoneCh := make(chan string, 512)
|
||
peerJobs := make(chan PeerStatus, maxPeerConnections*4)
|
||
|
||
for i := 0; i < maxPeerConnections; i++ {
|
||
workersWG.Add(1)
|
||
go func() {
|
||
defer workersWG.Done()
|
||
for {
|
||
select {
|
||
case <-workerCtx.Done():
|
||
return
|
||
case peer, ok := <-peerJobs:
|
||
if !ok {
|
||
return
|
||
}
|
||
e.runPeerWorker(workerCtx, tf, peer, pieceWriter, scheduler)
|
||
workerKey := peerStatusKey(peer.Address, peer.Port)
|
||
select {
|
||
case workerDoneCh <- workerKey:
|
||
default:
|
||
}
|
||
}
|
||
}
|
||
}()
|
||
}
|
||
|
||
startWorkers := func(peers []PeerStatus) {
|
||
for _, peer := range peers {
|
||
key := peerStatusKey(peer.Address, peer.Port)
|
||
if _, exists := startedWorkers[key]; exists {
|
||
continue
|
||
}
|
||
startedWorkers[key] = struct{}{}
|
||
select {
|
||
case <-workerCtx.Done():
|
||
return
|
||
case peerJobs <- peer:
|
||
}
|
||
}
|
||
}
|
||
|
||
var cleanupOnce sync.Once
|
||
cleanupWorkers := func() {
|
||
cleanupOnce.Do(func() {
|
||
workerCancel()
|
||
close(peerJobs)
|
||
scheduler.Stop()
|
||
workersWG.Wait()
|
||
})
|
||
}
|
||
defer cleanupWorkers()
|
||
|
||
e.setPhase("downloading")
|
||
logger.Info("ENGINE", "download started: pieces=%d peers=%d", len(tf.PieceHashes), len(e.snapshotPeers()))
|
||
lastProgressAt := time.Now()
|
||
lastReannounceAt := time.Now()
|
||
startWorkers(e.snapshotPeers())
|
||
heartbeatTicker := time.NewTicker(idleRetryDelay)
|
||
defer heartbeatTicker.Stop()
|
||
downloadCompleted := false
|
||
|
||
for !downloadCompleted {
|
||
select {
|
||
case <-ctx.Done():
|
||
e.setPhase("stopped")
|
||
e.setError(nil)
|
||
return
|
||
case _, ok := <-scheduler.Progress():
|
||
if !ok {
|
||
downloadCompleted = true
|
||
continue
|
||
}
|
||
lastProgressAt = time.Now()
|
||
debugf("piece completed: %d/%d", e.currentCompletedPieces(), len(tf.PieceHashes))
|
||
case <-scheduler.Done():
|
||
downloadCompleted = true
|
||
case workerKey := <-workerDoneCh:
|
||
delete(startedWorkers, workerKey)
|
||
debugf("worker finished for peer %s", workerKey)
|
||
case <-heartbeatTicker.C:
|
||
e.updateSpeed(time.Now())
|
||
// Сохраняем bitmap периодически
|
||
e.asyncSaveBitmap(partPath + ".bitmap")
|
||
nextReannounce := reannounceInterval
|
||
if e.currentActivePeers() < 4 {
|
||
nextReannounce = fastReannounceInterval
|
||
}
|
||
if time.Since(lastReannounceAt) >= nextReannounce {
|
||
lastReannounceAt = time.Now()
|
||
// Запускаем reannounce в фоне — не блокируем heartbeat
|
||
go func() {
|
||
e.setPhase("querying_trackers")
|
||
refreshCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
|
||
added := e.refreshPeersFromTrackers(refreshCtx, tf)
|
||
cancel()
|
||
e.setPhase("downloading")
|
||
if added > 0 {
|
||
debugf("reannounce complete, added peers=%d total=%d", added, len(e.snapshotPeers()))
|
||
}
|
||
}()
|
||
}
|
||
startWorkers(e.snapshotPeers())
|
||
|
||
if time.Since(lastProgressAt) >= downloadStallTimeout {
|
||
logger.Info("ENGINE", "download stalled: no progress for %s", downloadStallTimeout.Round(time.Second))
|
||
e.setTerminalError("stalled", fmt.Errorf("download stalled: no piece progress for %s", downloadStallTimeout.Round(time.Second)))
|
||
return
|
||
}
|
||
}
|
||
}
|
||
|
||
cleanupWorkers()
|
||
if err := pieceWriter.Close(); err != nil {
|
||
logger.Info("ENGINE", "failed to flush piece writer: %v", err)
|
||
e.setTerminalError("failed", fmt.Errorf("flush buffered pieces: %w", err))
|
||
return
|
||
}
|
||
|
||
e.setPhase("writing_files")
|
||
if err := materializeDownloadedFiles(tf, partPath, e.outputRoot); err != nil {
|
||
logger.Info("ENGINE", "failed to materialize files: %v", err)
|
||
e.setTerminalError("failed", fmt.Errorf("write output files: %w", err))
|
||
return
|
||
}
|
||
|
||
// Smart Sorter
|
||
e.setPhase("sorting_files")
|
||
if newRoot, err := sorter.SortAndMove(e.outputRoot, tf.Name); err != nil {
|
||
logger.Warn("ENGINE", "Smart Sorter failed to move files: %v", err)
|
||
} else if newRoot != "" {
|
||
e.mu.Lock()
|
||
e.outputRoot = newRoot
|
||
e.mu.Unlock()
|
||
}
|
||
|
||
if err := os.Remove(partPath); err != nil && !errors.Is(err, os.ErrNotExist) {
|
||
logger.Info("ENGINE", "failed to remove part file %s: %v", partPath, err)
|
||
}
|
||
// Удаляем bitmap вместе с .part
|
||
_ = os.Remove(partPath + ".bitmap")
|
||
|
||
e.mu.Lock()
|
||
e.downloadedBytes = int64(tf.Length)
|
||
e.completedPieces = len(tf.PieceHashes)
|
||
if len(e.pieceStates) != len(tf.PieceHashes) {
|
||
e.pieceStates = make([]PieceState, len(tf.PieceHashes))
|
||
}
|
||
for i := range e.pieceStates {
|
||
e.pieceStates[i] = PieceCompleted
|
||
}
|
||
e.phase = "seeding"
|
||
e.lastErr = nil
|
||
e.downloadSpeed = 0
|
||
e.uploadSpeed = 0
|
||
e.mu.Unlock()
|
||
logger.Info("ENGINE", "download completed, transitioning to seeding: %s", e.outputPath)
|
||
|
||
e.runSeeding(ctx, tf)
|
||
}
|
||
|
||
func (e *Engine) runSeeding(ctx context.Context, tf *torrentfile.TorrentFile) {
|
||
logger.Info("ENGINE", "entering seeding phase")
|
||
|
||
pieceReader := storage.NewPieceReader(tf.Files, tf.PieceLength, tf.Length, e.outputRoot)
|
||
defer pieceReader.Close()
|
||
|
||
workerCtx, workerCancel := context.WithCancel(ctx)
|
||
defer workerCancel()
|
||
|
||
var workersWG sync.WaitGroup
|
||
startedWorkers := make(map[string]struct{})
|
||
workerDoneCh := make(chan string, 512)
|
||
peerJobs := make(chan PeerStatus, maxPeerConnections*4)
|
||
|
||
for i := 0; i < maxPeerConnections; i++ {
|
||
workersWG.Add(1)
|
||
go func() {
|
||
defer workersWG.Done()
|
||
for {
|
||
select {
|
||
case <-workerCtx.Done():
|
||
return
|
||
case peer, ok := <-peerJobs:
|
||
if !ok {
|
||
return
|
||
}
|
||
e.runSeedingWorker(workerCtx, tf, peer, pieceReader)
|
||
workerKey := peerStatusKey(peer.Address, peer.Port)
|
||
select {
|
||
case workerDoneCh <- workerKey:
|
||
default:
|
||
}
|
||
}
|
||
}
|
||
}()
|
||
}
|
||
|
||
startWorkers := func(peers []PeerStatus) {
|
||
for _, peer := range peers {
|
||
key := peerStatusKey(peer.Address, peer.Port)
|
||
if _, exists := startedWorkers[key]; exists {
|
||
continue
|
||
}
|
||
startedWorkers[key] = struct{}{}
|
||
select {
|
||
case <-workerCtx.Done():
|
||
return
|
||
case peerJobs <- peer:
|
||
}
|
||
}
|
||
}
|
||
|
||
heartbeatTicker := time.NewTicker(idleRetryDelay)
|
||
defer heartbeatTicker.Stop()
|
||
|
||
startWorkers(e.snapshotPeers())
|
||
lastReannounceAt := time.Now()
|
||
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
e.setPhase("stopped")
|
||
e.setError(nil)
|
||
return
|
||
case workerKey := <-workerDoneCh:
|
||
delete(startedWorkers, workerKey)
|
||
case incomingConn := <-e.incomingConns:
|
||
e.handleIncomingSeeding(workerCtx, tf, incomingConn, pieceReader)
|
||
case <-heartbeatTicker.C:
|
||
e.updateSpeed(time.Now())
|
||
if time.Since(lastReannounceAt) >= reannounceInterval {
|
||
lastReannounceAt = time.Now()
|
||
go func() {
|
||
refreshCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
|
||
e.refreshPeersFromTrackers(refreshCtx, tf)
|
||
cancel()
|
||
}()
|
||
}
|
||
startWorkers(e.snapshotPeers())
|
||
}
|
||
}
|
||
}
|
||
|
||
func (e *Engine) runSeedingWorker(ctx context.Context, tf *torrentfile.TorrentFile, peer PeerStatus, reader *storage.PieceReader) {
|
||
addr := fmt.Sprintf("%s:%d", peer.Address, peer.Port)
|
||
debugf("seeding worker dialing %s", addr)
|
||
|
||
dialCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||
pc, err := newPeerClient(dialCtx, addr, tf.InfoHash, e.peerID, len(tf.PieceHashes))
|
||
cancel()
|
||
if err != nil {
|
||
e.setPeerError(peer.Address, peer.Port, err.Error())
|
||
return
|
||
}
|
||
defer pc.Close()
|
||
|
||
pc.OnPex = func(peers []tracker.Peer) {
|
||
added := e.AddDiscoveredPeers(peers, "pex")
|
||
if added > 0 {
|
||
debugf("PEX: received %d new peers from %s", added, addr)
|
||
}
|
||
}
|
||
|
||
workerCtx, cancelWorker := context.WithCancel(ctx)
|
||
defer cancelWorker()
|
||
|
||
go func() {
|
||
ticker := time.NewTicker(60 * time.Second)
|
||
defer ticker.Stop()
|
||
for {
|
||
select {
|
||
case <-workerCtx.Done():
|
||
return
|
||
case <-ticker.C:
|
||
if sample := e.GetActivePeersSample(50); len(sample) > 0 {
|
||
_ = pc.SendPex(workerCtx, sample)
|
||
}
|
||
}
|
||
}
|
||
}()
|
||
|
||
e.setPeerState(peer.Address, peer.Port, "seeding", "")
|
||
|
||
bitfield := make([]byte, (len(tf.PieceHashes)+7)/8)
|
||
for i := range bitfield {
|
||
bitfield[i] = 0xFF
|
||
}
|
||
lastPieces := len(tf.PieceHashes) % 8
|
||
if lastPieces != 0 {
|
||
mask := byte(0xFF << (8 - lastPieces))
|
||
bitfield[len(bitfield)-1] &= mask
|
||
}
|
||
|
||
if err := pc.SendBitfield(ctx, bitfield); err != nil {
|
||
return
|
||
}
|
||
|
||
if err := pc.SendUnchoke(ctx); err != nil {
|
||
return
|
||
}
|
||
|
||
for {
|
||
if err := ctx.Err(); err != nil {
|
||
return
|
||
}
|
||
|
||
msg, ptr, err := pc.ReadMessage(ctx)
|
||
if err != nil {
|
||
if isTimeout(err) {
|
||
continue
|
||
}
|
||
return
|
||
}
|
||
|
||
switch msg.ID {
|
||
case 6: // msgRequest
|
||
if len(msg.Payload) < 12 {
|
||
if ptr != nil {
|
||
wireMsgPool.Put(ptr)
|
||
}
|
||
continue
|
||
}
|
||
idx := int(binary.BigEndian.Uint32(msg.Payload[0:4]))
|
||
begin := int(binary.BigEndian.Uint32(msg.Payload[4:8]))
|
||
length := int(binary.BigEndian.Uint32(msg.Payload[8:12]))
|
||
|
||
if ptr != nil {
|
||
wireMsgPool.Put(ptr)
|
||
}
|
||
|
||
if length > 16384*2 {
|
||
return
|
||
}
|
||
|
||
data, err := reader.ReadBlock(idx, begin, length)
|
||
if err != nil {
|
||
return
|
||
}
|
||
|
||
if err := pc.SendPiece(ctx, idx, begin, data); err != nil {
|
||
return
|
||
}
|
||
|
||
e.addUploadedBytes(int64(length))
|
||
default:
|
||
if ptr != nil {
|
||
wireMsgPool.Put(ptr)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
func (e *Engine) runPeerWorker(
|
||
ctx context.Context,
|
||
tf *torrentfile.TorrentFile,
|
||
peer PeerStatus,
|
||
pieceWriter *bufferedPieceWriter,
|
||
scheduler *pieceScheduler,
|
||
) {
|
||
peerKey := peerStatusKey(peer.Address, peer.Port)
|
||
defer scheduler.ReleasePeer(peerKey)
|
||
|
||
e.setPeerState(peer.Address, peer.Port, "connecting", "")
|
||
|
||
peerAddr := net.JoinHostPort(peer.Address, strconv.Itoa(int(peer.Port)))
|
||
logger.Info("ENGINE", "peer connection attempt: %s", peerAddr)
|
||
debugf("starting peer worker for %s", peerAddr)
|
||
client, err := newPeerClient(ctx, peerAddr, tf.InfoHash, e.peerID, len(tf.PieceHashes))
|
||
if err != nil {
|
||
if ctx.Err() != nil {
|
||
e.setPeerState(peer.Address, peer.Port, "stopped", "")
|
||
} else {
|
||
logger.Info("ENGINE", "peer failed: %s err=%v", peerAddr, err)
|
||
e.setPeerError(peer.Address, peer.Port, err.Error())
|
||
}
|
||
return
|
||
}
|
||
defer client.Close()
|
||
|
||
client.OnPex = func(peers []tracker.Peer) {
|
||
added := e.AddDiscoveredPeers(peers, "pex")
|
||
if added > 0 {
|
||
debugf("PEX: received %d new peers from %s", added, peerAddr)
|
||
}
|
||
}
|
||
|
||
workerCtx, cancelWorker := context.WithCancel(ctx)
|
||
defer cancelWorker()
|
||
|
||
go func() {
|
||
ticker := time.NewTicker(60 * time.Second)
|
||
defer ticker.Stop()
|
||
for {
|
||
select {
|
||
case <-workerCtx.Done():
|
||
return
|
||
case <-ticker.C:
|
||
if sample := e.GetActivePeersSample(50); len(sample) > 0 {
|
||
_ = client.SendPex(workerCtx, sample)
|
||
}
|
||
}
|
||
}
|
||
}()
|
||
|
||
logger.Info("ENGINE", "peer connected: %s", peerAddr)
|
||
|
||
e.setPeerState(peer.Address, peer.Port, "ready", "")
|
||
consecutiveFailures := 0
|
||
|
||
endgameCancelCh := scheduler.SubscribeCancel(peerKey)
|
||
defer scheduler.UnsubscribeCancel(peerKey)
|
||
|
||
for {
|
||
if ctx.Err() != nil {
|
||
e.setPeerState(peer.Address, peer.Port, "stopped", "")
|
||
return
|
||
}
|
||
|
||
if backoff := e.peerBackoff(peer.Address, peer.Port); backoff > 0 {
|
||
select {
|
||
case <-ctx.Done():
|
||
e.setPeerState(peer.Address, peer.Port, "stopped", "")
|
||
return
|
||
case <-time.After(backoff):
|
||
}
|
||
}
|
||
|
||
have, hasInfo := client.PieceAvailability()
|
||
task, ok, err := scheduler.Acquire(ctx, peerKey, have, hasInfo)
|
||
if err != nil {
|
||
if ctx.Err() != nil {
|
||
e.setPeerState(peer.Address, peer.Port, "stopped", "")
|
||
} else {
|
||
logger.Info("ENGINE", "peer %s scheduler acquire failed: %v", peerAddr, err)
|
||
e.setPeerError(peer.Address, peer.Port, err.Error())
|
||
}
|
||
return
|
||
}
|
||
|
||
if !ok {
|
||
select {
|
||
case <-ctx.Done():
|
||
e.setPeerState(peer.Address, peer.Port, "stopped", "")
|
||
return
|
||
case <-scheduler.Done():
|
||
e.setPeerState(peer.Address, peer.Port, "done", "")
|
||
return
|
||
case <-time.After(250 * time.Millisecond):
|
||
}
|
||
continue
|
||
}
|
||
e.setPieceState(task.Index, PieceDownloading)
|
||
|
||
e.setPeerState(peer.Address, peer.Port, "requesting", "")
|
||
pieceSize := pieceSizeForIndex(tf, task.Index)
|
||
|
||
pieceCancelCh := make(chan struct{})
|
||
pieceDoneCh := make(chan struct{})
|
||
go func() {
|
||
for {
|
||
select {
|
||
case canceledPiece := <-endgameCancelCh:
|
||
if canceledPiece == task.Index {
|
||
close(pieceCancelCh)
|
||
return
|
||
}
|
||
case <-pieceDoneCh:
|
||
return
|
||
case <-ctx.Done():
|
||
return
|
||
}
|
||
}
|
||
}()
|
||
|
||
pieceData, transferStats, err := client.DownloadPiece(ctx, task.Index, pieceSize, pieceCancelCh)
|
||
close(pieceDoneCh)
|
||
|
||
if err != nil {
|
||
if _, reportErr := scheduler.Report(ctx, task.Index, false); reportErr != nil && ctx.Err() == nil {
|
||
logger.Info("ENGINE", "scheduler report failure for piece %d after peer error: %v", task.Index, reportErr)
|
||
}
|
||
e.setPieceState(task.Index, PieceMissing)
|
||
if ctx.Err() != nil {
|
||
e.setPeerState(peer.Address, peer.Port, "stopped", "")
|
||
return
|
||
}
|
||
|
||
if errors.Is(err, ErrPieceCanceled) {
|
||
// Soft cancel: piece was downloaded by another peer, continue.
|
||
continue
|
||
}
|
||
|
||
consecutiveFailures++
|
||
logger.Info("ENGINE", "peer %s disconnected: %v", peerAddr, err)
|
||
e.setPeerError(peer.Address, peer.Port, err.Error())
|
||
if shouldDropPeer(err) || consecutiveFailures >= 3 {
|
||
return
|
||
}
|
||
continue
|
||
}
|
||
|
||
hash := sha1.Sum(pieceData)
|
||
if hash != tf.PieceHashes[task.Index] {
|
||
if _, reportErr := scheduler.Report(ctx, task.Index, false); reportErr != nil && ctx.Err() == nil {
|
||
logger.Info("ENGINE", "scheduler report hash mismatch for piece %d failed: %v", task.Index, reportErr)
|
||
}
|
||
e.setPieceState(task.Index, PieceMissing)
|
||
consecutiveFailures++
|
||
logger.Info("ENGINE", "peer %s piece %d hash mismatch", peerAddr, task.Index)
|
||
e.setPeerError(peer.Address, peer.Port, fmt.Sprintf("piece %d hash mismatch", task.Index))
|
||
if consecutiveFailures >= 3 {
|
||
return
|
||
}
|
||
continue
|
||
}
|
||
|
||
if err := pieceWriter.WritePiece(ctx, task.Index, pieceData); err != nil {
|
||
if _, reportErr := scheduler.Report(ctx, task.Index, false); reportErr != nil && ctx.Err() == nil {
|
||
logger.Info("ENGINE", "scheduler report write failure for piece %d failed: %v", task.Index, reportErr)
|
||
}
|
||
e.setPieceState(task.Index, PieceMissing)
|
||
logger.Info("ENGINE", "io error buffering piece %d from peer %s: %v", task.Index, peerAddr, err)
|
||
e.setPeerError(peer.Address, peer.Port, fmt.Sprintf("write piece %d: %v", task.Index, err))
|
||
return
|
||
}
|
||
|
||
accepted, reportErr := scheduler.Report(ctx, task.Index, true)
|
||
if reportErr != nil {
|
||
if ctx.Err() != nil {
|
||
e.setPeerState(peer.Address, peer.Port, "stopped", "")
|
||
} else {
|
||
e.setPeerError(peer.Address, peer.Port, reportErr.Error())
|
||
}
|
||
return
|
||
}
|
||
if accepted {
|
||
e.recordPieceComplete(peer.Address, peer.Port, task.Index, len(pieceData), transferStats)
|
||
debugf("peer %s committed piece %d", peerAddr, task.Index)
|
||
} else {
|
||
e.setPieceState(task.Index, PieceCompleted)
|
||
debugf("scheduler rejected piece %d report from %s", task.Index, peerAddr)
|
||
}
|
||
|
||
consecutiveFailures = 0
|
||
e.setPeerState(peer.Address, peer.Port, "active", "")
|
||
}
|
||
}
|
||
|
||
func shouldDropPeer(err error) bool {
|
||
if err == nil {
|
||
return false
|
||
}
|
||
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||
return true
|
||
}
|
||
if isTimeout(err) {
|
||
return false
|
||
}
|
||
|
||
msg := strings.ToLower(err.Error())
|
||
switch {
|
||
case strings.Contains(msg, "peer choked"):
|
||
return false
|
||
case strings.Contains(msg, "peer did not unchoke"):
|
||
return false
|
||
case strings.Contains(msg, "hash mismatch"):
|
||
return false
|
||
default:
|
||
return true
|
||
}
|
||
}
|
||
|
||
func (e *Engine) snapshotPeers() []PeerStatus {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
return append([]PeerStatus(nil), e.peers...)
|
||
}
|
||
|
||
func (e *Engine) AddDiscoveredPeers(newPeers []tracker.Peer, source string) int {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
|
||
if e.peerIdx == nil {
|
||
e.peerIdx = make(map[string]int)
|
||
}
|
||
|
||
added := 0
|
||
for _, p := range newPeers {
|
||
host := p.IP.String()
|
||
key := peerStatusKey(host, p.Port)
|
||
|
||
if _, exists := e.peerIdx[key]; !exists {
|
||
e.peerIdx[key] = len(e.peers)
|
||
e.peers = append(e.peers, PeerStatus{
|
||
Address: host,
|
||
Port: p.Port,
|
||
Source: source,
|
||
State: "pending",
|
||
Score: 0.5,
|
||
})
|
||
added++
|
||
}
|
||
}
|
||
return added
|
||
}
|
||
|
||
func (e *Engine) GetActivePeersSample(limit int) []tracker.Peer {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
|
||
var active []tracker.Peer
|
||
for _, p := range e.peers {
|
||
if p.State == "active" || p.State == "seeding" {
|
||
if ip := net.ParseIP(p.Address); ip != nil {
|
||
active = append(active, tracker.Peer{IP: ip, Port: p.Port})
|
||
}
|
||
}
|
||
}
|
||
|
||
if len(active) > limit {
|
||
mathrand.Shuffle(len(active), func(i, j int) { active[i], active[j] = active[j], active[i] })
|
||
return active[:limit]
|
||
}
|
||
return active
|
||
}
|
||
|
||
func (e *Engine) refreshPeersFromTrackers(ctx context.Context, tf *torrentfile.TorrentFile) int {
|
||
peers, trackerStatuses := e.queryTrackers(ctx, collectTrackersToTry(tf), tf.InfoHash, tf.Length, tracker.AnnounceOptions{
|
||
PeerID: e.peerID,
|
||
Port: trackerPort,
|
||
Downloaded: e.currentDownloadedBytes(),
|
||
NumWant: trackerNumWant,
|
||
Timeout: trackerAnnounceTimeout,
|
||
})
|
||
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
|
||
if len(trackerStatuses) > 0 {
|
||
e.trackers = trackerStatuses
|
||
}
|
||
if len(peers) == 0 {
|
||
waitErr := fmt.Errorf("waiting for peers via %d trackers", len(trackerStatuses))
|
||
if trErr := firstTrackerError(trackerStatuses); trErr != nil {
|
||
waitErr = fmt.Errorf("waiting for peers: %w", trErr)
|
||
}
|
||
e.lastErr = waitErr
|
||
debugf("reannounce returned no peers")
|
||
return 0
|
||
}
|
||
e.lastErr = nil
|
||
|
||
// Используем peerIdx для O(1) дедупликации
|
||
if e.peerIdx == nil {
|
||
e.peerIdx = buildPeerIdx(e.peers)
|
||
}
|
||
|
||
added := 0
|
||
for _, peer := range peers {
|
||
key := peerStatusKey(peer.Address, peer.Port)
|
||
if _, ok := e.peerIdx[key]; ok {
|
||
continue
|
||
}
|
||
e.peerIdx[key] = len(e.peers)
|
||
e.peers = append(e.peers, peer)
|
||
added++
|
||
}
|
||
|
||
if added > 0 {
|
||
sort.Slice(e.peers, func(i, j int) bool {
|
||
if e.peers[i].Address == e.peers[j].Address {
|
||
return e.peers[i].Port < e.peers[j].Port
|
||
}
|
||
return e.peers[i].Address < e.peers[j].Address
|
||
})
|
||
// Пересобираем индекс после сортировки
|
||
e.peerIdx = buildPeerIdx(e.peers)
|
||
logger.Info("ENGINE", "discovered %d new peers (total=%d)", added, len(e.peers))
|
||
}
|
||
|
||
return added
|
||
}
|
||
|
||
func (e *Engine) currentDownloadedBytes() int64 {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
return e.downloadedBytes
|
||
}
|
||
|
||
func (e *Engine) currentCompletedPieces() int {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
return e.completedPieces
|
||
}
|
||
|
||
func (e *Engine) currentActivePeers() int {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
return countActivePeers(e.peers)
|
||
}
|
||
|
||
func (e *Engine) recordPieceComplete(address string, port uint16, pieceIndex int, pieceSize int, transfer pieceTransferStats) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
now := time.Now()
|
||
e.downloadedBytes += int64(pieceSize)
|
||
if transfer.UploadedBytes > 0 {
|
||
e.uploadedBytes += transfer.UploadedBytes
|
||
}
|
||
e.completedPieces++
|
||
if pieceIndex >= 0 && pieceIndex < len(e.pieceStates) {
|
||
e.pieceStates[pieceIndex] = PieceCompleted
|
||
}
|
||
e.updateSpeedLocked(now)
|
||
key := peerStatusKey(address, port)
|
||
if peer := e.peerByKey(key); peer != nil {
|
||
peer.DownloadedPieces++
|
||
peer.DownloadedBytes += int64(pieceSize)
|
||
if transfer.UploadedBytes > 0 {
|
||
peer.UploadedBytes += transfer.UploadedBytes
|
||
}
|
||
if transfer.AvgBlockLatency > 0 {
|
||
peer.LatencyMS = int(transfer.AvgBlockLatency / time.Millisecond)
|
||
}
|
||
if transfer.Duration > 0 {
|
||
sampleSpeed := float64(pieceSize) / transfer.Duration.Seconds()
|
||
if peer.DownloadSpeed <= 0 {
|
||
peer.DownloadSpeed = sampleSpeed
|
||
} else {
|
||
peer.DownloadSpeed = (peer.DownloadSpeed * 0.65) + (sampleSpeed * 0.35)
|
||
}
|
||
}
|
||
peer.Score = computePeerScore(peer.DownloadSpeed, time.Duration(peer.LatencyMS)*time.Millisecond, peer.ErrorCount)
|
||
peer.Error = ""
|
||
}
|
||
}
|
||
|
||
// asyncSaveBitmap сохраняет bitmap в фоне (вызывается из runDownload heartbeat).
|
||
func (e *Engine) asyncSaveBitmap(bitmapPath string) {
|
||
e.mu.RLock()
|
||
completed := make([]int, 0, e.completedPieces)
|
||
total := len(e.pieceStates)
|
||
for i, s := range e.pieceStates {
|
||
if s == PieceCompleted {
|
||
completed = append(completed, i)
|
||
}
|
||
}
|
||
e.mu.RUnlock()
|
||
go func() { _ = saveBitmap(bitmapPath, completed, total) }()
|
||
}
|
||
|
||
// peerByKey возвращает указатель на PeerStatus по ключу (O(1)).
|
||
// Вызывать только под e.mu.
|
||
func (e *Engine) peerByKey(key string) *PeerStatus {
|
||
if e.peerIdx == nil {
|
||
return nil
|
||
}
|
||
if i, ok := e.peerIdx[key]; ok && i < len(e.peers) {
|
||
return &e.peers[i]
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (e *Engine) setPeerState(address string, port uint16, state, errText string) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
key := peerStatusKey(address, port)
|
||
if p := e.peerByKey(key); p != nil {
|
||
p.State = state
|
||
p.Error = errText
|
||
p.Score = computePeerScore(p.DownloadSpeed, time.Duration(p.LatencyMS)*time.Millisecond, p.ErrorCount)
|
||
}
|
||
}
|
||
|
||
func (e *Engine) setPeerError(address string, port uint16, errText string) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
key := peerStatusKey(address, port)
|
||
if p := e.peerByKey(key); p != nil {
|
||
p.State = "error"
|
||
p.Error = errText
|
||
p.ErrorCount++
|
||
p.Score = computePeerScore(p.DownloadSpeed, time.Duration(p.LatencyMS)*time.Millisecond, p.ErrorCount)
|
||
}
|
||
}
|
||
|
||
func (e *Engine) setPieceState(pieceIndex int, state PieceState) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
if pieceIndex < 0 || pieceIndex >= len(e.pieceStates) {
|
||
return
|
||
}
|
||
if e.pieceStates[pieceIndex] == PieceCompleted {
|
||
return
|
||
}
|
||
e.pieceStates[pieceIndex] = state
|
||
}
|
||
|
||
func (e *Engine) peerBackoff(address string, port uint16) time.Duration {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
key := peerStatusKey(address, port)
|
||
p := e.peerByKey(key)
|
||
if p == nil {
|
||
return 0
|
||
}
|
||
if p.ErrorCount >= 6 {
|
||
return 1200 * time.Millisecond
|
||
}
|
||
switch {
|
||
case p.Score <= 0 && p.ErrorCount == 0:
|
||
return 0
|
||
case p.Score < 0.20:
|
||
return 900 * time.Millisecond
|
||
case p.Score < 0.35:
|
||
return 550 * time.Millisecond
|
||
case p.Score < 0.50:
|
||
return 250 * time.Millisecond
|
||
default:
|
||
return 0
|
||
}
|
||
}
|
||
|
||
func computePeerScore(downloadSpeed float64, latency time.Duration, errorCount int) float64 {
|
||
if downloadSpeed < 0 {
|
||
downloadSpeed = 0
|
||
}
|
||
if latency < 0 {
|
||
latency = 0
|
||
}
|
||
if errorCount < 0 {
|
||
errorCount = 0
|
||
}
|
||
if downloadSpeed == 0 && latency == 0 && errorCount == 0 {
|
||
return 0.7
|
||
}
|
||
|
||
speedFactor := 0.0
|
||
if downloadSpeed > 0 {
|
||
speedFactor = downloadSpeed / (downloadSpeed + 600*1024)
|
||
}
|
||
|
||
latencyFactor := 1.0 / (1.0 + latency.Seconds()*5.0)
|
||
errorPenalty := 1.0 / (1.0 + float64(errorCount)*0.35)
|
||
|
||
score := (0.70*speedFactor + 0.30*latencyFactor) * errorPenalty
|
||
if score < 0 {
|
||
return 0
|
||
}
|
||
if score > 1 {
|
||
return 1
|
||
}
|
||
return score
|
||
}
|
||
|
||
func countActivePeers(peers []PeerStatus) int {
|
||
active := 0
|
||
for _, peer := range peers {
|
||
switch peer.State {
|
||
case "connecting", "ready", "requesting", "active":
|
||
active++
|
||
}
|
||
}
|
||
return active
|
||
}
|
||
|
||
func (e *Engine) setTerminalError(phase string, err error) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
e.phase = phase
|
||
e.lastErr = err
|
||
e.downloadSpeed = 0
|
||
e.uploadSpeed = 0
|
||
}
|
||
|
||
func (e *Engine) progressUnsafe() float64 {
|
||
if e.torrent != nil && e.torrent.Length > 0 {
|
||
if e.phase == "completed" {
|
||
return 1
|
||
}
|
||
if e.downloadedBytes > 0 {
|
||
progress := float64(e.downloadedBytes) / float64(e.torrent.Length)
|
||
if progress > 1 {
|
||
return 1
|
||
}
|
||
return progress
|
||
}
|
||
}
|
||
|
||
switch e.phase {
|
||
case "idle", "failed", "stopped", "stalled", "tracker_errors":
|
||
return 0
|
||
case "loading_metadata":
|
||
return 0.05
|
||
case "querying_trackers":
|
||
return 0.12
|
||
case "ready":
|
||
return 0.2
|
||
case "preparing_download":
|
||
return 0.25
|
||
case "downloading":
|
||
return 0.3
|
||
case "writing_files":
|
||
return 0.95
|
||
default:
|
||
return 0
|
||
}
|
||
}
|
||
|
||
func (e *Engine) resetStateForNewLoad() {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
e.torrent = nil
|
||
e.peers = nil
|
||
e.peerIdx = nil
|
||
e.trackers = nil
|
||
e.lastErr = nil
|
||
e.phase = "idle"
|
||
e.downloadedBytes = 0
|
||
e.uploadedBytes = 0
|
||
e.completedPieces = 0
|
||
e.totalPieces = 0
|
||
e.outputPath = ""
|
||
e.outputRoot = ""
|
||
e.downloadSpeed = 0
|
||
e.uploadSpeed = 0
|
||
e.speedAt = time.Time{}
|
||
e.downloadAtBytes = 0
|
||
e.uploadAtBytes = 0
|
||
e.pieceStates = nil
|
||
}
|
||
|
||
// buildPeerIdx строит map key→index для быстрого поиска пира.
|
||
func buildPeerIdx(peers []PeerStatus) map[string]int {
|
||
idx := make(map[string]int, len(peers))
|
||
for i, p := range peers {
|
||
idx[peerStatusKey(p.Address, p.Port)] = i
|
||
}
|
||
return idx
|
||
}
|
||
|
||
func (e *Engine) updateSpeed(now time.Time) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
e.updateSpeedLocked(now)
|
||
}
|
||
|
||
func (e *Engine) updateSpeedLocked(now time.Time) {
|
||
if e.speedAt.IsZero() {
|
||
e.speedAt = now
|
||
e.downloadAtBytes = e.downloadedBytes
|
||
e.uploadAtBytes = e.uploadedBytes
|
||
e.downloadSpeed = 0
|
||
e.uploadSpeed = 0
|
||
return
|
||
}
|
||
elapsed := now.Sub(e.speedAt).Seconds()
|
||
if elapsed < 0.8 {
|
||
return
|
||
}
|
||
downloadDelta := e.downloadedBytes - e.downloadAtBytes
|
||
if downloadDelta < 0 {
|
||
downloadDelta = 0
|
||
}
|
||
uploadDelta := e.uploadedBytes - e.uploadAtBytes
|
||
if uploadDelta < 0 {
|
||
uploadDelta = 0
|
||
}
|
||
// EWMA α=0.3: новый сэмпл весит 30%, история — 70%
|
||
const alpha = 0.3
|
||
sampleDown := float64(downloadDelta) / elapsed
|
||
sampleUp := float64(uploadDelta) / elapsed
|
||
if e.downloadSpeed <= 0 {
|
||
e.downloadSpeed = sampleDown
|
||
} else {
|
||
e.downloadSpeed = alpha*sampleDown + (1-alpha)*e.downloadSpeed
|
||
}
|
||
if e.uploadSpeed <= 0 {
|
||
e.uploadSpeed = sampleUp
|
||
} else {
|
||
e.uploadSpeed = alpha*sampleUp + (1-alpha)*e.uploadSpeed
|
||
}
|
||
|
||
e.speedAt = now
|
||
e.downloadAtBytes = e.downloadedBytes
|
||
e.uploadAtBytes = e.uploadedBytes
|
||
}
|
||
|
||
func (e *Engine) setError(err error) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
e.lastErr = err
|
||
}
|
||
|
||
func (e *Engine) setPhase(phase string) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
e.phase = phase
|
||
}
|
||
|
||
func (e *Engine) swapCancel(cancel context.CancelFunc) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
if e.cancel != nil {
|
||
e.cancel()
|
||
}
|
||
e.cancel = cancel
|
||
}
|
||
|
||
func (e *Engine) clearCancel() {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
if e.cancel != nil {
|
||
e.cancel()
|
||
e.cancel = nil
|
||
}
|
||
}
|
||
|
||
func generatePeerID() [20]byte {
|
||
const prefix = "-ZT0001-"
|
||
const alphabet = "0123456789abcdefghijklmnopqrstuvwxyz"
|
||
|
||
var id [20]byte
|
||
copy(id[:], prefix)
|
||
|
||
buf := make([]byte, len(id)-len(prefix))
|
||
if _, err := rand.Read(buf); err != nil {
|
||
now := time.Now().UnixNano()
|
||
for i := range buf {
|
||
buf[i] = byte(now >> (i * 8))
|
||
}
|
||
}
|
||
|
||
for i, b := range buf {
|
||
id[len(prefix)+i] = alphabet[int(b)%len(alphabet)]
|
||
}
|
||
|
||
return id
|
||
}
|
||
|
||
func addPeer(peerMap map[string]PeerStatus, peer tracker.Peer, source string) {
|
||
key := peerStatusKey(peer.IP.String(), peer.Port)
|
||
if _, exists := peerMap[key]; exists {
|
||
return
|
||
}
|
||
peerMap[key] = PeerStatus{
|
||
Address: peer.IP.String(),
|
||
Port: peer.Port,
|
||
Source: source,
|
||
State: "discovered",
|
||
Score: 0.5,
|
||
}
|
||
}
|
||
|
||
func peerStatusKey(address string, port uint16) string {
|
||
return net.JoinHostPort(address, strconv.Itoa(int(port)))
|
||
}
|
||
|
||
func collectTrackersToTry(tf *torrentfile.TorrentFile) []string {
|
||
const fallbackTrackers = `
|
||
udp://open.stealth.si:80/announce
|
||
udp://tracker.opentrackr.org:1337/announce
|
||
udp://tracker.openbittorrent.com:6969/announce
|
||
udp://tracker.torrent.eu.org:451/announce
|
||
https://tracker.opentrackr.org:443/announce
|
||
http://tracker.opentrackr.org:1337/announce
|
||
`
|
||
|
||
seen := make(map[string]struct{})
|
||
list := make([]string, 0, len(tf.Trackers)+8)
|
||
|
||
add := func(tr string) {
|
||
if tr == "" {
|
||
return
|
||
}
|
||
if _, ok := seen[tr]; ok {
|
||
return
|
||
}
|
||
seen[tr] = struct{}{}
|
||
list = append(list, tr)
|
||
}
|
||
|
||
for _, tr := range tf.Trackers {
|
||
add(tr)
|
||
}
|
||
for _, tr := range splitLines(fallbackTrackers) {
|
||
add(tr)
|
||
}
|
||
|
||
return list
|
||
}
|
||
|
||
func splitLines(s string) []string {
|
||
parts := strings.Split(s, "\n")
|
||
out := make([]string, 0, len(parts))
|
||
for _, part := range parts {
|
||
trimmed := strings.TrimSpace(part)
|
||
if trimmed == "" {
|
||
continue
|
||
}
|
||
out = append(out, trimmed)
|
||
}
|
||
return out
|
||
}
|
||
|
||
func firstTrackerError(statuses []TrackerStatus) error {
|
||
for _, st := range statuses {
|
||
if st.State == "error" && st.Error != "" {
|
||
return errors.New(st.Error)
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func createOrOpenPartFile(tf *torrentfile.TorrentFile, outputRoot string) (string, *os.File, []int, error) {
|
||
partsDir := filepath.Join(outputRoot, ".ztorrent-parts")
|
||
if err := os.MkdirAll(partsDir, 0o755); err != nil {
|
||
return "", nil, nil, err
|
||
}
|
||
|
||
hashPrefix := hex.EncodeToString(tf.InfoHash[:6])
|
||
name := sanitizePathPart(tf.Name) + "-" + hashPrefix + ".part"
|
||
partPath := filepath.Join(partsDir, name)
|
||
bitmapPath := partPath + ".bitmap"
|
||
|
||
// Пробуем открыть существующий .part файл (resume)
|
||
existing, err := os.OpenFile(partPath, os.O_RDWR, 0o644)
|
||
if err == nil {
|
||
// Файл существует — пробуем загрузить bitmap и верифицировать куски
|
||
var completed []int
|
||
if bitmapPieces, bErr := loadBitmap(bitmapPath, len(tf.PieceHashes)); bErr == nil && len(bitmapPieces) > 0 {
|
||
// Быстрый путь: bitmap есть, только верифицируем упомянутые куски
|
||
completed = verifyPieces(existing, tf, bitmapPieces)
|
||
logger.Info("ENGINE", "resume via bitmap: %d pieces verified", len(completed))
|
||
} else {
|
||
// Медленный путь: сканируем все куски через SHA-1
|
||
completed = verifyAllPieces(existing, tf)
|
||
logger.Info("ENGINE", "resume via full scan: %d/%d pieces verified", len(completed), len(tf.PieceHashes))
|
||
}
|
||
// Перезаписываем bitmap актуальными данными
|
||
_ = saveBitmap(bitmapPath, completed, len(tf.PieceHashes))
|
||
return partPath, existing, completed, nil
|
||
}
|
||
|
||
// Файла нет — создаём новый
|
||
f, err := os.OpenFile(partPath, os.O_CREATE|os.O_RDWR, 0o644)
|
||
if err != nil {
|
||
return "", nil, nil, err
|
||
}
|
||
if err := f.Truncate(int64(tf.Length)); err != nil {
|
||
f.Close()
|
||
return "", nil, nil, err
|
||
}
|
||
return partPath, f, nil, nil
|
||
}
|
||
|
||
// verifyAllPieces сканирует весь .part файл и возвращает индексы валидных кусков.
|
||
func verifyAllPieces(f *os.File, tf *torrentfile.TorrentFile) []int {
|
||
var completed []int
|
||
for i, expectedHash := range tf.PieceHashes {
|
||
pieceLen := pieceSizeForIndex(tf, i)
|
||
if pieceLen <= 0 {
|
||
continue
|
||
}
|
||
buf := make([]byte, pieceLen)
|
||
n, err := f.ReadAt(buf, int64(i)*int64(tf.PieceLength))
|
||
if err != nil || n != pieceLen {
|
||
continue
|
||
}
|
||
if sha1.Sum(buf) == expectedHash {
|
||
completed = append(completed, i)
|
||
}
|
||
}
|
||
return completed
|
||
}
|
||
|
||
// verifyPieces верифицирует только куски из списка (используется при наличии bitmap).
|
||
func verifyPieces(f *os.File, tf *torrentfile.TorrentFile, indices []int) []int {
|
||
var completed []int
|
||
for _, i := range indices {
|
||
if i < 0 || i >= len(tf.PieceHashes) {
|
||
continue
|
||
}
|
||
pieceLen := pieceSizeForIndex(tf, i)
|
||
if pieceLen <= 0 {
|
||
continue
|
||
}
|
||
buf := make([]byte, pieceLen)
|
||
n, err := f.ReadAt(buf, int64(i)*int64(tf.PieceLength))
|
||
if err != nil || n != pieceLen {
|
||
continue
|
||
}
|
||
if sha1.Sum(buf) == tf.PieceHashes[i] {
|
||
completed = append(completed, i)
|
||
}
|
||
}
|
||
return completed
|
||
}
|
||
|
||
// bitmapData — структура для сериализации в JSON.
|
||
type bitmapData struct {
|
||
Completed []int `json:"completed"`
|
||
Total int `json:"total"`
|
||
}
|
||
|
||
// saveBitmap записывает список завершённых кусков в JSON файл.
|
||
func saveBitmap(path string, completed []int, total int) error {
|
||
data := bitmapData{Completed: completed, Total: total}
|
||
b, err := json.Marshal(data)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return os.WriteFile(path, b, 0o644)
|
||
}
|
||
|
||
// loadBitmap читает список завершённых кусков из JSON файла.
|
||
func loadBitmap(path string, expectedTotal int) ([]int, error) {
|
||
b, err := os.ReadFile(path)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
var data bitmapData
|
||
if err := json.Unmarshal(b, &data); err != nil {
|
||
return nil, err
|
||
}
|
||
if data.Total != expectedTotal {
|
||
return nil, fmt.Errorf("bitmap total mismatch: got %d want %d", data.Total, expectedTotal)
|
||
}
|
||
return data.Completed, nil
|
||
}
|
||
|
||
func materializeDownloadedFiles(tf *torrentfile.TorrentFile, partPath, outputRoot string) error {
|
||
partFile, err := os.Open(partPath)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer partFile.Close()
|
||
|
||
if err := os.MkdirAll(outputRoot, 0o755); err != nil {
|
||
return err
|
||
}
|
||
|
||
offset := int64(0)
|
||
for _, f := range tf.Files {
|
||
targetPath, err := safeOutputPath(outputRoot, f.Path)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if err := os.MkdirAll(filepath.Dir(targetPath), 0o755); err != nil {
|
||
return err
|
||
}
|
||
|
||
out, err := os.OpenFile(targetPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
_, copyErr := io.CopyN(out, io.NewSectionReader(partFile, offset, int64(f.Length)), int64(f.Length))
|
||
closeErr := out.Close()
|
||
if copyErr != nil {
|
||
return copyErr
|
||
}
|
||
if closeErr != nil {
|
||
return closeErr
|
||
}
|
||
|
||
offset += int64(f.Length)
|
||
}
|
||
|
||
if offset != int64(tf.Length) {
|
||
return fmt.Errorf("written bytes mismatch: expected %d, wrote %d", tf.Length, offset)
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
func safeOutputPath(root, rel string) (string, error) {
|
||
clean := filepath.Clean(filepath.FromSlash(rel))
|
||
if clean == "." || clean == "" {
|
||
return "", errors.New("invalid output file path")
|
||
}
|
||
if filepath.IsAbs(clean) || strings.HasPrefix(clean, "..") {
|
||
return "", fmt.Errorf("unsafe output path %q", rel)
|
||
}
|
||
return filepath.Join(root, clean), nil
|
||
}
|
||
|
||
func buildOutputPath(tf *torrentfile.TorrentFile, outputRoot string) string {
|
||
return filepath.Join(outputRoot, tf.Name)
|
||
}
|
||
|
||
func sanitizePathPart(s string) string {
|
||
s = strings.TrimSpace(s)
|
||
if s == "" {
|
||
return "torrent"
|
||
}
|
||
|
||
var b strings.Builder
|
||
for _, r := range s {
|
||
switch {
|
||
case unicode.IsLetter(r), unicode.IsDigit(r), r == '-', r == '_', r == '.':
|
||
b.WriteRune(r)
|
||
default:
|
||
b.WriteByte('_')
|
||
}
|
||
}
|
||
|
||
out := strings.Trim(b.String(), "._")
|
||
if out == "" {
|
||
return "torrent"
|
||
}
|
||
return out
|
||
}
|
||
|
||
func normalizeOutputRoot(root string) string {
|
||
root = strings.TrimSpace(root)
|
||
if root == "" {
|
||
cwd, err := os.Getwd()
|
||
if err == nil && cwd != "" {
|
||
return cwd
|
||
}
|
||
return "."
|
||
}
|
||
clean := filepath.Clean(root)
|
||
return clean
|
||
}
|
||
|
||
// listenIncoming принимает входящие BitTorrent соединения на порту 6881.
|
||
func (e *Engine) listenIncoming(ctx context.Context, tf *torrentfile.TorrentFile) {
|
||
ln, err := net.Listen("tcp", ":6881")
|
||
if err != nil {
|
||
// Порт занят — тихо выходим (не критично)
|
||
debugf("incoming listener skipped: %v", err)
|
||
return
|
||
}
|
||
defer ln.Close()
|
||
logger.Info("ENGINE", "listening for incoming peers on %s", ln.Addr())
|
||
|
||
// Open UPnP mapping
|
||
go e.portManager.OpenPort(ctx, 6881, "Ztorrent")
|
||
|
||
go func() {
|
||
<-ctx.Done()
|
||
e.portManager.ClosePort()
|
||
ln.Close()
|
||
}()
|
||
|
||
for {
|
||
conn, err := ln.Accept()
|
||
if err != nil {
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
default:
|
||
return
|
||
}
|
||
}
|
||
go e.handleIncoming(ctx, conn, tf)
|
||
}
|
||
}
|
||
|
||
// handleIncoming обрабатывает входящее соединение: handshake + регистрация пира.
|
||
func (e *Engine) handleIncoming(ctx context.Context, rawConn net.Conn, tf *torrentfile.TorrentFile) {
|
||
_ = rawConn.SetDeadline(time.Now().Add(peerReadTimeout))
|
||
head := make([]byte, 1)
|
||
if _, err := io.ReadFull(rawConn, head); err != nil {
|
||
rawConn.Close()
|
||
return
|
||
}
|
||
pstrlen := int(head[0])
|
||
if pstrlen <= 0 || pstrlen > 64 {
|
||
rawConn.Close()
|
||
return
|
||
}
|
||
rest := make([]byte, pstrlen+48)
|
||
if _, err := io.ReadFull(rawConn, rest); err != nil {
|
||
rawConn.Close()
|
||
return
|
||
}
|
||
if string(rest[:pstrlen]) != wireProtocolString {
|
||
rawConn.Close()
|
||
return
|
||
}
|
||
infoHashOffset := pstrlen + 8
|
||
var gotInfoHash [20]byte
|
||
copy(gotInfoHash[:], rest[infoHashOffset:infoHashOffset+20])
|
||
if gotInfoHash != tf.InfoHash {
|
||
rawConn.Close()
|
||
return
|
||
}
|
||
|
||
_ = rawConn.SetDeadline(time.Now().Add(peerWriteTimeout))
|
||
e.mu.RLock()
|
||
myID := e.peerID
|
||
e.mu.RUnlock()
|
||
reply := make([]byte, 1+pstrlen+48)
|
||
reply[0] = head[0]
|
||
copy(reply[1:], wireProtocolString)
|
||
copy(reply[1+pstrlen+8:], tf.InfoHash[:])
|
||
copy(reply[1+pstrlen+28:], myID[:])
|
||
if _, err := rawConn.Write(reply); err != nil {
|
||
rawConn.Close()
|
||
return
|
||
}
|
||
_ = rawConn.SetDeadline(time.Time{})
|
||
|
||
remoteAddr := rawConn.RemoteAddr().String()
|
||
host, portStr, err := net.SplitHostPort(remoteAddr)
|
||
if err != nil {
|
||
rawConn.Close()
|
||
return
|
||
}
|
||
portNum, _ := strconv.Atoi(portStr)
|
||
incomingPort := uint16(portNum)
|
||
incomingKey := peerStatusKey(host, incomingPort)
|
||
|
||
e.mu.Lock()
|
||
if e.peerIdx == nil {
|
||
e.peerIdx = make(map[string]int)
|
||
}
|
||
if _, exists := e.peerIdx[incomingKey]; !exists {
|
||
e.peerIdx[incomingKey] = len(e.peers)
|
||
e.peers = append(e.peers, PeerStatus{
|
||
Address: host,
|
||
Port: incomingPort,
|
||
Source: "incoming",
|
||
State: "connecting",
|
||
Score: 0.6,
|
||
})
|
||
}
|
||
e.mu.Unlock()
|
||
|
||
logger.Info("ENGINE", "incoming peer: %s", remoteAddr)
|
||
|
||
// Delegate to the active loop via channel without closing
|
||
select {
|
||
case e.incomingConns <- rawConn:
|
||
default:
|
||
rawConn.Close()
|
||
}
|
||
}
|
||
|
||
func (e *Engine) handleIncomingSeeding(ctx context.Context, tf *torrentfile.TorrentFile, conn net.Conn, reader *storage.PieceReader) {
|
||
go func() {
|
||
defer conn.Close()
|
||
// It was already handshaked, but we need a peerClient wrapper that bypasses the handshake.
|
||
// Actually, I can use the newIncomingPeerClient I created in peerwire.go
|
||
|
||
// Wait, peerwire.go has newIncomingPeerClient which STILL expects to do sendExtendedHandshake and readInitialMessages.
|
||
// Yes, that's correct.
|
||
// Extensions bit is in the handshake, but wait! We didn't save the extensions bit from their handshake!
|
||
// In handleIncoming, we didn't check their extensions bit.
|
||
// Let's assume extensions=true for now, or just false. We don't strictly need extensions for seeding basic pieces.
|
||
|
||
pc, err := newIncomingPeerClient(ctx, conn, false, len(tf.PieceHashes))
|
||
if err != nil {
|
||
return
|
||
}
|
||
defer pc.Close()
|
||
|
||
pc.OnPex = func(peers []tracker.Peer) {
|
||
added := e.AddDiscoveredPeers(peers, "pex")
|
||
if added > 0 {
|
||
debugf("PEX: received %d new peers from incoming conn", added)
|
||
}
|
||
}
|
||
|
||
bitfield := make([]byte, (len(tf.PieceHashes)+7)/8)
|
||
for i := range bitfield {
|
||
bitfield[i] = 0xFF
|
||
}
|
||
lastPieces := len(tf.PieceHashes) % 8
|
||
if lastPieces != 0 {
|
||
mask := byte(0xFF << (8 - lastPieces))
|
||
bitfield[len(bitfield)-1] &= mask
|
||
}
|
||
|
||
if err := pc.SendBitfield(ctx, bitfield); err != nil {
|
||
return
|
||
}
|
||
|
||
if err := pc.SendUnchoke(ctx); err != nil {
|
||
return
|
||
}
|
||
|
||
for {
|
||
if err := ctx.Err(); err != nil {
|
||
return
|
||
}
|
||
|
||
msg, ptr, err := pc.ReadMessage(ctx)
|
||
if err != nil {
|
||
if isTimeout(err) {
|
||
continue
|
||
}
|
||
return
|
||
}
|
||
|
||
switch msg.ID {
|
||
case 6: // msgRequest
|
||
if len(msg.Payload) < 12 {
|
||
if ptr != nil {
|
||
wireMsgPool.Put(ptr)
|
||
}
|
||
continue
|
||
}
|
||
idx := int(binary.BigEndian.Uint32(msg.Payload[0:4]))
|
||
begin := int(binary.BigEndian.Uint32(msg.Payload[4:8]))
|
||
length := int(binary.BigEndian.Uint32(msg.Payload[8:12]))
|
||
|
||
if ptr != nil {
|
||
wireMsgPool.Put(ptr)
|
||
}
|
||
|
||
if length > 16384*2 {
|
||
return
|
||
}
|
||
|
||
data, err := reader.ReadBlock(idx, begin, length)
|
||
if err != nil {
|
||
return
|
||
}
|
||
|
||
if err := pc.SendPiece(ctx, idx, begin, data); err != nil {
|
||
return
|
||
}
|
||
|
||
e.addUploadedBytes(int64(length))
|
||
default:
|
||
if ptr != nil {
|
||
wireMsgPool.Put(ptr)
|
||
}
|
||
}
|
||
}
|
||
}()
|
||
}
|
||
|
||
// SetSequentialMode toggles sequential piece downloading mode on the fly
|
||
func (e *Engine) SetSequentialMode(mode bool) {
|
||
e.mu.RLock()
|
||
scheduler := e.scheduler
|
||
e.mu.RUnlock()
|
||
if scheduler != nil {
|
||
scheduler.SetSequential(mode)
|
||
}
|
||
}
|
||
|
||
// HasPiece is a thread-safe check to see if a piece is fully downloaded
|
||
func (e *Engine) HasCompletedPiece(index int) bool {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
if index < 0 || index >= len(e.pieceStates) {
|
||
return false
|
||
}
|
||
return e.pieceStates[index] == PieceCompleted
|
||
}
|
||
|
||
// ToggleFilePriority toggles a file's priority between 0 (skip) and 1 (normal).
|
||
func (e *Engine) ToggleFilePriority(fileIdx int) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
|
||
if e.torrent == nil || fileIdx < 0 || fileIdx >= len(e.torrent.Files) {
|
||
return
|
||
}
|
||
|
||
if e.torrent.Files[fileIdx].Priority == 0 {
|
||
e.torrent.Files[fileIdx].Priority = 1
|
||
} else {
|
||
e.torrent.Files[fileIdx].Priority = 0
|
||
}
|
||
}
|
||
|
||
// PartFilePath returns the path to the current .part file
|
||
func (e *Engine) PartFilePath() string {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
return e.outputPath
|
||
}
|