Ztorrent/internal/torrent/engine.go
2026-03-06 12:52:01 +03:00

367 lines
7 KiB
Go

package torrent
import (
"context"
"crypto/rand"
"errors"
"fmt"
"net"
"sync"
"time"
"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
trackers []TrackerStatus
peerID [20]byte
cancel context.CancelFunc
lastErr error
phase string
}
type Status struct {
Loaded bool
Name string
Length int
PieceLength int
PieceCount int
Files []torrentfile.File
Announce string
PeerCount int
Peers []PeerStatus
Trackers []TrackerStatus
PeerID string
Progress float64
Phase string
LastError string
}
type PeerStatus struct {
Address string
Port uint16
Source string
}
type TrackerStatus struct {
URL string
State string
PeerCount int
Error string
}
func NewEngine() *Engine {
return &Engine{
peerID: generatePeerID(),
phase: "idle",
}
}
func (e *Engine) LoadTorrent(path string) error {
e.setPhase("loading_metadata")
tf, err := torrentfile.Open(path)
if err != nil {
e.setError(err)
e.setPhase("failed")
return err
}
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
e.swapCancel(cancel)
defer e.clearCancel()
e.setPhase("querying_trackers")
peers, trackerStatuses := e.queryTrackers(ctx, tf, tracker.AnnounceOptions{
PeerID: e.peerID,
Port: 6881,
NumWant: 200,
Timeout: 6 * time.Second,
})
var finalErr error
if len(peers) == 0 {
if err := firstTrackerError(trackerStatuses); err != nil {
finalErr = fmt.Errorf("no peers found: %w", err)
} else {
finalErr = fmt.Errorf("no peers found via %d trackers", len(trackerStatuses))
}
}
e.mu.Lock()
e.torrent = tf
e.peers = peers
e.trackers = trackerStatuses
e.lastErr = finalErr
if len(peers) > 0 {
e.phase = "ready"
} else {
e.phase = "waiting_peers"
}
e.mu.Unlock()
if finalErr != nil {
e.mu.Lock()
if e.phase == "waiting_peers" {
e.phase = "tracker_errors"
}
e.mu.Unlock()
return finalErr
}
return nil
}
func (e *Engine) Stop() {
e.mu.Lock()
defer e.mu.Unlock()
if e.cancel != nil {
e.cancel()
e.cancel = nil
}
}
func (e *Engine) Progress() float64 {
e.mu.RLock()
defer e.mu.RUnlock()
switch e.phase {
case "idle":
return 0
case "loading_metadata":
return 0.1
case "querying_trackers":
return 0.25
case "ready":
return 0.4
case "waiting_peers", "tracker_errors":
return 0.3
default:
return 0
}
}
func (e *Engine) Status() Status {
e.mu.RLock()
defer e.mu.RUnlock()
status := Status{
Loaded: e.torrent != nil,
PeerCount: len(e.peers),
PeerID: string(e.peerID[:]),
Progress: e.progressUnsafe(),
Phase: e.phase,
Peers: append([]PeerStatus(nil), e.peers...),
Trackers: append([]TrackerStatus(nil), e.trackers...),
}
if e.torrent != nil {
status.Name = e.torrent.Name
status.Length = e.torrent.Length
status.PieceLength = e.torrent.PieceLength
status.PieceCount = len(e.torrent.PieceHashes)
status.Files = append(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, tf *torrentfile.TorrentFile, opts tracker.AnnounceOptions) ([]PeerStatus, []TrackerStatus) {
trackers := collectTrackersToTry(tf)
peerMap := make(map[string]PeerStatus)
statuses := make([]TrackerStatus, 0, len(trackers))
for _, announce := range trackers {
perTrackerCtx, cancel := context.WithTimeout(ctx, opts.Timeout)
peers, err := tracker.GetPeersFromURL(perTrackerCtx, announce, tf, opts)
cancel()
st := TrackerStatus{
URL: announce,
}
if err != nil {
st.State = "error"
st.Error = err.Error()
statuses = append(statuses, st)
continue
}
if len(peers) == 0 {
st.State = "empty"
} else {
st.State = "ok"
st.PeerCount = len(peers)
}
statuses = append(statuses, st)
for _, peer := range peers {
addPeer(peerMap, peer, announce)
}
}
peerList := make([]PeerStatus, 0, len(peerMap))
for _, peer := range peerMap {
peerList = append(peerList, peer)
}
return peerList, statuses
}
func addPeer(peerMap map[string]PeerStatus, peer tracker.Peer, source string) {
key := net.JoinHostPort(peer.IP.String(), fmt.Sprintf("%d", peer.Port))
if _, exists := peerMap[key]; exists {
return
}
peerMap[key] = PeerStatus{
Address: peer.IP.String(),
Port: peer.Port,
Source: source,
}
}
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 {
lines := make([]string, 0)
current := make([]byte, 0, len(s))
flush := func() {
if len(current) == 0 {
return
}
lines = append(lines, string(current))
current = current[:0]
}
for i := 0; i < len(s); i++ {
if s[i] == '\n' {
flush()
continue
}
if s[i] == '\r' {
continue
}
if s[i] == ' ' || s[i] == '\t' {
if len(current) == 0 {
continue
}
}
current = append(current, s[i])
}
flush()
return lines
}
func firstTrackerError(statuses []TrackerStatus) error {
for _, st := range statuses {
if st.State == "error" && st.Error != "" {
return errors.New(st.Error)
}
}
return nil
}
func (e *Engine) progressUnsafe() float64 {
switch e.phase {
case "idle":
return 0
case "loading_metadata":
return 0.1
case "querying_trackers":
return 0.25
case "ready":
return 0.4
case "waiting_peers", "tracker_errors":
return 0.3
default:
return 0
}
}
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
}