diff --git a/ssh/scp.go b/ssh/scp.go index a83cd9fe..3ada1f19 100644 --- a/ssh/scp.go +++ b/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 -P -J 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 } diff --git a/ssh/scp_test.go b/ssh/scp_test.go index 4f396037..b902e905 100644 --- a/ssh/scp_test.go +++ b/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, diff --git a/ssh/ssh.go b/ssh/ssh.go index 641f09b0..54a084f6 100644 --- a/ssh/ssh.go +++ b/ssh/ssh.go @@ -26,6 +26,7 @@ type SSHArgs struct { PrivateKeyPath string Port string MaxRetries int + Compress bool ProxyJump *ProxyJumpArgs } diff --git a/starlark/copy_from.go b/starlark/copy_from.go index cdbcf0fc..9a46a7fe 100644 --- a/starlark/copy_from.go +++ b/starlark/copy_from.go @@ -22,9 +22,12 @@ 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=, resources=resources, workdir=path]) +// If compressed flag is true, the command will compress the file sent to the server. +// +// Starlark format: copy_from([] [,path=, 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( @@ -32,6 +35,7 @@ func copyFromFunc(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tu "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) { diff --git a/starlark/copy_from_test.go b/starlark/copy_from_test.go index c66645e3..20667126 100644 --- a/starlark/copy_from_test.go +++ b/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 {