package torrent import ( "context" "log/slog" "strings" "git.kmsign.ru/royalcat/tstor/pkg/kvsingle" "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, torrentFileProperties kv.Store[string, FileProperties], storage TorrentFileDeleter) *Controller { return &Controller{ t: t, storage: storage, fileProperties: torrentFileProperties, 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) ([]*FileController, error) { ctx, span := tracer.Start(ctx, "Files") defer span.End() 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 := make([]*FileController, 0) for _, v := range s.t.Files() { if strings.Contains(v.Path(), "/.pad/") { continue } props := kvsingle.New(s.fileProperties, v.Path()) ctl := NewFileController(v, props) files = append(files, ctl) } return files, nil } func (s *Controller) GetFile(ctx context.Context, file string) (*FileController, error) { files, err := s.Files(ctx) if err != nil { return nil, err } for _, v := range files { if v.Path() == file { return v, nil } } return nil, 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") err := s.fileProperties.Edit(ctx, f.Path(), func(ctx context.Context, v FileProperties) (FileProperties, error) { v.Excluded = true return v, nil }) if err == kv.ErrKeyNotFound { err := s.fileProperties.Set(ctx, f.Path(), FileProperties{Excluded: true}) if err != nil { return err } } else if err != nil { return err } return s.storage.DeleteFile(f) } 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) SetPriority(ctx context.Context, priority types.PiecePriority) error { // log := c.log.With(slog.Int("priority", int(priority))) for _, f := range c.t.Files() { err := c.setFilePriority(ctx, f, priority) if err != nil { return err } } return nil } const defaultPriority = types.PiecePriorityNone func (c *Controller) Priority(ctx context.Context) (types.PiecePriority, error) { prio := defaultPriority files, err := c.Files(ctx) if err != nil { return 0, err } for _, v := range files { props, err := v.Properties(ctx) if err != nil { return 0, err } if props.Priority > prio { prio = props.Priority } } return prio, nil } func (c *Controller) setFilePriority(ctx context.Context, file *torrent.File, priority types.PiecePriority) error { err := c.fileProperties.Edit(ctx, file.Path(), func(ctx context.Context, v FileProperties) (FileProperties, error) { v.Priority = priority return v, nil }) if err == kv.ErrKeyNotFound { seterr := c.fileProperties.Set(ctx, file.Path(), FileProperties{Priority: priority}) if seterr != nil { return seterr } err = nil } if err != nil { return err } file.SetPriority(priority) return nil } func (c *Controller) initializeTorrentPriories(ctx context.Context) error { ctx, span := tracer.Start(ctx, "initializeTorrentPriories") defer span.End() log := c.log files, err := c.Files(ctx) if err != nil { return err } for _, file := range files { props, err := file.Properties(ctx) if err != nil { log.Error(ctx, "failed to get file properties", rlog.Error(err)) continue } log = log.With(slog.Int("priority", int(props.Priority))) file.file.SetPriority(props.Priority) } log.Info(ctx, "torrent initialization complete", slog.String("infohash", c.InfoHash()), slog.String("torrent_name", c.Name())) return nil }