tstor/src/host/service/service.go

package service

import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"time"

	"git.kmsign.ru/royalcat/tstor/src/host/controller"
	"git.kmsign.ru/royalcat/tstor/src/host/filestorage"
	"git.kmsign.ru/royalcat/tstor/src/host/store"
	"git.kmsign.ru/royalcat/tstor/src/host/vfs"

	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/bencode"
	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/types"
	"github.com/anacrolix/torrent/types/infohash"
)

type Service struct {
	c             *torrent.Client
	excludedFiles *store.ExlcudedFiles
	infoBytes     *store.InfoBytes

	torrentLoaded chan struct{}

	// stats *Stats
	DefaultPriority types.PiecePriority
	Storage         *filestorage.FileStorage
	SourceDir       string

	log                     *slog.Logger
	addTimeout, readTimeout int
}

func NewService(sourceDir string, c *torrent.Client, storage *filestorage.FileStorage, excludedFiles *store.ExlcudedFiles, infoBytes *store.InfoBytes, addTimeout, readTimeout int) *Service {
	s := &Service{
		log:             slog.With("component", "torrent-service"),
		c:               c,
		DefaultPriority: types.PiecePriorityNone,
		excludedFiles:   excludedFiles,
		infoBytes:       infoBytes,
		Storage:         storage,
		SourceDir:       sourceDir,
		torrentLoaded:   make(chan struct{}),
		// stats:       newStats(), // TODO persistent
		addTimeout:  addTimeout,
		readTimeout: readTimeout,
	}

	go func() {
		err := s.loadTorrentFiles(context.Background())
		if err != nil {
			s.log.Error("initial torrent load failed", "error", err)
		}
		close(s.torrentLoaded)
	}()

	return s
}

var _ vfs.FsFactory = (*Service)(nil).NewTorrentFs

func (s *Service) AddTorrent(ctx context.Context, f vfs.File) (*torrent.Torrent, error) {
	defer f.Close()

	stat, err := f.Stat()
	if err != nil {
		return nil, fmt.Errorf("call stat failed: %w", err)
	}

	mi, err := metainfo.Load(f)
	if err != nil {
		return nil, fmt.Errorf("loading torrent metadata from file %s, error: %w", stat.Name(), err)
	}
	t, ok := s.c.Torrent(mi.HashInfoBytes())
	if !ok {
		spec, err := torrent.TorrentSpecFromMetaInfoErr(mi)
		if err != nil {
			return nil, fmt.Errorf("parse spec from metadata: %w", err)
		}
		infoBytes := spec.InfoBytes

		if !isValidInfoHashBytes(infoBytes) {
			infoBytes = nil
		}

		if len(infoBytes) == 0 {
			infoBytes, err = s.infoBytes.GetBytes(spec.InfoHash)
			if err != nil && err != store.ErrNotFound {
				return nil, fmt.Errorf("get info bytes from database: %w", err)
			}
		}

		var info metainfo.Info
		err = bencode.Unmarshal(infoBytes, &info)
		if err != nil {
			infoBytes = nil
		} else {
			for _, t := range s.c.Torrents() {
				if t.Name() == info.BestName() && t.InfoHash() != spec.InfoHash {
					<-t.GotInfo()
					if !isTorrentCompatable(*t.Info(), info) {
						return nil, fmt.Errorf(
							"torrent with name '%s' not compatable existing infohash: %s, new: %s",
							t.Name(), t.InfoHash().HexString(), spec.InfoHash.HexString(),
						)
					}
				}
			}
		}

		t, _ = s.c.AddTorrentOpt(torrent.AddTorrentOpts{
			InfoHash:  spec.InfoHash,
			Storage:   s.Storage,
			InfoBytes: infoBytes,
			ChunkSize: spec.ChunkSize,
		})
		t.AllowDataDownload()
		t.AllowDataUpload()
		t.DownloadAll()

		select {
		case <-ctx.Done():
			return nil, fmt.Errorf("creating torrent timed out")
		case <-t.GotInfo():
			err := s.infoBytes.Set(t.InfoHash(), t.Metainfo())
			if err != nil {
				s.log.Error("error setting info bytes for torrent %s: %s", t.Name(), err.Error())
			}
			for _, f := range t.Files() {
				f.SetPriority(s.DefaultPriority)
			}

		}
	}

	return t, nil
}

func isTorrentCompatable(existingInfo, newInfo metainfo.Info) bool {
	existingFiles := slices.Clone(existingInfo.Files)
	newFiles := slices.Clone(newInfo.Files)

	pathCmp := func(a, b metainfo.FileInfo) int {
		return slices.Compare(a.BestPath(), b.BestPath())
	}
	slices.SortStableFunc(existingFiles, pathCmp)
	slices.SortStableFunc(newFiles, pathCmp)

	// torrents basically equals
	if slices.EqualFunc(existingFiles, newFiles, func(fi1, fi2 metainfo.FileInfo) bool {
		return fi1.Length == fi2.Length && slices.Equal(fi1.BestPath(), fi1.BestPath())
	}) {
		return true
	}

	if len(newFiles) > len(existingFiles) {
		all := append(existingFiles, newFiles...)
		slices.SortStableFunc(all, pathCmp)
		slices.CompactFunc(all, func(fi1, fi2 metainfo.FileInfo) bool {
			return slices.Equal(fi1.BestPath(), fi2.BestPath()) && fi1.Length == fi2.Length
		})
	}

	return false
}

func isValidInfoHashBytes(d []byte) bool {
	var info metainfo.Info
	err := bencode.Unmarshal(d, &info)
	return err == nil
}

func (s *Service) NewTorrentFs(f vfs.File) (vfs.Filesystem, error) {
	ctx, cancel := context.WithTimeout(context.TODO(), time.Second*time.Duration(s.addTimeout))
	defer cancel()
	defer f.Close()

	t, err := s.AddTorrent(ctx, f)
	if err != nil {
		return nil, err
	}

	return vfs.NewTorrentFs(controller.NewTorrent(t, s.excludedFiles), s.readTimeout), nil
}

func (s *Service) Stats() (*Stats, error) {
	return &Stats{}, nil
}

func (s *Service) GetStats() torrent.ConnStats {
	return s.c.ConnStats()
}

func (s *Service) loadTorrentFiles(ctx context.Context) error {
	return filepath.Walk(s.SourceDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return fmt.Errorf("fs walk error: %w", err)
		}

		if ctx.Err() != nil {
			return ctx.Err()
		}

		if info.IsDir() {
			return nil
		}

		if strings.HasSuffix(path, ".torrent") {
			file := vfs.NewLazyOsFile(path)
			defer file.Close()

			_, err = s.AddTorrent(ctx, file)
			if err != nil {
				s.log.Error("failed adding torrent", "error", err)
			}
		}

		return nil
	})
}

func (s *Service) ListTorrents(ctx context.Context) ([]*controller.Torrent, error) {
	<-s.torrentLoaded

	out := []*controller.Torrent{}
	for _, v := range s.c.Torrents() {
		out = append(out, controller.NewTorrent(v, s.excludedFiles))
	}
	return out, nil
}

func (s *Service) GetTorrent(infohashHex string) (*controller.Torrent, error) {
	<-s.torrentLoaded

	t, ok := s.c.Torrent(infohash.FromHexString(infohashHex))
	if !ok {
		return nil, nil
	}

	return controller.NewTorrent(t, s.excludedFiles), nil
}