Improve pgs filesystem support (and fix local storage)

picosh · Nov 9, 2023 · a8ab7f3 · a8ab7f3
1 parent dc6a7d0
commit a8ab7f3
Show file tree

Hide file tree

Showing 10 changed files with 66 additions and 34 deletions.
diff --git a/filehandlers/assets/handler.go b/filehandlers/assets/handler.go
@@ -102,13 +102,16 @@ func (h *UploadAssetHandler) Read(s ssh.Session, entry *utils.FileEntry) (os.Fil
 	return fileInfo, reader, nil
 }
 
-func (h *UploadAssetHandler) List(s ssh.Session, fpath string) ([]os.FileInfo, error) {
+func (h *UploadAssetHandler) List(s ssh.Session, fpath string, isDir bool) ([]os.FileInfo, error) {
 	var fileList []os.FileInfo
+
 	user, err := getUser(s)
 	if err != nil {
 		return fileList, err
 	}
-	cleanFilename := filepath.Base(fpath)
+
+	cleanFilename := fpath
+
 	bucketName := shared.GetAssetBucketName(user.ID)
 	bucket, err := h.Storage.GetBucket(bucketName)
 	if err != nil {
@@ -125,12 +128,19 @@ func (h *UploadAssetHandler) List(s ssh.Session, fpath string) ([]os.FileInfo, e
 			FName:  name,
 			FIsDir: true,
 		}
+
 		fileList = append(fileList, info)
 	} else {
-		fileList, err = h.Storage.ListFiles(bucket, fpath, false)
+		if cleanFilename != "/" && isDir {
+			cleanFilename += "/"
+		}
+
+		foundList, err := h.Storage.ListFiles(bucket, cleanFilename, false)
 		if err != nil {
 			return fileList, err
 		}
+
+		fileList = append(fileList, foundList...)
 	}
 
 	return fileList, nil

diff --git a/filehandlers/imgs/handler.go b/filehandlers/imgs/handler.go
@@ -111,7 +111,7 @@ func (h *UploadImgHandler) Read(s ssh.Session, entry *utils.FileEntry) (os.FileI
 	return fileInfo, reader, nil
 }
 
-func (h *UploadImgHandler) List(s ssh.Session, fpath string) ([]os.FileInfo, error) {
+func (h *UploadImgHandler) List(s ssh.Session, fpath string, isDir bool) ([]os.FileInfo, error) {
 	var fileList []os.FileInfo
 	user, err := getUser(s)
 	if err != nil {

diff --git a/filehandlers/post_handler.go b/filehandlers/post_handler.go
@@ -89,7 +89,7 @@ func (h *ScpUploadHandler) Read(s ssh.Session, entry *utils.FileEntry) (os.FileI
 	return fileInfo, reader, nil
 }
 
-func (h *ScpUploadHandler) List(s ssh.Session, fpath string) ([]os.FileInfo, error) {
+func (h *ScpUploadHandler) List(s ssh.Session, fpath string, isDir bool) ([]os.FileInfo, error) {
 	var fileList []os.FileInfo
 	user, err := getUser(s)
 	if err != nil {

diff --git a/shared/storage/fs.go b/shared/storage/fs.go
@@ -3,8 +3,8 @@ package storage
 import (
 	"fmt"
 	"io"
-	"io/fs"
 	"os"
+	"path"
 	"path/filepath"
 	"strings"
 
@@ -123,26 +123,46 @@ func (s *StorageFS) DeleteFile(bucket Bucket, fpath string) error {
 
 func (s *StorageFS) ListFiles(bucket Bucket, dir string, recursive bool) ([]os.FileInfo, error) {
 	var fileList []os.FileInfo
-	fpath := filepath.Join(bucket.Path, dir)
-	err := filepath.WalkDir(fpath, func(s string, d fs.DirEntry, err error) error {
+
+	fpath := path.Join(bucket.Path, dir)
+
+	info, err := os.Stat(fpath)
+	if err != nil {
+		return fileList, err
+	}
+
+	if info.IsDir() && !strings.HasSuffix(dir, "/") {
+		fileList = append(fileList, &utils.VirtualFile{
+			FName:    "",
+			FIsDir:   info.IsDir(),
+			FSize:    info.Size(),
+			FModTime: info.ModTime(),
+		})
+
+		return fileList, err
+	}
+
+	files, err := os.ReadDir(fpath)
+	if err != nil {
+		fileList = append(fileList, info)
+		return fileList, nil
+	}
+
+	for _, f := range files {
+		info, err := f.Info()
 		if err != nil {
-			return err
+			return fileList, err
 		}
-		if !d.IsDir() {
-			fileInfo, err := os.Stat(s)
-			if err != nil {
-				return err
-			}
-			info := &utils.VirtualFile{
-				FName:    strings.Replace(s, bucket.Path, "", 1),
-				FIsDir:   d.IsDir(),
-				FSize:    fileInfo.Size(),
-				FModTime: fileInfo.ModTime(),
-			}
-			fileList = append(fileList, info)
+
+		i := &utils.VirtualFile{
+			FName:    f.Name(),
+			FIsDir:   f.IsDir(),
+			FSize:    info.Size(),
+			FModTime: info.ModTime(),
 		}
-		return nil
-	})
+
+		fileList = append(fileList, i)
+	}
 
 	return fileList, err
 }
diff --git a/shared/storage/minio.go b/shared/storage/minio.go
@@ -99,6 +99,7 @@ func (s *StorageMinio) ListFiles(bucket Bucket, dir string, recursive bool) ([]o
 	var fileList []os.FileInfo
 
 	resolved := strings.TrimPrefix(dir, "/")
+
 	opts := minio.ListObjectsOptions{Prefix: resolved, Recursive: recursive}
 	for obj := range s.Client.ListObjects(context.Background(), bucket.Name, opts) {
 		if obj.Err != nil {
@@ -108,8 +109,9 @@ func (s *StorageMinio) ListFiles(bucket Bucket, dir string, recursive bool) ([]o
 		if obj.Size == 0 {
 			isDir = true
 		}
+
 		info := &utils.VirtualFile{
-			FName:    strings.TrimSuffix(obj.Key, "/"),
+			FName:    strings.TrimSuffix(strings.TrimPrefix(obj.Key, resolved), "/"),
 			FIsDir:   isDir,
 			FSize:    obj.Size,
 			FModTime: obj.LastModified,

diff --git a/wish/cmd/server/main.go b/wish/cmd/server/main.go
@@ -42,7 +42,7 @@ func (h *handler) Read(session ssh.Session, entry *utils.FileEntry) (os.FileInfo
 	}, data, nil
 }
 
-func (h *handler) List(session ssh.Session, fpath string) ([]os.FileInfo, error) {
+func (h *handler) List(session ssh.Session, fpath string, isDir bool) ([]os.FileInfo, error) {
 	return nil, nil
 }
 

diff --git a/wish/list/list.go b/wish/list/list.go
@@ -18,7 +18,7 @@ func Middleware(writeHandler utils.CopyFromClientHandler) wish.Middleware {
 				return
 			}
 
-			fileList, err := writeHandler.List(session, "/")
+			fileList, err := writeHandler.List(session, "/", true)
 			if err != nil {
 				utils.ErrorHandler(session, err)
 				return

diff --git a/wish/send/rsync/rsync.go b/wish/send/rsync/rsync.go
@@ -26,7 +26,7 @@ func (h *handler) Skip(file *rsyncutils.ReceiverFile) bool {
 }
 
 func (h *handler) List(path string) ([]fs.FileInfo, error) {
-	list, err := h.writeHandler.List(h.session, path)
+	list, err := h.writeHandler.List(h.session, path, true)
 	if err != nil {
 		return nil, err
 	}

diff --git a/wish/send/sftp/handler.go b/wish/send/sftp/handler.go
@@ -40,23 +40,23 @@ func (f *handler) Filecmd(r *sftp.Request) error {
 		_, err := f.writeHandler.Write(f.session, entry)
 
 		return err
-	case "Setstat":
+	case "Setstat", "Mkdir":
 		return nil
 	}
 	return errors.New("unsupported")
 }
 
 func (f *handler) Filelist(r *sftp.Request) (sftp.ListerAt, error) {
 	switch r.Method {
-	case "List":
-		fallthrough
-	case "Stat":
-		listData, err := f.writeHandler.List(f.session, r.Filepath)
+	case "List", "Stat":
+		list := r.Method == "List"
+
+		listData, err := f.writeHandler.List(f.session, r.Filepath, list)
 		if err != nil {
 			return nil, err
 		}
 
-		if r.Filepath != "/" {
+		if list {
 			listData = slices.DeleteFunc(listData, func(f os.FileInfo) bool {
 				return f.Name() == "/"
 			})

diff --git a/wish/send/utils/utils.go b/wish/send/utils/utils.go
@@ -58,7 +58,7 @@ type CopyFromClientHandler interface {
 	// Write should write the given file.
 	Write(ssh.Session, *FileEntry) (string, error)
 	Read(ssh.Session, *FileEntry) (os.FileInfo, io.ReaderAt, error)
-	List(ssh.Session, string) ([]os.FileInfo, error)
+	List(ssh.Session, string, bool) ([]os.FileInfo, error)
 	Validate(ssh.Session) error
 }