tstor/src/vfs/archive.go

package vfs

import (
	"archive/zip"
	"context"
	"fmt"
	"io"
	"io/fs"
	"path"
	"strings"
	"sync"
	"time"

	"git.kmsign.ru/royalcat/tstor/pkg/ioutils"
	"github.com/bodgit/sevenzip"
	"github.com/nwaples/rardecode/v2"
	"github.com/royalcat/ctxio"
)

var ArchiveFactories = map[string]FsFactory{
	".zip": func(ctx context.Context, sourcePath string, f File) (Filesystem, error) {
		stat, err := f.Info()
		if err != nil {
			return nil, err
		}
		return NewArchive(ctx, sourcePath, stat.Name(), f, stat.Size(), ZipLoader)
	},
	".rar": func(ctx context.Context, sourcePath string, f File) (Filesystem, error) {
		stat, err := f.Info()
		if err != nil {
			return nil, err
		}
		return NewArchive(ctx, sourcePath, stat.Name(), f, stat.Size(), RarLoader)
	},
	".7z": func(ctx context.Context, sourcePath string, f File) (Filesystem, error) {
		stat, err := f.Info()
		if err != nil {
			return nil, err
		}
		return NewArchive(ctx, sourcePath, stat.Name(), f, stat.Size(), SevenZipLoader)
	},
}

type archiveLoader func(ctx context.Context, archivePath string, r File, size int64) (map[string]fileEntry, error)

var _ Filesystem = &ArchiveFS{}

type fileEntry struct {
	fs.FileInfo
	open func(ctx context.Context) (File, error)
}

type ArchiveFS struct {
	name string

	size int64

	files map[string]fileEntry
}

// Rename implements Filesystem.
func (a *ArchiveFS) Rename(ctx context.Context, oldpath string, newpath string) error {
	return ErrNotImplemented
}

// ModTime implements Filesystem.
func (a *ArchiveFS) ModTime() time.Time {
	return time.Time{}
}

// Mode implements Filesystem.
func (a *ArchiveFS) Mode() fs.FileMode {
	return fs.ModeDir
}

// Size implements Filesystem.
func (a *ArchiveFS) Size() int64 {
	return int64(a.size)
}

// Sys implements Filesystem.
func (a *ArchiveFS) Sys() any {
	return nil
}

// FsName implements Filesystem.
func (a *ArchiveFS) FsName() string {
	return "archivefs"
}

func NewArchive(ctx context.Context, archivePath, name string, f File, size int64, loader archiveLoader) (*ArchiveFS, error) {
	archiveFiles, err := loader(ctx, archivePath, f, size)
	if err != nil {
		return nil, err
	}

	// TODO make optional
	singleDir := true
	for k := range archiveFiles {
		if !strings.HasPrefix(k, "/"+name+"/") {
			singleDir = false
			break
		}
	}

	files := make(map[string]fileEntry, len(archiveFiles))
	for k, v := range archiveFiles {
		// TODO make optional
		if strings.Contains(k, "/__MACOSX/") {
			continue
		}

		if singleDir {
			k, _ = strings.CutPrefix(k, "/"+name)
		}

		files[k] = v
	}

	// FIXME configurable
	files["/.forcegallery"] = fileEntry{
		FileInfo: NewFileInfo("/.forcegallery", 0),
		open: func(ctx context.Context) (File, error) {
			return NewMemoryFile(".forcegallery", []byte{}), nil
		},
	}

	return &ArchiveFS{
		name:  name,
		size:  size,
		files: files,
	}, nil
}

// Unlink implements Filesystem.
func (a *ArchiveFS) Unlink(ctx context.Context, filename string) error {
	return ErrNotImplemented
}

func (a *ArchiveFS) Open(ctx context.Context, filename string) (File, error) {
	if filename == Separator {
		return NewDirFile(filename), nil
	}

	f, ok := a.files[filename]
	if ok {
		return f.open(ctx)
	}

	for p := range a.files {
		if strings.HasPrefix(p, filename) {
			return NewDirFile(filename), nil
		}
	}

	return nil, ErrNotExist
}

func (a *ArchiveFS) ReadDir(ctx context.Context, path string) ([]fs.DirEntry, error) {
	infos := make(map[string]fs.FileInfo, len(a.files))
	for k, v := range a.files {
		infos[k] = v
	}

	return ListDirFromInfo(infos, path)
}

// Stat implements Filesystem.
func (afs *ArchiveFS) Stat(ctx context.Context, filename string) (fs.FileInfo, error) {
	if entry, ok := afs.files[filename]; ok {
		return entry, nil
	}

	for p, _ := range afs.files {
		if strings.HasPrefix(p, filename) {
			return NewDirInfo(path.Base(filename)), nil
		}
	}

	return nil, ErrNotExist
}

// Info implements Filesystem.
func (a *ArchiveFS) Info() (fs.FileInfo, error) {
	return a, nil
}

// IsDir implements Filesystem.
func (a *ArchiveFS) IsDir() bool {
	return true
}

// Name implements Filesystem.
func (a *ArchiveFS) Name() string {
	return a.name
}

// Type implements Filesystem.
func (a *ArchiveFS) Type() fs.FileMode {
	return fs.ModeDir
}

var _ File = (*archiveFile)(nil)

func newArchiveFile(name string, size int64, rr *randomReaderFromLinear) *archiveFile {
	return &archiveFile{
		name: name,
		size: size,
		rr:   rr,
	}
}

type archiveFile struct {
	name string
	size int64

	m      sync.Mutex
	offset int64

	rr *randomReaderFromLinear
}

// Seek implements File.
func (d *archiveFile) Seek(offset int64, whence int) (int64, error) {
	switch whence {
	case io.SeekStart:
		d.offset = offset

	case io.SeekCurrent:
		d.offset += offset
	case io.SeekEnd:
		d.offset = d.size + offset
	}
	return d.offset, nil
}

// Name implements File.
func (d *archiveFile) Name() string {
	return d.name
}

// Type implements File.
func (d *archiveFile) Type() fs.FileMode {
	return ModeFileRO
}

func (d *archiveFile) Info() (fs.FileInfo, error) {
	return NewFileInfo(d.name, d.size), nil
}

func (d *archiveFile) Size() int64 {
	return d.size
}

func (d *archiveFile) IsDir() bool {
	return false
}

func (d *archiveFile) Read(ctx context.Context, p []byte) (n int, err error) {
	ctx, span := tracer.Start(ctx, "archive.File.Read")
	defer span.End()

	n, err = d.rr.ReadAt(ctx, p, d.offset)
	d.offset += int64(n)
	return n, err
}

func (d *archiveFile) ReadAt(ctx context.Context, p []byte, off int64) (n int, err error) {
	d.m.Lock()
	defer d.m.Unlock()

	return d.rr.ReadAt(ctx, p, off)
}

func (d *archiveFile) Close(ctx context.Context) error {
	// FIXME close should do nothing as archive fs currently reuse the same file instances
	return nil
}

type archiveFileReaderFactory func(ctx context.Context) (ctxio.ReadCloser, error)

var _ archiveLoader = ZipLoader

func ZipLoader(ctx context.Context, archivePath string, f File, size int64) (map[string]fileEntry, error) {
	hash, err := FileHash(ctx, f)
	if err != nil {
		return nil, err
	}

	reader := ctxio.IoReaderAt(ctx, f)
	zr, err := zip.NewReader(reader, size)
	if err != nil {
		return nil, err
	}

	out := make(map[string]fileEntry)
	for i := range zr.File {
		zipFile := zr.File[i]
		if zipFile.FileInfo().IsDir() {
			continue
		}

		i := i
		af := func(ctx context.Context) (ctxio.ReadCloser, error) {
			reader := ctxio.IoReaderAt(ctx, f)

			zr, err := zip.NewReader(reader, size)
			if err != nil {
				return nil, fmt.Errorf("failed to create zip reader: %w", err)
			}

			rc, err := zr.File[i].Open()
			if err != nil {
				return nil, fmt.Errorf("failed to open file in zip archive: %w", err)
			}

			return ctxio.WrapIoReadCloser(rc), nil
		}

		info := zipFile.FileInfo()

		rr := newRandomReaderFromLinear(archiveFileIndex{archiveHash: hash, filename: zipFile.Name}, info.Size(), af)

		out[AbsPath(zipFile.Name)] = fileEntry{
			FileInfo: info,
			open: func(ctx context.Context) (File, error) {
				return newArchiveFile(info.Name(), info.Size(), rr), nil
			},
		}
	}

	return out, nil
}

var _ archiveLoader = SevenZipLoader

func SevenZipLoader(ctx context.Context, archivePath string, ctxreader File, size int64) (map[string]fileEntry, error) {
	hash, err := FileHash(ctx, ctxreader)
	if err != nil {
		return nil, err
	}

	reader := ctxio.IoReaderAt(ctx, ctxreader)
	r, err := sevenzip.NewReader(reader, size)
	if err != nil {
		return nil, err
	}

	out := make(map[string]fileEntry)
	for i, f := range r.File {
		f := f
		if f.FileInfo().IsDir() {
			continue
		}

		i := i
		af := func(ctx context.Context) (ctxio.ReadCloser, error) {
			reader := ctxio.IoReaderAt(ctx, ctxreader)
			zr, err := sevenzip.NewReader(reader, size)
			if err != nil {
				return nil, err
			}

			rc, err := zr.File[i].Open()
			if err != nil {
				return nil, err
			}

			return ctxio.WrapIoReadCloser(rc), nil
		}

		info := f.FileInfo()

		rr := newRandomReaderFromLinear(archiveFileIndex{archiveHash: hash, filename: f.Name}, info.Size(), af)

		out[AbsPath(f.Name)] = fileEntry{
			FileInfo: f.FileInfo(),
			open: func(ctx context.Context) (File, error) {
				return newArchiveFile(info.Name(), info.Size(), rr), nil
			},
		}
	}

	return out, nil
}

var _ archiveLoader = RarLoader

func RarLoader(ctx context.Context, archivePath string, f File, size int64) (map[string]fileEntry, error) {
	hash, err := FileHash(ctx, f)
	if err != nil {
		return nil, err
	}

	reader := ioutils.WrapIoReadSeeker(ctx, f, size)

	r, err := rardecode.NewReader(reader)
	if err != nil {
		return nil, err
	}

	out := make(map[string]fileEntry)
	for {
		header, err := r.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}

		name := header.Name
		af := func(ctx context.Context) (ctxio.ReadCloser, error) {
			reader := ioutils.WrapIoReadSeeker(ctx, f, size)
			r, err := rardecode.NewReader(reader)
			if err != nil {
				return nil, err
			}

			for header, err := r.Next(); err != io.EOF; header, err = r.Next() {
				if err != nil {
					return nil, err
				}
				if header.Name == name {
					return ctxio.NopCloser(ctxio.WrapIoReader(r)), nil
				}
			}
			return nil, fmt.Errorf("file with name '%s' not found", name)
		}

		rr := newRandomReaderFromLinear(archiveFileIndex{archiveHash: hash, filename: header.Name}, header.UnPackedSize, af)

		out[AbsPath(header.Name)] = fileEntry{
			FileInfo: NewFileInfo(header.Name, header.UnPackedSize),
			open: func(ctx context.Context) (File, error) {
				return newArchiveFile(header.Name, header.UnPackedSize, rr), nil
			},
		}
	}

	return out, nil
}