package torrent import ( "context" "log/slog" "slices" "strings" "git.kmsign.ru/royalcat/tstor/pkg/rlog" "github.com/anacrolix/torrent" "github.com/anacrolix/torrent/types" "github.com/royalcat/kv" ) type TorrentFileDeleter interface { DeleteFile(file *torrent.File) error } type FileProperties struct { Excluded bool `json:"excluded"` Priority types.PiecePriority `json:"priority"` } type Controller struct { torrentFilePath string t *torrent.Torrent storage TorrentFileDeleter fileProperties kv.Store[string, FileProperties] log *rlog.Logger } func newController(t *torrent.Torrent, fileProperties kv.Store[string, FileProperties], storage TorrentFileDeleter) *Controller { return &Controller{ t: t, storage: storage, fileProperties: fileProperties, log: rlog.Component("torrent-client", "controller").With(slog.String("infohash", t.InfoHash().HexString())), } } func (s *Controller) TorrentFilePath() string { return s.torrentFilePath } func (s *Controller) Torrent() *torrent.Torrent { return s.t } func (c *Controller) Name() string { <-c.t.GotInfo() if name := c.t.Name(); name != "" { return name } return c.InfoHash() } func (s *Controller) InfoHash() string { <-s.t.GotInfo() return s.t.InfoHash().HexString() } func (s *Controller) BytesCompleted() int64 { <-s.t.GotInfo() return s.t.BytesCompleted() } func (s *Controller) BytesMissing() int64 { <-s.t.GotInfo() return s.t.BytesMissing() } func (s *Controller) Length() int64 { <-s.t.GotInfo() return s.t.Length() } func (s *Controller) Files(ctx context.Context) ([]*torrent.File, error) { fps := map[string]FileProperties{} err := s.fileProperties.Range(ctx, func(k string, v FileProperties) error { fps[k] = v return nil }) if err != nil { return nil, err } select { case <-ctx.Done(): return nil, ctx.Err() case <-s.t.GotInfo(): } files := s.t.Files() files = slices.DeleteFunc(files, func(file *torrent.File) bool { if file == nil { return true } p := file.Path() if strings.Contains(p, "/.pad/") { return true } if props, ok := fps[p]; ok && props.Excluded { return true } return false }) return files, nil } func Map[T, U any](ts []T, f func(T) U) []U { us := make([]U, len(ts)) for i := range ts { us[i] = f(ts[i]) } return us } func (s *Controller) ExcludeFile(ctx context.Context, f *torrent.File) error { log := s.log.With(slog.String("file", f.Path())) log.Info(ctx, "excluding file") return s.fileProperties.Edit(ctx, f.Path(), func(ctx context.Context, v FileProperties) (FileProperties, error) { v.Excluded = true return v, nil }) } func (s *Controller) isFileComplete(startIndex int, endIndex int) bool { for i := startIndex; i < endIndex; i++ { if !s.t.Piece(i).State().Complete { return false } } return true } func (s *Controller) ValidateTorrent(ctx context.Context) error { select { case <-ctx.Done(): return ctx.Err() case <-s.t.GotInfo(): } for i := 0; i < s.t.NumPieces(); i++ { if ctx.Err() != nil { return ctx.Err() } s.t.Piece(i).VerifyData() } return nil } func (c *Controller) SetFilePriority(ctx context.Context, file *torrent.File, priority types.PiecePriority) error { log := c.log.With(slog.String("file", file.Path()), slog.Int("priority", int(priority))) log.Info(ctx, "set pritority for file") err := c.fileProperties.Edit(ctx, file.Path(), func(ctx context.Context, v FileProperties) (FileProperties, error) { v.Priority = priority return v, nil }) if err != nil { return err } file.SetPriority(priority) return nil } func (c *Controller) initializeTorrentPriories(ctx context.Context) error { log := c.log.WithComponent("initializeTorrentPriories") files, err := c.Files(ctx) if err != nil { return err } for _, file := range files { if file == nil { continue } props, err := c.fileProperties.Get(ctx, file.Path()) if err != nil { if err == kv.ErrKeyNotFound { continue } log.Error(ctx, "failed to get file priority", rlog.Error(err)) } file.SetPriority(props.Priority) } log.Info(ctx, "torrent initialization complete", slog.String("infohash", c.InfoHash()), slog.String("torrent_name", c.Name())) return nil }