Adding support for copy-compress

Signed-off-by: Vladimir Vivien <[email protected]>

ssh/scp.go

@@ -151,6 +151,13 @@ func makeSCPCmdStr(progName string, args SSHArgs) (string, error) {
 		return fmt.Sprintf("-P %s", args.Port)
 	}
 
+	compress := func() string {
+		if args.Compress {
+			return "-C"
+		}
+		return ""
+	}
+
 	proxyJump := func() string {
 		if args.ProxyJump != nil {
 			return fmt.Sprintf("-J %s@%s", args.ProxyJump.User, args.ProxyJump.Host)
@@ -160,8 +167,8 @@ func makeSCPCmdStr(progName string, args SSHArgs) (string, error) {
 	// build command as
 	// scp -i <pkpath> -P <port> -J <proxyjump> user@host:path
 	cmd := fmt.Sprintf(
-		`%s %s %s %s`,
-		scpCmdPrefix(), pkPath(), port(), proxyJump(),
+		`%s %s %s %s %s`,
+		scpCmdPrefix(), pkPath(), port(), compress(), proxyJump(),
 	)
 	return cmd, nil
 }

ssh/scp_test.go

@@ -26,6 +26,20 @@ func TestCopyFrom(t *testing.T) {
 			srcFile:     "foo.txt",
 			fileContent: "FooBar",
 		},
+		{
+			name: "copy single compress",
+			sshArgs: SSHArgs{
+				User:           support.CurrentUsername(),
+				PrivateKeyPath: support.PrivateKeyPath(),
+				Host:           "127.0.0.1",
+				Compress:       true,
+				Port:           support.PortValue(),
+				MaxRetries:     support.MaxConnectionRetries(),
+			},
+			remoteFiles: map[string]string{"foo.txt": "FooBar"},
+			srcFile:     "foo.txt",
+			fileContent: "FooBar",
+		},
 		{
 			name:        "copy single file in dir",
 			sshArgs:     testSSHArgs,

ssh/ssh.go

@@ -26,6 +26,7 @@ type SSHArgs struct {
 	PrivateKeyPath string
 	Port           string
 	MaxRetries     int
+	Compress       bool
 	ProxyJump      *ProxyJumpArgs
 }
 

starlark/copy_from.go

@@ -22,16 +22,20 @@ import (
 // If resources and workdir are not provided, copyFromFunc uses defaults from starlark thread generated
 // by previous calls to resources(), ssh_config, and crashd_config().
 //
-// Starlark format: copy_from([<path>] [,path=<list>, resources=resources, workdir=path])
+// If compressed flag is true, the command will compress the file sent to the server.
+//
+// Starlark format: copy_from([<path>] [,path=<list>, resources=resources, workdir=path, compress=False])
 func copyFromFunc(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var sourcePath, workdir string
+	var compress bool
 	var resources *starlark.List
 
 	if err := starlark.UnpackArgs(
 		identifiers.capture, args, kwargs,
 		"path", &sourcePath,
 		"resources?", &resources,
 		"workdir?", &workdir,
+		"compress?", &compress,
 	); err != nil {
 		return starlark.None, fmt.Errorf("%s: %s", identifiers.capture, err)
 	}
@@ -66,7 +70,7 @@ func copyFromFunc(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tu
 		}
 	}
 
-	results, err := execCopyFrom(workdir, sourcePath, agent, resources)
+	results, err := execCopyFrom(workdir, sourcePath, compress, agent, resources)
 	if err != nil {
 		return starlark.None, fmt.Errorf("%s: %s", identifiers.copyFrom, err)
 	}
@@ -83,7 +87,7 @@ func copyFromFunc(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tu
 	return starlark.NewList(resultList), nil
 }
 
-func execCopyFrom(rootPath string, path string, agent ssh.Agent, resources *starlark.List) ([]commandResult, error) {
+func execCopyFrom(rootPath string, path string, compress bool, agent ssh.Agent, resources *starlark.List) ([]commandResult, error) {
 	if resources == nil {
 		return nil, fmt.Errorf("%s: missing resources", identifiers.copyFrom)
 	}
@@ -117,7 +121,7 @@ func execCopyFrom(rootPath string, path string, agent ssh.Agent, resources *star
 
 		switch {
 		case string(kind) == identifiers.hostResource && string(transport) == "ssh":
-			result, err := execSCPCopyFrom(host, rootDir, path, agent, res)
+			result, err := execSCPCopyFrom(host, rootDir, path, compress, agent, res)
 			if err != nil {
 				logrus.Errorf("%s: failed to copyFrom %s: %s", identifiers.copyFrom, path, err)
 			}
@@ -131,7 +135,7 @@ func execCopyFrom(rootPath string, path string, agent ssh.Agent, resources *star
 	return results, nil
 }
 
-func execSCPCopyFrom(host, rootDir, path string, agent ssh.Agent, res *starlarkstruct.Struct) (commandResult, error) {
+func execSCPCopyFrom(host, rootDir, path string, compress bool, agent ssh.Agent, res *starlarkstruct.Struct) (commandResult, error) {
 	sshCfg := starlarkstruct.FromKeywords(starlarkstruct.Default, makeDefaultSSHConfig())
 	if val, err := res.Attr(identifiers.sshCfg); err == nil {
 		if cfg, ok := val.(*starlarkstruct.Struct); ok {
@@ -144,6 +148,7 @@ func execSCPCopyFrom(host, rootDir, path string, agent ssh.Agent, res *starlarks
 		return commandResult{}, err
 	}
 	args.Host = host
+	args.Compress = compress
 
 	// create dir for the host
 	if err := os.MkdirAll(rootDir, 0744); err != nil && !os.IsExist(err) {

starlark/copy_from_test.go

@@ -74,7 +74,58 @@ func testCopyFromFuncForHostResources(t *testing.T, port, privateKey, username s
 				defer os.RemoveAll(expected)
 			},
 		},
+		{
+			name:        "single machine compress",
+			remoteFiles: map[string]string{"foo.txt": "FooBar"},
+			args:        func(t *testing.T) starlark.Tuple { return starlark.Tuple{starlark.String("foo.txt")} },
+			kwargs: func(t *testing.T) []starlark.Tuple {
+				sshCfg := makeTestSSHConfig(privateKey, port, username)
+				resources := starlark.NewList([]starlark.Value{makeTestSSHHostResource("127.0.0.1", sshCfg)})
+				return []starlark.Tuple{
+					[]starlark.Value{starlark.String("resources"), resources},
+					[]starlark.Value{starlark.String("compress"), starlark.Bool(true)},
+				}
+			},
 
+			eval: func(t *testing.T, args starlark.Tuple, kwargs []starlark.Tuple) {
+
+				val, err := copyFromFunc(newTestThreadLocal(t), nil, args, kwargs)
+				if err != nil {
+					t.Fatal(err)
+				}
+				resource := ""
+				cpErr := ""
+				result := ""
+				if strct, ok := val.(*starlarkstruct.Struct); ok {
+					if val, err := strct.Attr("resource"); err == nil {
+						if r, ok := val.(starlark.String); ok {
+							resource = string(r)
+						}
+					}
+					if val, err := strct.Attr("err"); err == nil {
+						if r, ok := val.(starlark.String); ok {
+							cpErr = string(r)
+						}
+					}
+					if val, err := strct.Attr("result"); err == nil {
+						if r, ok := val.(starlark.String); ok {
+							result = string(r)
+						}
+					}
+				}
+
+				if cpErr != "" {
+					t.Fatal(cpErr)
+				}
+
+				expected := filepath.Join(defaults.workdir, sanitizeStr(resource), "foo.txt")
+				if result != expected {
+					t.Errorf("unexpected file name copied: %s", result)
+				}
+
+				defer os.RemoveAll(expected)
+			},
+		},
 		{
 			name:        "multiple machines single files",
 			remoteFiles: map[string]string{"bar/bar.txt": "BarBar", "bar/foo.txt": "FooBar", "baz.txt": "BazBuz"},
@@ -203,7 +254,13 @@ func testCopyFromFuncForHostResources(t *testing.T, port, privateKey, username s
 		},
 	}
 
-	sshArgs := ssh.SSHArgs{User: username, Host: "127.0.0.1", Port: port, PrivateKeyPath: privateKey}
+	sshArgs := ssh.SSHArgs{
+		User:           username,
+		Host:           "127.0.0.1",
+		Port:           port,
+		PrivateKeyPath: privateKey,
+		MaxRetries:     testSupport.MaxConnectionRetries(),
+	}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			for file, content := range test.remoteFiles {