tstor/src/host/service/service.go

package service

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

	"git.kmsign.ru/royalcat/tstor/src/config"
	"git.kmsign.ru/royalcat/tstor/src/host/controller"
	"git.kmsign.ru/royalcat/tstor/src/host/datastorage"
	"git.kmsign.ru/royalcat/tstor/src/host/store"
	"git.kmsign.ru/royalcat/tstor/src/host/vfs"
	"go.uber.org/multierr"
	"golang.org/x/exp/maps"

	"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"
	"github.com/royalcat/kv"
)

type DirAquire struct {
	Name   string
	Hashes []infohash.T
}

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

	torrentLoaded chan struct{}

	// stats *Stats
	DefaultPriority types.PiecePriority
	Storage         datastorage.DataStorage
	SourceDir       string

	dirsAquire kv.Store[string, DirAquire]

	log                     *slog.Logger
	addTimeout, readTimeout int
}

func NewService(sourceDir string, cfg config.TorrentClient, c *torrent.Client,
	storage datastorage.DataStorage, excludedFiles *store.FilesMappings, infoBytes *store.InfoBytes,
	addTimeout, readTimeout int,
) (*Service, error) {
	dirsAcquire, err := kv.NewBadgerKV[string, DirAquire](path.Join(cfg.MetadataFolder, "dir-acquire"))
	if err != nil {
		return nil, err
	}

	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{}),
		dirsAquire:      dirsAcquire,
		// 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, nil
}

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

func (s *Service) Close() error {
	err := multierr.Combine(s.c.Close()...)
	err = multierr.Append(err, s.Storage.Close())
	return err
}

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 {
			compatable, _, err := s.checkTorrentCompatable(ctx, spec.InfoHash, info)
			if err != nil {
				return nil, err
			}
			if !compatable {
				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()

		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 (s *Service) checkTorrentCompatable(ctx context.Context, ih infohash.T, info metainfo.Info) (compatable bool, tryLater bool, err error) {
	log := s.log.With("new-name", info.BestName(), "new-infohash", ih.String())

	name := info.BestName()

	aq, found, err := s.dirsAquire.Get(ctx, info.BestName())
	if err != nil {
		return false, false, err
	}
	if !found {
		err = s.dirsAquire.Set(ctx, name, DirAquire{
			Name:   name,
			Hashes: slices.Compact([]infohash.T{ih}),
		})
		if err != nil {
			return false, false, err
		}

		log.Debug("acquiring was not found, so created")
		return true, false, nil
	}

	if slices.Contains(aq.Hashes, ih) {
		log.Debug("hash already know to be compatable")
		return true, false, nil
	}

	for _, existingTorrent := range s.c.Torrents() {
		if existingTorrent.Name() != name || existingTorrent.InfoHash() == ih {
			continue
		}

		existingInfo := existingTorrent.Info()

		existingFiles := slices.Clone(existingInfo.Files)
		newFiles := slices.Clone(info.Files)

		if !s.checkTorrentFilesCompatable(aq, existingFiles, newFiles) {
			return false, false, nil
		}

		aq.Hashes = slicesUnique(append(aq.Hashes, ih))
		err = s.dirsAquire.Set(ctx, aq.Name, aq)
		if err != nil {
			log.Warn("torrent not compatible")
			return false, false, err
		}

	}

	if slices.Contains(aq.Hashes, ih) {
		log.Debug("hash is compatable")
		return true, false, nil
	}

	log.Debug("torrent with same name not found, try later")
	return false, true, nil
}

func (s *Service) checkTorrentFilesCompatable(aq DirAquire, existingFiles, newFiles []metainfo.FileInfo) bool {
	log := s.log.With("name", aq.Name)

	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) {
		type fileInfo struct {
			Path   string
			Length int64
		}
		mapInfo := func(fi metainfo.FileInfo) fileInfo {
			return fileInfo{
				Path:   strings.Join(fi.BestPath(), "/"),
				Length: fi.Length,
			}
		}

		existingFiles := apply(existingFiles, mapInfo)
		newFiles := apply(newFiles, mapInfo)

		for _, n := range newFiles {
			if slices.Contains(existingFiles, n) {
				continue
			}

			for _, e := range existingFiles {
				if e.Path == n.Path && e.Length != n.Length {
					log.Warn("torrents not compatible, has files with different length", "path", n.Path, "existing-length", e.Length, "new-length", e.Length)
					return false
				}
			}
		}
	}

	return true
}

func (s *Service) getTorrentsByName(name string) []*torrent.Torrent {
	out := []*torrent.Torrent{}
	for _, t := range s.c.Torrents() {
		if t.Name() == name {
			out = append(out, t)
		}
	}
	return out
}

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()

	info, err := f.Stat()
	if err != nil {
		return nil, err
	}

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

	return vfs.NewTorrentFs(info.Name(), 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
}

func slicesUnique[S ~[]E, E comparable](in S) S {
	m := map[E]struct{}{}
	for _, v := range in {
		m[v] = struct{}{}
	}

	return maps.Keys(m)
}

func apply[I, O any](in []I, f func(e I) O) []O {
	out := []O{}
	for _, v := range in {
		out = append(out, f(v))
	}
	return out
}