add interruptibleReader,can interrupt file upload or download

Author: leijux

internal/pkg/engin/engin.go

@@ -151,13 +151,13 @@ func (e *Engin) Close() error {
 	return errs
 }
 
-func execInternalCommand(client *ssh.Client, command string) (string, error) {
+func execInternalCommand(ctx context.Context, client *ssh.Client, command string) (string, error) {
 	if client == nil {
 		command = strings.Join([]string{"rscript.exec", command}, " ")
 	}
 	commands := strings.Split(strings.TrimSpace(command), " ")
 	if iFunc, ok := internalFuncMap[commands[0]]; ok {
-		output, err := iFunc(client, commands[1:]...)
+		output, err := iFunc(ctx, client, commands[1:]...)
 		if err != nil {
 			return "", err
 		}
@@ -180,7 +180,7 @@ func execExternalCommand(client *ssh.Client, command string) (string, error) {
 
 func (e *Engin) execCommand(client *ssh.Client, command string) (string, error) {
 	if strings.HasPrefix(command, prefix) {
-		return execInternalCommand(client, command)
+		return execInternalCommand(e.ctx, client, command)
 	}
 	return execExternalCommand(client, command)
 }

internal/pkg/engin/funcs.go

@@ -1,6 +1,7 @@
 package engin
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -15,7 +16,7 @@ import (
 
 const prefix = "rscript"
 
-type internalFunc func(*ssh.Client, ...string) (string, error)
+type internalFunc func(context.Context, *ssh.Client, ...string) (string, error)
 
 var internalFuncMap = map[string]internalFunc{
 	prefix + ".uploadFile":   UploadFileWithFs(nil),
@@ -28,7 +29,7 @@ func SetInternalFunc(name string, f internalFunc) {
 }
 
 // remoteFilePath, localPath string
-func downloadFile(client *ssh.Client, args ...string) (string, error) {
+func downloadFile(ctx context.Context, client *ssh.Client, args ...string) (string, error) {
 	if len(args) != 2 {
 		return "", errors.New("args error")
 	}
@@ -61,7 +62,7 @@ func downloadFile(client *ssh.Client, args ...string) (string, error) {
 	}
 	defer localFile.Close()
 
-	_, err = io.Copy(localFile, remoteFile)
+	_, err = io.Copy(localFile, NewInterruptibleReader(ctx, remoteFile))
 	if err != nil {
 		return "", fmt.Errorf("io.Copy: %w", err)
 	}
@@ -70,7 +71,7 @@ func downloadFile(client *ssh.Client, args ...string) (string, error) {
 }
 
 func UploadFileWithFs(f fs.FS) internalFunc {
-	return func(client *ssh.Client, args ...string) (string, error) {
+	return func(ctx context.Context, client *ssh.Client, args ...string) (string, error) {
 		if len(args) != 2 {
 			return "", errors.New("args error")
 		}
@@ -113,7 +114,7 @@ func UploadFileWithFs(f fs.FS) internalFunc {
 		}
 		defer remoteFile.Close()
 
-		_, err = io.Copy(remoteFile, localFile)
+		_, err = io.Copy(remoteFile, NewInterruptibleReader(ctx, localFile))
 		if err != nil {
 			return "", fmt.Errorf("io.Copy: %w", err)
 		}
@@ -122,7 +123,7 @@ func UploadFileWithFs(f fs.FS) internalFunc {
 	}
 }
 
-func execLocalCommand(_ *ssh.Client, args ...string) (string, error) {
+func execLocalCommand(ctx context.Context, _ *ssh.Client, args ...string) (string, error) {
 	var (
 		name    string
 		cmdArgs []string
@@ -147,3 +148,20 @@ func execLocalCommand(_ *ssh.Client, args ...string) (string, error) {
 	}
 	return out.String(), nil
 }
+
+type InterruptibleReader func(p []byte) (n int, err error)
+
+func (r InterruptibleReader) Read(p []byte) (n int, err error) {
+	return r(p)
+}
+
+func NewInterruptibleReader(ctx context.Context, r io.Reader) io.Reader {
+	return InterruptibleReader(func(p []byte) (n int, err error) {
+		select {
+		case <-ctx.Done():
+			return 0, io.EOF
+		default:
+			return r.Read(p)
+		}
+	})
+}