tstor/pkg/ctxbilly/uring.go

package ctxbilly

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	securejoin "github.com/cyphar/filepath-securejoin"
	"github.com/iceber/iouring-go"
)

func NewURingFS() (*UringFS, error) {
	ur, err := iouring.New(64, iouring.WithAsync())
	if err != nil {
		return nil, err
	}

	return &UringFS{
		ur: ur,
	}, nil
}

var _ Filesystem = (*UringFS)(nil)

const (
	defaultDirectoryMode = 0o755
	defaultCreateMode    = 0o666
)

// UringFS is a fs implementation based on the OS filesystem which is bound to
// a base dir.
// Prefer this fs implementation over ChrootOS.
//
// Behaviours of note:
//  1. Read and write operations can only be directed to files which descends
//     from the base dir.
//  2. Symlinks don't have their targets modified, and therefore can point
//     to locations outside the base dir or to non-existent paths.
//  3. Readlink and Lstat ensures that the link file is located within the base
//     dir, evaluating any symlinks that file or base dir may contain.
type UringFS struct {
	ur      *iouring.IOURing
	baseDir string
}

func newBoundOS(d string) *UringFS {
	return &UringFS{baseDir: d}
}

func (fs *UringFS) Create(ctx context.Context, filename string) (File, error) {
	return fs.OpenFile(ctx, filename, os.O_RDWR|os.O_CREATE|os.O_TRUNC, defaultCreateMode)
}

func (fs *UringFS) OpenFile(ctx context.Context, filename string, flag int, perm os.FileMode) (File, error) {
	fn, err := fs.abs(filename)
	if err != nil {
		return nil, err
	}

	f, err := os.OpenFile(fn, flag, perm)
	if err != nil {
		return nil, err
	}

	return newFile(fs.ur, f)
}

func (fs *UringFS) ReadDir(ctx context.Context, path string) ([]os.FileInfo, error) {
	dir, err := fs.abs(path)
	if err != nil {
		return nil, err
	}

	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, err
	}
	infos := make([]os.FileInfo, 0, len(entries))
	for _, v := range entries {
		info, err := v.Info()
		if err != nil {
			return nil, err
		}

		infos = append(infos, info)
	}

	return infos, nil
}

func (fs *UringFS) Rename(ctx context.Context, from, to string) error {
	f, err := fs.abs(from)
	if err != nil {
		return err
	}
	t, err := fs.abs(to)
	if err != nil {
		return err
	}

	// MkdirAll for target name.
	if err := fs.createDir(t); err != nil {
		return err
	}

	return os.Rename(f, t)
}

func (fs *UringFS) MkdirAll(ctx context.Context, path string, perm os.FileMode) error {
	dir, err := fs.abs(path)
	if err != nil {
		return err
	}
	return os.MkdirAll(dir, perm)
}

func (fs *UringFS) Stat(ctx context.Context, filename string) (os.FileInfo, error) {
	filename, err := fs.abs(filename)
	if err != nil {
		return nil, err
	}
	return os.Stat(filename)
}

func (fs *UringFS) Remove(ctx context.Context, filename string) error {
	fn, err := fs.abs(filename)
	if err != nil {
		return err
	}
	return os.Remove(fn)
}

func (fs *UringFS) Join(elem ...string) string {
	return filepath.Join(elem...)
}

func (fs *UringFS) RemoveAll(path string) error {
	dir, err := fs.abs(path)
	if err != nil {
		return err
	}
	return os.RemoveAll(dir)
}

func (fs *UringFS) Symlink(ctx context.Context, target, link string) error {
	ln, err := fs.abs(link)
	if err != nil {
		return err
	}
	// MkdirAll for containing dir.
	if err := fs.createDir(ln); err != nil {
		return err
	}
	return os.Symlink(target, ln)
}

func (fs *UringFS) Lstat(ctx context.Context, filename string) (os.FileInfo, error) {
	filename = filepath.Clean(filename)
	if !filepath.IsAbs(filename) {
		filename = filepath.Join(fs.baseDir, filename)
	}
	if ok, err := fs.insideBaseDirEval(filename); !ok {
		return nil, err
	}
	return os.Lstat(filename)
}

func (fs *UringFS) Readlink(ctx context.Context, link string) (string, error) {
	if !filepath.IsAbs(link) {
		link = filepath.Clean(filepath.Join(fs.baseDir, link))
	}
	if ok, err := fs.insideBaseDirEval(link); !ok {
		return "", err
	}
	return os.Readlink(link)
}

// Chroot returns a new OS filesystem, with the base dir set to the
// result of joining the provided path with the underlying base dir.
// func (fs *UringFS) Chroot(path string) (Filesystem, error) {
// 	joined, err := securejoin.SecureJoin(fs.baseDir, path)
// 	if err != nil {
// 		return nil, err
// 	}
// 	return newBoundOS(joined), nil
// }

// Root returns the current base dir of the billy.Filesystem.
// This is required in order for this implementation to be a drop-in
// replacement for other upstream implementations (e.g. memory and osfs).
func (fs *UringFS) Root() string {
	return fs.baseDir
}

func (fs *UringFS) createDir(fullpath string) error {
	dir := filepath.Dir(fullpath)
	if dir != "." {
		if err := os.MkdirAll(dir, defaultDirectoryMode); err != nil {
			return err
		}
	}

	return nil
}

// abs transforms filename to an absolute path, taking into account the base dir.
// Relative paths won't be allowed to ascend the base dir, so `../file` will become
// `/working-dir/file`.
//
// Note that if filename is a symlink, the returned address will be the target of the
// symlink.
func (fs *UringFS) abs(filename string) (string, error) {
	if filename == fs.baseDir {
		filename = string(filepath.Separator)
	}

	path, err := securejoin.SecureJoin(fs.baseDir, filename)
	if err != nil {
		return "", nil
	}

	return path, nil
}

// insideBaseDirEval checks whether filename is contained within
// a dir that is within the fs.baseDir, by first evaluating any symlinks
// that either filename or fs.baseDir may contain.
func (fs *UringFS) insideBaseDirEval(filename string) (bool, error) {
	dir, err := filepath.EvalSymlinks(filepath.Dir(filename))
	if dir == "" || os.IsNotExist(err) {
		dir = filepath.Dir(filename)
	}
	wd, err := filepath.EvalSymlinks(fs.baseDir)
	if wd == "" || os.IsNotExist(err) {
		wd = fs.baseDir
	}
	if filename != wd && dir != wd && !strings.HasPrefix(dir, wd+string(filepath.Separator)) {
		return false, fmt.Errorf("path outside base dir")
	}
	return true, nil
}

func newFile(fsur *iouring.IOURing, f *os.File) (*URingFile, error) {
	ur, err := iouring.New(64, iouring.WithAttachWQ(fsur))
	if err != nil {
		return nil, err
	}

	return &URingFile{
		ur: ur,
		f:  f,
	}, nil
}

type URingFile struct {
	ur *iouring.IOURing
	f  *os.File
}

// Close implements File.
func (o *URingFile) Close(ctx context.Context) error {
	return errors.Join(o.ur.UnregisterFile(o.f), o.Close(ctx))
}

// Name implements File.
func (o *URingFile) Name() string {
	return o.f.Name()
}

// Read implements File.
func (o *URingFile) Read(ctx context.Context, p []byte) (n int, err error) {
	req, err := o.ur.Read(o.f, p, nil)
	if err != nil {
		return 0, err
	}
	defer req.Cancel()

	select {
	case <-req.Done():
		return req.GetRes()
	case <-ctx.Done():
		req.Cancel()
		<-req.Done()
		return req.GetRes()
	}
}

// ReadAt implements File.
func (o *URingFile) ReadAt(ctx context.Context, p []byte, off int64) (n int, err error) {
	req, err := o.ur.Pread(o.f, p, uint64(off), nil)
	if err != nil {
		return 0, err
	}
	defer req.Cancel()

	select {
	case <-req.Done():
		return req.GetRes()
	case <-ctx.Done():
		req.Cancel()
		<-req.Done()
		return req.GetRes()
	}
}

// Write implements File.
func (o *URingFile) Write(ctx context.Context, p []byte) (n int, err error) {
	req, err := o.ur.Write(o.f, p, nil)
	if err != nil {
		return 0, err
	}
	defer req.Cancel()

	select {
	case <-req.Done():
		return req.GetRes()
	case <-ctx.Done():
		req.Cancel()
		<-req.Done()
		return req.GetRes()
	}
}

// WriteAt implements File.
func (o *URingFile) WriteAt(ctx context.Context, p []byte, off int64) (n int, err error) {
	req, err := o.ur.Pwrite(o.f, p, uint64(off), nil)
	if err != nil {
		return 0, err
	}
	defer req.Cancel()

	select {
	case <-req.Done():
		return req.GetRes()
	case <-ctx.Done():
		req.Cancel()
		<-req.Done()
		return req.GetRes()
	}
}

// Seek implements File.
func (o *URingFile) Seek(offset int64, whence int) (int64, error) {
	return o.f.Seek(offset, whence)
}

// Truncate implements File.
func (o *URingFile) Truncate(ctx context.Context, size int64) error {
	return o.f.Truncate(size)
}

var _ File = (*URingFile)(nil)