Ztorrent/internal/torrentfile/torrentfile.go

package torrentfile

import (
	"bytes"
	"crypto/sha1"
	"errors"
	"fmt"
	"os"
	"path"

	"github.com/jackpal/bencode-go"
)

type bencodeTorrent struct {
	Announce     string      `bencode:"announce"`
	AnnounceList [][]string  `bencode:"announce-list"`
	Info         bencodeInfo `bencode:"info"`
}

type bencodeInfo struct {
	Pieces      string `bencode:"pieces"`
	PieceLength int    `bencode:"piece length"`
	Length      int    `bencode:"length"`
	Files       []file `bencode:"files"`
	Name        string `bencode:"name"`
}

type file struct {
	Length int      `bencode:"length"`
	Path   []string `bencode:"path"`
}

type File struct {
	Path   string
	Length int
}

type TorrentFile struct {
	Announce    string
	Trackers    []string
	InfoHash    [20]byte
	PieceHashes [][20]byte
	PieceLength int
	Length      int
	Name        string
	Files       []File
}

func Open(path string) (*TorrentFile, error) {
	rawData, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}

	infoBytes, err := extractInfoBytes(rawData)
	if err != nil {
		return nil, err
	}

	var bto bencodeTorrent

	if err := bencode.Unmarshal(bytes.NewReader(rawData), &bto); err != nil {
		return nil, err
	}

	// Если announce пустой, берём первый трекер из announce-list (BEP 12)
	primaryAnnounce := bto.Announce
	if primaryAnnounce == "" {
		for _, tier := range bto.AnnounceList {
			for _, tr := range tier {
				if tr != "" {
					primaryAnnounce = tr
					break
				}
			}
			if primaryAnnounce != "" {
				break
			}
		}
	}

	if err := validateMetadata(bto.Info); err != nil {
		return nil, err
	}

	pieceHashes, err := splitPieceHashes(bto.Info.Pieces)
	if err != nil {
		return nil, err
	}

	files, totalLength, err := deriveFilesAndLength(bto.Info)
	if err != nil {
		return nil, err
	}

	if err := validatePieceCount(totalLength, bto.Info.PieceLength, len(pieceHashes)); err != nil {
		return nil, err
	}

	infoHash := sha1.Sum(infoBytes)
	trackers := collectTrackers(primaryAnnounce, bto.AnnounceList)

	tf := TorrentFile{
		Announce:    primaryAnnounce,
		Trackers:    trackers,
		InfoHash:    infoHash,
		PieceHashes: pieceHashes,
		PieceLength: bto.Info.PieceLength,
		Length:      totalLength,
		Name:        bto.Info.Name,
		Files:       files,
	}

	return &tf, nil
}

func FromMetadata(infoBytes []byte, trackers []string) (*TorrentFile, error) {
	var info bencodeInfo
	if err := bencode.Unmarshal(bytes.NewReader(infoBytes), &info); err != nil {
		return nil, err
	}

	primaryAnnounce := ""
	if len(trackers) > 0 {
		primaryAnnounce = trackers[0]
	}

	if err := validateMetadata(info); err != nil {
		return nil, err
	}

	pieceHashes, err := splitPieceHashes(info.Pieces)
	if err != nil {
		return nil, err
	}

	files, totalLength, err := deriveFilesAndLength(info)
	if err != nil {
		return nil, err
	}

	if err := validatePieceCount(totalLength, info.PieceLength, len(pieceHashes)); err != nil {
		return nil, err
	}

	infoHash := sha1.Sum(infoBytes)

	tf := TorrentFile{
		Announce:    primaryAnnounce,
		Trackers:    trackers,
		InfoHash:    infoHash,
		PieceHashes: pieceHashes,
		PieceLength: info.PieceLength,
		Length:      totalLength,
		Name:        info.Name,
		Files:       files,
	}

	return &tf, nil
}

func validateMetadata(info bencodeInfo) error {
	switch {
	case info.Name == "":
		return errors.New("torrent name is empty")
	case info.PieceLength <= 0:
		return errors.New("torrent piece length must be greater than zero")
	}

	switch {
	case info.Length > 0 && len(info.Files) > 0:
		return errors.New("torrent contains both single-file and multi-file metadata")
	case info.Length <= 0 && len(info.Files) == 0:
		return errors.New("torrent missing both length and files metadata")
	}

	return nil
}

func splitPieceHashes(rawPieces string) ([][20]byte, error) {
	pieces := []byte(rawPieces)
	if len(pieces) == 0 {
		return nil, errors.New("torrent has no piece hashes")
	}

	if len(pieces)%sha1.Size != 0 {
		return nil, fmt.Errorf("invalid pieces data size %d (must be multiple of %d)", len(pieces), sha1.Size)
	}

	hashes := make([][20]byte, len(pieces)/sha1.Size)
	for i := range hashes {
		start := i * sha1.Size
		copy(hashes[i][:], pieces[start:start+sha1.Size])
	}

	return hashes, nil
}

func deriveFilesAndLength(info bencodeInfo) ([]File, int, error) {
	if info.Length > 0 {
		return []File{{Path: info.Name, Length: info.Length}}, info.Length, nil
	}

	files := make([]File, 0, len(info.Files))
	totalLength := 0
	for _, f := range info.Files {
		if f.Length <= 0 {
			return nil, 0, errors.New("torrent file length must be greater than zero")
		}
		if len(f.Path) == 0 {
			return nil, 0, errors.New("torrent file path is empty")
		}

		parts := make([]string, 0, len(f.Path)+1)
		parts = append(parts, info.Name)
		parts = append(parts, f.Path...)
		filePath := path.Join(parts...)
		if filePath == "." {
			return nil, 0, errors.New("torrent file path is invalid")
		}

		files = append(files, File{
			Path:   filePath,
			Length: f.Length,
		})
		totalLength += f.Length
	}

	if totalLength <= 0 {
		return nil, 0, errors.New("torrent total length must be greater than zero")
	}

	return files, totalLength, nil
}

func validatePieceCount(totalLength, pieceLength, actualPieces int) error {
	expectedPieces := (totalLength + pieceLength - 1) / pieceLength
	if expectedPieces != actualPieces {
		return fmt.Errorf("piece hash count mismatch: expected %d, got %d", expectedPieces, actualPieces)
	}
	return nil
}

func extractInfoBytes(data []byte) ([]byte, error) {
	if len(data) == 0 || data[0] != 'd' {
		return nil, errors.New("torrent root must be a bencoded dictionary")
	}

	idx := 1
	for idx < len(data) {
		if data[idx] == 'e' {
			break
		}

		key, next, err := parseBencodeString(data, idx)
		if err != nil {
			return nil, err
		}
		idx = next

		valueStart := idx
		valueEnd, err := skipBencodeValue(data, idx)
		if err != nil {
			return nil, err
		}

		if string(key) == "info" {
			return data[valueStart:valueEnd], nil
		}
		idx = valueEnd
	}

	return nil, errors.New("torrent does not contain top-level info dictionary")
}

func skipBencodeValue(data []byte, idx int) (int, error) {
	if idx >= len(data) {
		return 0, errors.New("unexpected end of bencoded data")
	}

	switch c := data[idx]; {
	case c == 'i':
		return skipBencodeInt(data, idx)
	case c == 'l':
		idx++
		for {
			if idx >= len(data) {
				return 0, errors.New("unexpected end while parsing bencoded list")
			}
			if data[idx] == 'e' {
				return idx + 1, nil
			}
			next, err := skipBencodeValue(data, idx)
			if err != nil {
				return 0, err
			}
			idx = next
		}
	case c == 'd':
		idx++
		for {
			if idx >= len(data) {
				return 0, errors.New("unexpected end while parsing bencoded dictionary")
			}
			if data[idx] == 'e' {
				return idx + 1, nil
			}
			_, next, err := parseBencodeString(data, idx)
			if err != nil {
				return 0, err
			}
			idx = next

			next, err = skipBencodeValue(data, idx)
			if err != nil {
				return 0, err
			}
			idx = next
		}
	case c >= '0' && c <= '9':
		_, next, err := parseBencodeString(data, idx)
		return next, err
	default:
		return 0, fmt.Errorf("invalid bencode token %q at index %d", c, idx)
	}
}

func skipBencodeInt(data []byte, idx int) (int, error) {
	if data[idx] != 'i' {
		return 0, fmt.Errorf("expected integer token at index %d", idx)
	}
	idx++
	if idx >= len(data) {
		return 0, errors.New("unexpected end while parsing bencoded integer")
	}

	if data[idx] == '-' {
		idx++
		if idx >= len(data) {
			return 0, errors.New("unexpected end after bencoded integer sign")
		}
	}

	if data[idx] < '0' || data[idx] > '9' {
		return 0, fmt.Errorf("invalid bencoded integer digit at index %d", idx)
	}

	for idx < len(data) && data[idx] != 'e' {
		if data[idx] < '0' || data[idx] > '9' {
			return 0, fmt.Errorf("invalid bencoded integer digit at index %d", idx)
		}
		idx++
	}

	if idx >= len(data) || data[idx] != 'e' {
		return 0, errors.New("unterminated bencoded integer")
	}

	return idx + 1, nil
}

func parseBencodeString(data []byte, idx int) ([]byte, int, error) {
	length, valueStart, err := parseBencodeStringLength(data, idx)
	if err != nil {
		return nil, 0, err
	}

	valueEnd := valueStart + length
	if valueEnd > len(data) {
		return nil, 0, errors.New("bencoded string length exceeds input size")
	}

	return data[valueStart:valueEnd], valueEnd, nil
}

func parseBencodeStringLength(data []byte, idx int) (int, int, error) {
	if idx >= len(data) {
		return 0, 0, errors.New("unexpected end while parsing bencoded string length")
	}
	if data[idx] < '0' || data[idx] > '9' {
		return 0, 0, fmt.Errorf("invalid bencoded string length token at index %d", idx)
	}

	length := 0
	for idx < len(data) && data[idx] >= '0' && data[idx] <= '9' {
		length = (length * 10) + int(data[idx]-'0')
		idx++
	}

	if idx >= len(data) || data[idx] != ':' {
		return 0, 0, errors.New("unterminated bencoded string length")
	}

	return length, idx + 1, nil
}

func collectTrackers(primary string, announceList [][]string) []string {
	seen := make(map[string]struct{})
	trackers := make([]string, 0, 1+len(announceList))

	add := func(tracker string) {
		if tracker == "" {
			return
		}
		if _, ok := seen[tracker]; ok {
			return
		}
		seen[tracker] = struct{}{}
		trackers = append(trackers, tracker)
	}

	add(primary)
	for _, tier := range announceList {
		for _, tracker := range tier {
			add(tracker)
		}
	}

	return trackers
}