perf: several dragged items cross to another place as one tar stream; older agents get them one by one

internal/agent/ship.go

@@ -99,7 +99,24 @@ func treeSize(ctx context.Context, p string) int64 {
 // packTar writes src as a tar stream whose entries start at name. Owners are
 // left out: the numbers mean other people on another machine.
 func packTar(ctx context.Context, w io.Writer, src, name string, m *shipMeter) error {
+	return packTars(ctx, w, []shipItem{{src, name}}, m)
+}
+
+// shipItem is one thing sent: where it is and the name it goes under.
+type shipItem struct{ src, name string }
+
+// packTars writes the items as one tar stream, each under its name.
+func packTars(ctx context.Context, w io.Writer, items []shipItem, m *shipMeter) error {
 	tw := tar.NewWriter(w)
+	for _, it := range items {
+		if err := packInto(ctx, tw, it.src, it.name, m); err != nil {
+			return err
+		}
+	}
+	return tw.Close()
+}
+
+func packInto(ctx context.Context, tw *tar.Writer, src, name string, m *shipMeter) error {
 	err := filepath.WalkDir(src, func(p string, d fs.DirEntry, err error) error {
 		if err != nil {
 			return err
@@ -152,10 +169,7 @@ func packTar(ctx context.Context, w io.Writer, src, name string, m *shipMeter) e
 		m.files.Add(1)
 		return nil
 	})
-	if err != nil {
-		return err
-	}
-	return tw.Close()
+	return err
 }
 
 type countReader struct {
@@ -175,10 +189,24 @@ func (a *Agent) doShipOut(parent context.Context, env proto.TaskEnvelope) proto.
 	if err := json.Unmarshal(env.Args, &args); err != nil {
 		return proto.TaskResult{Error: err.Error()}
 	}
-	if err := shipName(args.Name); err != nil {
-		return proto.TaskResult{Error: err.Error()}
+	names := args.ShipNames()
+	for _, n := range names {
+		if err := shipName(n); err != nil {
+			return proto.TaskResult{Error: err.Error()}
+		}
+	}
+	if len(args.More) > 0 && args.SSH != nil {
+		return proto.TaskResult{Error: "по SSH-адресу - по одному"}
 	}
 	src := filepath.Clean(args.Src)
+	srcs := []string{args.Src}
+	for _, p := range args.More {
+		// The others: of the same folder only.
+		if path.Dir(path.Clean(p)) != path.Dir(path.Clean(args.Src)) {
+			return proto.TaskResult{Error: "несколько - только из одной папки"}
+		}
+		srcs = append(srcs, p)
+	}
 	var conn *sshfsConn
 	if args.Conn != "" {
 		// The files are on an SSH host this agent holds a connection to.
@@ -195,14 +223,20 @@ func (a *Agent) doShipOut(parent context.Context, env proto.TaskEnvelope) proto.
 	m := &shipMeter{}
 	m.note.Store("подсчёт")
 	go a.reportShip(ctx, args.Job, m, cancel)
-	if conn != nil {
-		m.total.Store(a.remoteSize(conn, src))
-	} else {
-		m.total.Store(treeSize(ctx, src))
+	var total int64
+	items := make([]shipItem, len(srcs))
+	for i, p := range srcs {
+		if conn != nil {
+			total += a.remoteSize(conn, path.Clean(p))
+		} else {
+			total += treeSize(ctx, filepath.Clean(p))
+		}
+		items[i] = shipItem{filepath.Clean(p), names[i]}
 	}
-	pack := func(w io.Writer) error { return packTar(ctx, w, src, args.Name, m) }
+	m.total.Store(total)
+	pack := func(w io.Writer) error { return packTars(ctx, w, items, m) }
 	if conn != nil {
-		pack = func(w io.Writer) error { return a.packRemote(ctx, w, conn, src, args.Name, m) }
+		pack = func(w io.Writer) error { return a.packRemote(ctx, w, conn, srcs, names, m) }
 	}
 
 	var err error
@@ -318,8 +352,11 @@ func (a *Agent) doShipIn(ctx context.Context, env proto.TaskEnvelope) proto.Task
 	if err := json.Unmarshal(env.Args, &args); err != nil {
 		return proto.TaskResult{Error: err.Error()}
 	}
-	if err := shipName(args.Name); err != nil {
-		return proto.TaskResult{Error: err.Error()}
+	names := args.ShipNames()
+	for _, n := range names {
+		if err := shipName(n); err != nil {
+			return proto.TaskResult{Error: err.Error()}
+		}
 	}
 	dir := filepath.Clean(args.Dir)
 	var conn *sshfsConn
@@ -331,8 +368,10 @@ func (a *Agent) doShipIn(ctx context.Context, env proto.TaskEnvelope) proto.Task
 		}
 		conn, dir = c, path.Clean(args.Dir)
 	} else {
-		if _, err := os.Lstat(filepath.Join(dir, args.Name)); err == nil && args.Conflict != "replace" && args.Conflict != "skip" {
-			return proto.TaskResult{Error: filepath.Join(dir, args.Name) + " уже существует"}
+		for _, n := range names {
+			if _, err := os.Lstat(filepath.Join(dir, n)); err == nil && args.Conflict != "replace" && args.Conflict != "skip" {
+				return proto.TaskResult{Error: filepath.Join(dir, n) + " уже существует"}
+			}
 		}
 		if err := os.MkdirAll(dir, 0o755); err != nil {
 			return proto.TaskResult{Error: err.Error()}
@@ -366,7 +405,7 @@ func (a *Agent) doShipIn(ctx context.Context, env proto.TaskEnvelope) proto.Task
 		}
 	}()
 	if conn != nil {
-		skipped, err := a.unpackRemote(conn, pr, dir, args.Name, args.Conflict)
+		skipped, err := a.unpackRemote(conn, pr, dir, names, args.Conflict)
 		if err == nil {
 			_, _ = io.Copy(io.Discard, pr)
 		}
@@ -375,13 +414,17 @@ func (a *Agent) doShipIn(ctx context.Context, env proto.TaskEnvelope) proto.Task
 			return proto.TaskResult{Error: err.Error()}
 		}
 		res := map[string]int{}
-		if skipped {
-			res["skipped"] = 1
+		if skipped > 0 {
+			res["skipped"] = skipped
 		}
 		data, _ := json.Marshal(res)
 		return proto.TaskResult{OK: true, Data: data}
 	}
-	u := &unpacker{dir: dir, name: args.Name, conflict: args.Conflict}
+	allowed := map[string]bool{}
+	for _, n := range names {
+		allowed[n] = true
+	}
+	u := &unpacker{dir: dir, names: allowed, conflict: args.Conflict}
 	if os.Geteuid() == 0 {
 		// New things belong to whoever owns the folder they land in.
 		if info, err := os.Stat(dir); err == nil {
@@ -412,7 +455,8 @@ func (a *Agent) doShipIn(ctx context.Context, env proto.TaskEnvelope) proto.Task
 // the conflict rule for what is there already: folders merge, and files are
 // replaced or kept.
 type unpacker struct {
-	dir, name, conflict string
+	dir, conflict string
+	names         map[string]bool // the top entries the stream may carry
 	uid, gid            int
 	chown               bool
 	done, skipped       int
@@ -429,7 +473,8 @@ func (u *unpacker) unpack(tr *tar.Reader) error {
 			return err
 		}
 		clean := path.Clean(hdr.Name)
-		if clean != u.name && !strings.HasPrefix(clean, u.name+"/") {
+		top, _, _ := strings.Cut(clean, "/")
+		if !u.names[top] || clean == ".." || strings.HasPrefix(clean, "../") {
 			return fmt.Errorf("в архиве чужой путь %q", hdr.Name)
 		}
 		target := filepath.Join(u.dir, filepath.FromSlash(clean))

internal/agent/shipmulti_test.go

@@ -0,0 +1,66 @@
+package agent
+
+import (
+	"archive/tar"
+	"bytes"
+	"context"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+)
+
+// Several items of one folder go as one tar and land side by side; an entry
+// of none of them is refused.
+func TestShipSeveralInOneTar(t *testing.T) {
+	src, dst := t.TempDir(), t.TempDir()
+	_ = os.WriteFile(filepath.Join(src, "a.txt"), []byte("A"), 0o644)
+	_ = os.MkdirAll(filepath.Join(src, "b"), 0o755)
+	_ = os.WriteFile(filepath.Join(src, "b", "c.txt"), []byte("C"), 0o644)
+	var buf bytes.Buffer
+	items := []shipItem{{filepath.Join(src, "a.txt"), "a.txt"}, {filepath.Join(src, "b"), "b"}}
+	if err := packTars(context.Background(), &buf, items, &shipMeter{}); err != nil {
+		t.Fatal(err)
+	}
+	u := &unpacker{dir: dst, names: map[string]bool{"a.txt": true, "b": true}}
+	if err := u.unpack(tar.NewReader(&buf)); err != nil {
+		t.Fatal(err)
+	}
+	if b, _ := os.ReadFile(filepath.Join(dst, "b", "c.txt")); string(b) != "C" {
+		t.Fatalf("b/c.txt: %q", b)
+	}
+	if b, _ := os.ReadFile(filepath.Join(dst, "a.txt")); string(b) != "A" {
+		t.Fatalf("a.txt: %q", b)
+	}
+	buf.Reset()
+	_ = packTars(context.Background(), &buf, items, &shipMeter{})
+	u = &unpacker{dir: t.TempDir(), names: map[string]bool{"a.txt": true}}
+	if err := u.unpack(tar.NewReader(&buf)); err == nil || !strings.Contains(err.Error(), "чужой путь") {
+		t.Fatalf("b let in: %v", err)
+	}
+}
+
+func TestRenameTarSeveral(t *testing.T) {
+	var src bytes.Buffer
+	tw := tar.NewWriter(&src)
+	for _, n := range []string{"a", "b/x", "c", "._a"} {
+		_ = tw.WriteHeader(&tar.Header{Name: n, Mode: 0o644, Typeflag: tar.TypeReg})
+	}
+	_ = tw.Close()
+	var out bytes.Buffer
+	if err := renameTar(&out, &src, map[string]string{"a": "a", "b": "b2"}); err != nil {
+		t.Fatal(err)
+	}
+	tr := tar.NewReader(&out)
+	var got []string
+	for {
+		h, err := tr.Next()
+		if err != nil {
+			break
+		}
+		got = append(got, h.Name)
+	}
+	if strings.Join(got, ",") != "a,b2/x" {
+		t.Fatalf("entries: %v", got)
+	}
+}

internal/agent/sshfs.go

@@ -641,11 +641,11 @@ func (a *Agent) doSSHFSPut(ctx context.Context, env proto.TaskEnvelope) proto.Ta
 	return proto.TaskResult{OK: true, Data: data}
 }
 
-// renameTar copies a tar stream, giving its top entry base - and everything
-// under it - the name name: what the host's tar packs under the file's own
-// name arrives under the name chosen for it. Entries other than base (the
+// renameTar copies a tar stream, giving each top entry named in names - and
+// everything under it - its new name: what the host's tar packs under the
+// file's own name arrives under the name chosen for it. Other entries (the
 // "._" companions a macOS tar adds) are left out.
-func renameTar(w io.Writer, r io.Reader, base, name string) error {
+func renameTar(w io.Writer, r io.Reader, names map[string]string) error {
 	tr, tw := tar.NewReader(r), tar.NewWriter(w)
 	for {
 		h, err := tr.Next()
@@ -661,7 +661,8 @@ func renameTar(w io.Writer, r io.Reader, base, name string) error {
 			top, rest = n[:i], n[i:]
 		}
 		// Only base and what is under it - nothing that climbs out by "..".
-		if top != base || strings.Contains("/"+rest+"/", "/../") {
+		name, ok := names[top]
+		if !ok || strings.Contains("/"+rest+"/", "/../") {
 			continue
 		}
 		h.Name = name + rest
@@ -677,7 +678,7 @@ func renameTar(w io.Writer, r io.Reader, base, name string) error {
 
 // packRemote writes src on the connection's host as a tar stream named name,
 // counting what goes for the progress.
-func (a *Agent) packRemote(ctx context.Context, w io.Writer, c *sshfsConn, src, name string, m *shipMeter) error {
+func (a *Agent) packRemote(ctx context.Context, w io.Writer, c *sshfsConn, srcs, names []string, m *shipMeter) error {
 	defer c.hold()()
 	s, err := a.sshfsSession(c)
 	if err != nil {
@@ -688,11 +689,19 @@ func (a *Agent) packRemote(ctx context.Context, w io.Writer, c *sshfsConn, src,
 	out, _ := s.StdoutPipe()
 	var errb bytes.Buffer
 	s.Stderr = &limitBuf{b: &errb, max: 4096}
-	src = path.Clean(src)
-	if err := s.Start("COPYFILE_DISABLE=1 tar -cf - -C " + shq(path.Dir(src)) + " -- " + shq(path.Base(src))); err != nil {
+	// All from one folder: packed from there by their own names.
+	dir := path.Dir(path.Clean(srcs[0]))
+	rename := map[string]string{}
+	cmd := "COPYFILE_DISABLE=1 tar -cf - -C " + shq(dir) + " --"
+	for i, p := range srcs {
+		b := path.Base(path.Clean(p))
+		rename[b] = names[i]
+		cmd += " " + shq(b)
+	}
+	if err := s.Start(cmd); err != nil {
 		return err
 	}
-	if err := renameTar(w, &countReader{r: out, n: &m.done}, path.Base(src), name); err != nil {
+	if err := renameTar(w, &countReader{r: out, n: &m.done}, rename); err != nil {
 		return err
 	}
 	if err := s.Wait(); err != nil {
@@ -721,34 +730,52 @@ func (a *Agent) remoteSize(c *sshfsConn, src string) int64 {
 // unpackRemote puts a tar stream into dir on the connection's host, by the
 // conflict rule: tar keeps or replaces what is there, an "error" rule
 // refuses when name is there already.
-func (a *Agent) unpackRemote(c *sshfsConn, r io.Reader, dir, name, conflict string) (skipped bool, err error) {
+func (a *Agent) unpackRemote(c *sshfsConn, r io.Reader, dir string, names []string, conflict string) (skipped int, err error) {
 	defer c.hold()()
-	target := strings.TrimSuffix(dir, "/") + "/" + name
-	out, _ := a.sshfsRun(c, "test -e "+shq(target)+" -o -L "+shq(target)+" && echo LPRA_EXISTS; tar --version 2>/dev/null | head -n 1", nil)
-	if strings.Contains(string(out), "LPRA_EXISTS") {
-		switch conflict {
-		case "skip":
-			// Left as it was, as on a machine.
-			return true, nil
-		case "replace":
-		default:
-			return false, errors.New(target + " уже существует")
-		}
+	base := strings.TrimSuffix(dir, "/") + "/"
+	check := ""
+	for _, n := range names {
+		check += "test -e " + shq(base+n) + " -o -L " + shq(base+n) + " && echo " + shq("LPRA_EXISTS "+n) + "; "
 	}
-	// Only name and what is under it reaches the host's tar.
+	out, _ := a.sshfsRun(c, check+"true", nil)
+	keep := map[string]string{}
+	var put []string
+	for _, n := range names {
+		if strings.Contains(string(out), "LPRA_EXISTS "+n+"\n") {
+			switch conflict {
+			case "skip":
+				// Left as it was, as on a machine.
+				skipped++
+				continue
+			case "replace":
+			default:
+				return 0, errors.New(base + n + " уже существует")
+			}
+		}
+		keep[n] = n
+		put = append(put, n)
+	}
+	if len(put) == 0 {
+		return skipped, nil
+	}
+	// Only the items sent (and not kept back) reach the host's tar.
 	pr, pw := io.Pipe()
-	go func() { pw.CloseWithError(renameTar(pw, r, name, name)) }()
+	go func() { pw.CloseWithError(renameTar(pw, r, keep)) }()
 	defer pr.Close()
-	_, err = a.sshfsRun(c, unpackCommand(dir, name, conflict, strings.Contains(string(out), "GNU tar")), pr)
-	return false, err
+	_, err = a.sshfsRun(c, unpackCommand(dir, put, conflict), pr)
+	return skipped, err
 }
 
-// unpackCommand unpacks a tar stream of name into dir; "replace" removes
-// what is there first, so the thing is replaced whole as on a machine.
-func unpackCommand(dir, name, conflict string, gnu bool) string {
+// unpackCommand unpacks a tar stream of names into dir; "replace" removes
+// what is there first, so each thing is replaced whole as on a machine.
+func unpackCommand(dir string, names []string, conflict string) string {
 	cmd := "mkdir -p -- " + shq(dir) + " && cd -- " + shq(dir) + " && tar -xf -"
 	if conflict == "replace" {
-		cmd = "rm -rf -- " + shq(strings.TrimSuffix(dir, "/")+"/"+name) + " && " + cmd
+		rm := "rm -rf --"
+		for _, n := range names {
+			rm += " " + shq(strings.TrimSuffix(dir, "/")+"/"+n)
+		}
+		cmd = rm + " && " + cmd
 	}
 	return cmd
 }

internal/agent/sshfs_test.go

@@ -62,7 +62,7 @@ func TestRenameTarEntries(t *testing.T) {
 	}
 	_ = tw.Close()
 	var out bytes.Buffer
-	if err := renameTar(&out, &src, "old", "new"); err != nil {
+	if err := renameTar(&out, &src, map[string]string{"old": "new"}); err != nil {
 		t.Fatal(err)
 	}
 	tr := tar.NewReader(&out)
@@ -90,7 +90,7 @@ func TestRenameTarDropsOtherEntries(t *testing.T) {
 	}
 	_ = tw.Close()
 	var out bytes.Buffer
-	if err := renameTar(&out, &src, "r.txt", "new.txt"); err != nil {
+	if err := renameTar(&out, &src, map[string]string{"r.txt": "new.txt"}); err != nil {
 		t.Fatal(err)
 	}
 	tr := tar.NewReader(&out)
@@ -165,7 +165,7 @@ func TestRenameTarKeepsInsideName(t *testing.T) {
 	}
 	_ = tw.Close()
 	var out bytes.Buffer
-	if err := renameTar(&out, &src, "d", "d"); err != nil {
+	if err := renameTar(&out, &src, map[string]string{"d": "d"}); err != nil {
 		t.Fatal(err)
 	}
 	tr := tar.NewReader(&out)
@@ -184,11 +184,11 @@ func TestRenameTarKeepsInsideName(t *testing.T) {
 
 // Into an SSH host, "replace" replaces the thing whole, as on a machine.
 func TestUnpackCommandReplaces(t *testing.T) {
-	cmd := unpackCommand("/srv", "d", "replace", true)
+	cmd := unpackCommand("/srv", []string{"d"}, "replace")
 	if !strings.HasPrefix(cmd, "rm -rf -- '/srv/d' && ") || strings.Contains(cmd, "--overwrite") {
 		t.Fatalf("replace: %s", cmd)
 	}
-	if cmd := unpackCommand("/srv", "d", "error", false); cmd != "mkdir -p -- '/srv' && cd -- '/srv' && tar -xf -" {
+	if cmd := unpackCommand("/srv", []string{"d"}, "error"); cmd != "mkdir -p -- '/srv' && cd -- '/srv' && tar -xf -" {
 		t.Fatalf("error: %s", cmd)
 	}
 }