tstor/src/host/datastorage/storage.go

package datastorage

import (
	"context"
	"log/slog"
	"os"
	"path"
	"path/filepath"
	"slices"

	"git.kmsign.ru/royalcat/tstor/src/host/controller"
	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/storage"
)

type DataStorage interface {
	storage.ClientImplCloser
	DeleteFile(file *torrent.File) error
	CleanupDirs(ctx context.Context, expected []*controller.Torrent, dryRun bool) (int, error)
	CleanupFiles(ctx context.Context, expected []*controller.Torrent, dryRun bool) (int, error)
}

// NewFileStorage creates a new ClientImplCloser that stores files using the OS native filesystem.
func NewFileStorage(baseDir string, pc storage.PieceCompletion) *FileStorage {
	return &FileStorage{
		baseDir: baseDir,
		ClientImplCloser: storage.NewFileOpts(storage.NewFileClientOpts{
			ClientBaseDir:   baseDir,
			PieceCompletion: pc,
			TorrentDirMaker: torrentDir,
			FilePathMaker:   filePath,
		}),
		pieceCompletion: pc,
		log:             slog.With("component", "torrent-client"),
	}
}

// File-based storage for torrents, that isn't yet bound to a particular torrent.
type FileStorage struct {
	baseDir string
	storage.ClientImplCloser
	pieceCompletion storage.PieceCompletion
	log             *slog.Logger
}

func (me *FileStorage) Close() error {
	return me.pieceCompletion.Close()
}

func torrentDir(baseDir string, info *metainfo.Info, infoHash metainfo.Hash) string {
	dirName := info.Name
	if dirName == "" {
		dirName = infoHash.HexString()
	}

	return filepath.Join(baseDir, dirName)
}

func filePath(opts storage.FilePathMakerOpts) string {
	return filepath.Join(opts.File.Path...)
}

func (fs *FileStorage) filePath(info *metainfo.Info, infoHash metainfo.Hash, fileInfo *metainfo.FileInfo) string {
	return filepath.Join(torrentDir(fs.baseDir, info, infoHash), filePath(storage.FilePathMakerOpts{
		Info: info,
		File: fileInfo,
	}))
}

func (fs *FileStorage) DeleteFile(file *torrent.File) error {
	info := file.Torrent().Info()
	infoHash := file.Torrent().InfoHash()
	torrentDir := torrentDir(fs.baseDir, info, infoHash)
	fileInfo := file.FileInfo()
	relFilePath := filePath(storage.FilePathMakerOpts{
		Info: info,
		File: &fileInfo,
	})
	filePath := path.Join(torrentDir, relFilePath)
	for i := file.BeginPieceIndex(); i < file.EndPieceIndex(); i++ {
		pk := metainfo.PieceKey{InfoHash: infoHash, Index: i}
		err := fs.pieceCompletion.Set(pk, false)
		if err != nil {
			return err
		}
	}
	return os.Remove(filePath)
}

func (fs *FileStorage) CleanupDirs(ctx context.Context, expected []*controller.Torrent, dryRun bool) (int, error) {
	log := fs.log.With("function", "CleanupDirs", "expectedTorrents", len(expected), "dryRun", dryRun)

	expectedEntries := []string{}
	for _, e := range expected {
		expectedEntries = append(expectedEntries, e.Torrent().Name())
	}

	entries, err := os.ReadDir(fs.baseDir)
	if err != nil {
		return 0, err
	}

	toDelete := []string{}
	for _, v := range entries {
		if !slices.Contains(expectedEntries, v.Name()) {
			toDelete = append(toDelete, v.Name())
		}
	}

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

	log.Info("deleting trash data", "dirsCount", len(toDelete))
	if !dryRun {
		for i, name := range toDelete {
			p := path.Join(fs.baseDir, name)
			log.Warn("deleting trash data", "path", p)
			err := os.RemoveAll(p)
			if err != nil {
				return i, err
			}
		}
	}

	return len(toDelete), nil
}

// func (fs *FileStorage) IsCompatable(ctx context.Context, addition *controller.Torrent, dryRun bool) (bool, error) {
// 	log := fs.log.With("function", "IsCompatable", "addition", addition.Name())

// 	ifp
// }

func (fs *FileStorage) CleanupFiles(ctx context.Context, expected []*controller.Torrent, dryRun bool) (int, error) {
	log := fs.log.With("function", "CleanupFiles", "expectedTorrents", len(expected), "dryRun", dryRun)

	expectedEntries := []string{}
	{
		for _, e := range expected {
			files, err := e.Files(ctx)
			if err != nil {
				return 0, err
			}

			for _, f := range files {
				expectedEntries = append(expectedEntries, fs.filePath(e.Torrent().Info(), e.Torrent().InfoHash(), ptr(f.FileInfo())))
			}
		}
	}

	entries := []string{}
	err := filepath.Walk(fs.baseDir,
		func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if ctx.Err() != nil {
				return ctx.Err()
			}

			if info.IsDir() {
				return nil
			}
			entries = append(entries, path)
			return nil
		})
	if err != nil {
		return 0, err
	}

	toDelete := []string{}
	for _, v := range entries {
		if !slices.Contains(expectedEntries, v) {
			toDelete = append(toDelete, v)
		}
	}

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

	log.Info("deleting trash data", "filesCount", len(toDelete))
	if !dryRun {
		for i, p := range toDelete {
			fs.log.Warn("deleting trash data", "path", p)
			err := os.Remove(p)
			if err != nil {
				return i, err
			}
		}
	}
	return len(toDelete), nil
}

func ptr[D any](v D) *D {
	return &v
}