Add support for SFTP remote file copies

1 files changed, 151 insertions, 86 deletions

diff --git a/internal/cp/cp.go b/internal/cp/cp.go
index c19627f..d944f9c 100644
--- a/internal/cp/cp.go
+++ b/internal/cp/cp.go
@@ -5,21 +5,36 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
-	"os"
-	"path/filepath"
+	"path"
+	"sync"
+
+	"gitlab.com/rhogenson/ccp/internal/wfs"
+	"gitlab.com/rhogenson/ccp/internal/wfs/sftpfs"
 )
 
 type Progress interface {
 	Max(int64)
 	Progress(int64)
-	FileStart(string)
+	FileStart(string, string)
 	FileDone(string, error)
 }
 
-func size(files []string) int64 {
+type FSPath struct {
+	FS   wfs.FS
+	Path string
+}
+
+func (p FSPath) String() string {
+	if fsys, ok := p.FS.(*sftpfs.FS); ok {
+		return fsys.User + "@" + fsys.Host + ":" + p.Path
+	}
+	return p.Path
+}
+
+func size(srcs []FSPath) int64 {
 	var n int64 = 0
-	for _, f := range files {
-		filepath.WalkDir(f, func(path string, d fs.DirEntry, err error) error {
+	for _, src := range srcs {
+		fs.WalkDir(src.FS, src.Path, func(_ string, d fs.DirEntry, err error) error {
 			if err != nil {
 				return nil
 			}
@@ -30,7 +45,7 @@ func size(files []string) int64 {
 					return nil
 				}
 				n += 1 + stat.Size()
-			case os.ModeSymlink, os.ModeDir:
+			case fs.ModeSymlink, fs.ModeDir:
 				n++
 			}
 			return nil
@@ -39,15 +54,40 @@ func size(files []string) int64 {
 	return n
 }
 
-func fileExists(path string) bool {
-	_, err := os.Lstat(path)
-	return !errors.Is(err, os.ErrNotExist)
+func fileExists(path FSPath) bool {
+	_, err := wfs.Lstat(path.FS, path.Path)
+	return !errors.Is(err, fs.ErrNotExist)
+}
+
+type copier struct {
+	p   Progress
+	sem chan struct{}
+
+	force bool
+}
+
+func (c *copier) g(fn func()) {
+	c.sem <- struct{}{}
+	go func() {
+		defer func() { <-c.sem }()
+		fn()
+	}()
+}
+
+func (c *copier) openWithRetry(path FSPath, fn func() error) error {
+	if err := fn(); err == nil || !c.force || !fileExists(path) {
+		return err
+	}
+	if err := wfs.RemoveAll(path.FS, path.Path); err != nil {
+		return err
+	}
+	return fn()
 }
 
-func copyRegularFile(src, dst string, force bool, progress Progress) (err error) {
-	progress.FileStart(src)
+func (c *copier) copyRegularFile(src, dst FSPath) error {
+	c.p.FileStart(src.String(), dst.String())
 
-	in, err := os.Open(src)
+	in, err := src.FS.Open(src.Path)
 	if err != nil {
 		return err
 	}
@@ -56,117 +96,142 @@ func copyRegularFile(src, dst string, force bool, progress Progress) (err error)
 	if err != nil {
 		return err
 	}
-	out, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, stat.Mode().Perm())
-	if err != nil {
-		if !force || !fileExists(dst) {
-			return err
-		}
-		if err := os.RemoveAll(dst); err != nil {
-			return err
-		}
-		out, err = os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, stat.Mode().Perm())
-		if err != nil {
-			return err
-		}
+	var out io.WriteCloser
+	if err := c.openWithRetry(dst, func() error {
+		var err error
+		out, err = dst.FS.Create(dst.Path, stat.Mode().Perm())
+		return err
+	}); err != nil {
+		return err
 	}
-	defer out.Close()
 	for {
 		n, err := io.CopyN(out, in, 1024*1024)
 		if n > 0 {
-			progress.Progress(n)
+			c.p.Progress(n)
 		}
 		if err != nil {
 			if err == io.EOF {
 				break
 			}
+			out.Close()
 			return err
 		}
 	}
 	if err := out.Close(); err != nil {
 		return err
 	}
-	progress.Progress(1)
+	c.p.Progress(1)
+	c.p.FileDone(src.String(), nil)
 	return nil
 }
 
-func copySpecialFile(src string, d fs.DirEntry, dst string, force bool, progress Progress) error {
-	switch d.Type() {
-	case fs.ModeSymlink:
-		target, err := os.Readlink(src)
-		if err != nil {
-			return err
+func (c *copier) copySymlink(src FSPath, d fs.DirEntry, dst FSPath) error {
+	target, err := wfs.ReadLink(src.FS, src.Path)
+	if err != nil {
+		return err
+	}
+	if err := c.openWithRetry(dst, func() error {
+		return dst.FS.Symlink(target, dst.Path)
+	}); err != nil {
+		return err
+	}
+	c.p.Progress(1)
+	return nil
+}
+
+func (c *copier) copyDir(src FSPath, d fs.DirEntry, dst FSPath) error {
+	stat, err := d.Info()
+	if err != nil {
+		return err
+	}
+	hasWritePerm := stat.Mode()&0300 == 0300
+	if err := c.openWithRetry(dst, func() error {
+		if hasWritePerm {
+			return wfs.MkdirMode(dst.FS, dst.Path, stat.Mode().Perm())
+		} else {
+			return dst.FS.Mkdir(dst.Path)
 		}
-		if err := os.Symlink(target, dst); err != nil {
-			if !force || !fileExists(dst) {
-				return err
-			}
-			if err := os.RemoveAll(dst); err != nil {
-				return err
-			}
-			if err := os.Symlink(target, dst); err != nil {
-				return err
+	}); err != nil {
+		return err
+	}
+	entries, err := fs.ReadDir(src.FS, src.Path)
+	wg := new(sync.WaitGroup)
+	<-c.sem
+	for _, entry := range entries {
+		wg.Add(1)
+		c.g(func() {
+			defer wg.Done()
+			src := FSPath{src.FS, path.Join(src.Path, entry.Name())}
+			if err := c.copyFile(src, entry, FSPath{dst.FS, path.Join(dst.Path, entry.Name())}); err != nil {
+				c.p.FileDone(src.String(), err)
 			}
-		}
-		progress.Progress(1)
-	case fs.ModeDir:
-		stat, err := d.Info()
-		if err != nil {
+		})
+	}
+	wg.Wait()
+	c.sem <- struct{}{}
+	if err != nil {
+		return err
+	}
+	if !hasWritePerm {
+		if err := dst.FS.Chmod(dst.Path, stat.Mode().Perm()); err != nil {
 			return err
 		}
-		if err := os.Mkdir(dst, stat.Mode().Perm()); err != nil {
-			if !force || !fileExists(dst) {
-				return err
-			}
-			if err := os.RemoveAll(dst); err != nil {
-				return err
-			}
-			if err := os.Mkdir(dst, stat.Mode().Perm()); err != nil {
-				return err
-			}
-		}
-		progress.Progress(1)
+	}
+	c.p.Progress(1)
+	return nil
+}
+
+func (c *copier) copyFile(src FSPath, d fs.DirEntry, dst FSPath) error {
+	switch d.Type() {
+	case 0: // regular file
+		return c.copyRegularFile(src, dst)
+	case fs.ModeDir:
+		return c.copyDir(src, d, dst)
+	case fs.ModeSymlink:
+		return c.copySymlink(src, d, dst)
 	default:
 		return fmt.Errorf("%s: unknown file type %s", src, d.Type())
 	}
-	return nil
 }
 
-func Copy(srcs []string, dstRoot string, force bool, progress Progress) {
+func Copy(progress Progress, srcs []FSPath, dstRoot FSPath, force bool) {
 	go func() {
 		progress.Max(size(srcs))
 	}()
+
+	dstIsDir := true
+	if len(srcs) == 1 {
+		stat, err := fs.Stat(dstRoot.FS, dstRoot.Path)
+		dstIsDir = err == nil && stat.IsDir()
+	}
+
 	const maxConcurrency = 10
 	// sem acts as a semaphore to limit the number of concurrent file copies
 	sem := make(chan struct{}, maxConcurrency)
+	c := &copier{
+		p:     progress,
+		sem:   sem,
+		force: force,
+	}
+	wg := new(sync.WaitGroup)
 	for _, srcRoot := range srcs {
-		filepath.WalkDir(srcRoot, func(src string, d fs.DirEntry, err error) error {
-			if err != nil {
-				progress.FileDone(src, err)
-				return nil
+		wg.Add(1)
+		c.g(func() {
+			defer wg.Done()
+			dstRoot := dstRoot
+			if dstIsDir {
+				dstRoot.Path = path.Join(dstRoot.Path, path.Base(srcRoot.Path))
 			}
-			relPath, err := filepath.Rel(srcRoot, src)
+			stat, err := fs.Stat(srcRoot.FS, srcRoot.Path)
 			if err != nil {
-				progress.FileDone(src, err)
-				return nil
+				progress.FileDone(srcRoot.String(), err)
+				return
 			}
-			dst := filepath.Join(dstRoot, relPath)
-			if d.Type().IsRegular() {
-				sem <- struct{}{}
-				go func() {
-					defer func() { <-sem }()
-					err := copyRegularFile(src, dst, force, progress)
-					progress.FileDone(src, err)
-				}()
-			} else {
-				if err := copySpecialFile(src, d, dst, force, progress); err != nil {
-					progress.FileDone(src, err)
-					return nil
-				}
+			if err := c.copyFile(srcRoot, fs.FileInfoToDirEntry(stat), dstRoot); err != nil {
+				progress.FileDone(srcRoot.String(), err)
+				return
 			}
-			return nil
 		})
 	}
-	for range maxConcurrency {
-		sem <- struct{}{}
-	}
+	wg.Wait()
 }