tstor/src/host/vfs/torrent.go

package vfs

import (
	"context"
	"io"
	"io/fs"
	"path"
	"slices"
	"strings"
	"sync"

	"git.kmsign.ru/royalcat/tstor/pkg/ctxio"
	"git.kmsign.ru/royalcat/tstor/src/host/controller"
	"github.com/anacrolix/torrent"
	"golang.org/x/exp/maps"
)

var _ Filesystem = &TorrentFs{}

type TorrentFs struct {
	name string

	mu      sync.Mutex
	Torrent *controller.Torrent

	filesCache map[string]File

	resolver *resolver
}

func NewTorrentFs(name string, c *controller.Torrent) *TorrentFs {
	return &TorrentFs{
		name:     name,
		Torrent:  c,
		resolver: newResolver(ArchiveFactories),
	}
}

var _ fs.DirEntry = (*TorrentFs)(nil)

// Name implements fs.DirEntry.
func (tfs *TorrentFs) Name() string {
	return tfs.name
}

// Info implements fs.DirEntry.
func (tfs *TorrentFs) Info() (fs.FileInfo, error) {
	return newDirInfo(tfs.name), nil
}

// IsDir implements fs.DirEntry.
func (tfs *TorrentFs) IsDir() bool {
	return true
}

// Type implements fs.DirEntry.
func (tfs *TorrentFs) Type() fs.FileMode {
	return fs.ModeDir
}

func (fs *TorrentFs) files(ctx context.Context) (map[string]File, error) {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	if fs.filesCache != nil {
		return fs.filesCache, nil
	}

	files, err := fs.Torrent.Files(context.Background())
	if err != nil {
		return nil, err
	}

	fs.filesCache = make(map[string]File)
	for _, file := range files {
		file.Download()
		p := AbsPath(file.Path())
		tf, err := openTorrentFile(ctx, path.Base(p), file)
		if err != nil {
			return nil, err
		}
		fs.filesCache[p] = tf
	}

	// TODO optional
	if len(fs.filesCache) == 1 && fs.resolver.isNestedFs(fs.Torrent.Name()) {
		filepath := "/" + fs.Torrent.Name()
		if file, ok := fs.filesCache[filepath]; ok {
			nestedFs, err := fs.resolver.nestedFs(ctx, filepath, file)
			if err != nil {
				return nil, err
			}
			if nestedFs == nil {
				goto DEFAULT_DIR // FIXME
			}
			fs.filesCache, err = listFilesRecursive(ctx, nestedFs, "/")
			if err != nil {
				return nil, err
			}

			return fs.filesCache, nil
		}

	}

DEFAULT_DIR:
	rootDir := "/" + fs.Torrent.Name() + "/"
	singleDir := true
	for k, _ := range fs.filesCache {
		if !strings.HasPrefix(k, rootDir) {
			singleDir = false
		}
	}
	if singleDir {
		for k, f := range fs.filesCache {
			delete(fs.filesCache, k)
			k, _ = strings.CutPrefix(k, rootDir)
			k = AbsPath(k)
			fs.filesCache[k] = f
		}
	}

	return fs.filesCache, nil
}

// func anyPeerHasFiles(file *torrent.File) bool {
// 	for _, conn := range file.Torrent().PeerConns() {
// 		if bitmapHaveFile(conn.PeerPieces(), file) {
// 			return true
// 		}
// 	}
// 	return false
// }

// func bitmapHaveFile(bitmap *roaring.Bitmap, file *torrent.File) bool {
// 	for i := file.BeginPieceIndex(); i < file.EndPieceIndex(); i++ {
// 		if !bitmap.ContainsInt(i) {
// 			return false
// 		}
// 	}
// 	return true
// }

func listFilesRecursive(ctx context.Context, vfs Filesystem, start string) (map[string]File, error) {
	out := make(map[string]File, 0)
	entries, err := vfs.ReadDir(ctx, start)
	if err != nil {
		return nil, err
	}
	for _, entry := range entries {
		filename := path.Join(start, entry.Name())
		if entry.IsDir() {
			rec, err := listFilesRecursive(ctx, vfs, filename)
			if err != nil {
				return nil, err
			}
			maps.Copy(out, rec)
		} else {
			file, err := vfs.Open(ctx, filename)
			if err != nil {
				return nil, err
			}
			out[filename] = file
		}
	}

	return out, nil
}

func (fs *TorrentFs) rawOpen(ctx context.Context, path string) (File, error) {
	files, err := fs.files(ctx)
	if err != nil {
		return nil, err
	}
	file, err := getFile(files, path)
	return file, err
}

func (fs *TorrentFs) rawStat(ctx context.Context, filename string) (fs.FileInfo, error) {
	files, err := fs.files(ctx)
	if err != nil {
		return nil, err
	}
	file, err := getFile(files, filename)
	if err != nil {
		return nil, err
	}
	return file.Stat()
}

// Stat implements Filesystem.
func (fs *TorrentFs) Stat(ctx context.Context, filename string) (fs.FileInfo, error) {
	if filename == Separator {
		return newDirInfo(filename), nil
	}

	fsPath, nestedFs, nestedFsPath, err := fs.resolver.resolvePath(ctx, filename, fs.rawOpen)
	if err != nil {
		return nil, err
	}
	if nestedFs != nil {
		return nestedFs.Stat(ctx, nestedFsPath)
	}

	return fs.rawStat(ctx, fsPath)
}

func (fs *TorrentFs) Open(ctx context.Context, filename string) (File, error) {
	fsPath, nestedFs, nestedFsPath, err := fs.resolver.resolvePath(ctx, filename, fs.rawOpen)
	if err != nil {
		return nil, err
	}
	if nestedFs != nil {
		return nestedFs.Open(ctx, nestedFsPath)
	}

	return fs.rawOpen(ctx, fsPath)
}

func (fs *TorrentFs) ReadDir(ctx context.Context, name string) ([]fs.DirEntry, error) {
	fsPath, nestedFs, nestedFsPath, err := fs.resolver.resolvePath(ctx, name, fs.rawOpen)
	if err != nil {
		return nil, err
	}
	if nestedFs != nil {
		return nestedFs.ReadDir(ctx, nestedFsPath)
	}
	files, err := fs.files(ctx)
	if err != nil {
		return nil, err
	}

	return listDirFromFiles(files, fsPath)
}

func (fs *TorrentFs) Unlink(ctx context.Context, name string) error {
	name = AbsPath(name)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	files, err := fs.files(ctx)
	if err != nil {
		return err
	}

	if !slices.Contains(maps.Keys(files), name) {
		return ErrNotExist
	}

	file := files[name]
	delete(fs.filesCache, name)

	tfile, ok := file.(*torrentFile)
	if !ok {
		return ErrNotImplemented
	}

	return fs.Torrent.ExcludeFile(context.Background(), tfile.file)
}

var _ File = &torrentFile{}

type torrentFile struct {
	name string

	mu sync.Mutex

	tr torrent.Reader

	file *torrent.File
}

func openTorrentFile(ctx context.Context, name string, file *torrent.File) (*torrentFile, error) {
	select {
	case <-file.Torrent().GotInfo():
		break
	case <-ctx.Done():
		return nil, ctx.Err()
	}

	r := file.NewReader()
	r.SetReadahead(4096) // TODO configurable
	r.SetResponsive()

	return &torrentFile{
		name: name,
		tr:   r,
		file: file,
	}, nil
}

func (tf *torrentFile) Stat() (fs.FileInfo, error) {
	return newFileInfo(tf.name, tf.file.Length()), nil
}

func (tf *torrentFile) Size() int64 {
	return tf.file.Length()
}

func (tf *torrentFile) IsDir() bool {
	return false
}

func (rw *torrentFile) Close(ctx context.Context) error {
	rw.mu.Lock()
	defer rw.mu.Unlock()

	return rw.tr.Close()
}

// Read implements ctxio.Reader.
func (tf *torrentFile) Read(ctx context.Context, p []byte) (n int, err error) {
	tf.mu.Lock()
	defer tf.mu.Unlock()

	return tf.tr.ReadContext(ctx, p)
}

func (yf *torrentFile) ReadAt(ctx context.Context, p []byte, off int64) (int, error) {
	yf.mu.Lock()
	defer yf.mu.Unlock()

	_, err := yf.tr.Seek(off, io.SeekStart)
	if err != nil {
		return 0, err
	}

	return readAtLeast(ctx, yf, p, len(p))
}

func readAtLeast(ctx context.Context, r ctxio.Reader, buf []byte, min int) (n int, err error) {
	if len(buf) < min {
		return 0, io.ErrShortBuffer
	}
	for n < min && err == nil {
		var nn int

		nn, err = r.Read(ctx, buf[n:])
		n += nn
	}
	if n >= min {
		err = nil
	} else if n > 0 && err == io.EOF {
		err = io.ErrUnexpectedEOF
	}
	return
}