lpra/internal/agent/ship.go

package agent

import (
	"archive/tar"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/x-grk/longpolling-remote-access/internal/proto"
)

// A file or folder goes to another machine as a tar stream: packed here on
// the fly, then either posted in pieces to a relay on the server for the
// other agent to unpack, or written by SSH straight into tar on that host.
// Nothing is staged on disk on either side.

// shipMeter counts what has gone out, for the panel's progress bar.
type shipMeter struct {
	total, done atomic.Int64
	files       atomic.Int64
	note        atomic.Value // string
}

func (m *shipMeter) snapshot() proto.ShipProgress {
	n, _ := m.note.Load().(string)
	return proto.ShipProgress{Total: m.total.Load(), Done: m.done.Load(), Files: int(m.files.Load()), Note: n}
}

// reportShip posts the progress to the job until ctx ends; the server
// answering 410 means the job was cancelled, and cancel is called.
func (a *Agent) reportShip(ctx context.Context, job string, m *shipMeter, cancel func()) {
	if job == "" {
		return
	}
	t := time.NewTicker(700 * time.Millisecond)
	defer t.Stop()
	for {
		body, _ := json.Marshal(m.snapshot())
		req, _ := http.NewRequestWithContext(ctx, http.MethodPost, a.host()+"/api/agent/fs/job/"+job+"/progress", bytes.NewReader(body))
		a.authHeaders(req)
		req.Header.Set("Content-Type", "application/json")
		if resp, err := a.http.Do(req); err == nil {
			io.Copy(io.Discard, resp.Body)
			resp.Body.Close()
			if resp.StatusCode == http.StatusGone {
				cancel()
				return
			}
		}
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		}
	}
}

// shipName checks the name the thing gets at the other end.
func shipName(name string) error {
	if name == "" || name == "." || name == ".." || strings.ContainsAny(name, "/\x00") {
		return errors.New("недопустимое имя")
	}
	return nil
}

// treeSize is how many bytes of files are under p.
func treeSize(ctx context.Context, p string) int64 {
	var n int64
	_ = filepath.WalkDir(p, func(_ string, d fs.DirEntry, err error) error {
		if err != nil || ctx.Err() != nil {
			return nil
		}
		if d.Type().IsRegular() {
			if info, err := d.Info(); err == nil {
				n += info.Size()
			}
		}
		return nil
	})
	return n
}

// 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
		}
		if err := ctx.Err(); err != nil {
			return err
		}
		info, err := os.Lstat(p)
		if err != nil {
			return err
		}
		link := ""
		switch {
		case info.Mode()&fs.ModeSymlink != 0:
			if link, err = os.Readlink(p); err != nil {
				return err
			}
		case info.IsDir(), info.Mode().IsRegular():
		default:
			return nil // sockets, devices, pipes: not sent
		}
		hdr, err := tar.FileInfoHeader(info, link)
		if err != nil {
			return err
		}
		rel, _ := filepath.Rel(src, p)
		hdr.Name = name
		if rel != "." {
			hdr.Name = path.Join(name, filepath.ToSlash(rel))
		}
		if info.IsDir() {
			hdr.Name += "/"
		}
		hdr.Uid, hdr.Gid, hdr.Uname, hdr.Gname = 0, 0, "", ""
		hdr.Format = tar.FormatPAX
		if err := tw.WriteHeader(hdr); err != nil {
			return err
		}
		if info.Mode().IsRegular() {
			f, err := os.Open(p)
			if err != nil {
				return err
			}
			_, err = io.CopyN(tw, &countReader{r: f, n: &m.done}, hdr.Size)
			f.Close()
			if err != nil {
				return fmt.Errorf("%s: %v", p, err)
			}
		}
		m.files.Add(1)
		return nil
	})
	return err
}

type countReader struct {
	r io.Reader
	n *atomic.Int64
}

func (c *countReader) Read(p []byte) (int, error) {
	n, err := c.r.Read(p)
	c.n.Add(int64(n))
	return n, err
}

// doShipOut sends Src to the other machine: by the relay, or by SSH.
func (a *Agent) doShipOut(parent context.Context, env proto.TaskEnvelope) proto.TaskResult {
	var args proto.FSShipArgs
	if err := json.Unmarshal(env.Args, &args); 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.
		c, err := sshfsGet(args.Conn)
		if err != nil {
			return proto.TaskResult{Error: err.Error()}
		}
		conn, src = c, path.Clean(args.Src)
	} else if _, err := os.Lstat(src); err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	ctx, cancel := context.WithCancel(parent)
	defer cancel()
	m := &shipMeter{}
	m.note.Store("подсчёт")
	go a.reportShip(ctx, args.Job, m, cancel)
	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]}
	}
	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, srcs, names, m) }
	}

	var err error
	if args.SSH != nil {
		err = a.shipSSH(ctx, args, src, m)
	} else {
		m.note.Store("")
		err = a.shipRelay(ctx, args, pack)
	}
	if err != nil {
		if ctx.Err() != nil && parent.Err() == nil {
			err = errors.New("отменено")
		}
		return proto.TaskResult{Error: err.Error()}
	}
	res := map[string]any{"files": m.files.Load(), "bytes": m.done.Load()}
	if args.Move && args.SSH != nil {
		if err := os.RemoveAll(src); err != nil {
			return proto.TaskResult{Error: "отправлено, но здесь не удалилось: " + err.Error()}
		}
	}
	data, _ := json.Marshal(res)
	return proto.TaskResult{OK: true, Data: data}
}

func (a *Agent) relayURL(id string, i int) string {
	return a.host() + "/api/agent/fs/relay/" + id + "/" + strconv.Itoa(i)
}

// shipRelay posts the tar stream in pieces, a few at once; the last carries
// X-Last.
func (a *Agent) shipRelay(ctx context.Context, args proto.FSShipArgs, pack func(io.Writer) error) error {
	if args.Relay == "" || args.Chunk <= 0 {
		return errors.New("не задана передача")
	}
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	pr, pw := io.Pipe()
	go func() { pw.CloseWithError(pack(pw)) }()
	defer pr.Close()

	type piece struct {
		i    int
		data []byte
		last bool
	}
	ch := make(chan piece, 1)
	var (
		mu    sync.Mutex
		first error
		wg    sync.WaitGroup
	)
	setErr := func(err error) {
		mu.Lock()
		if first == nil {
			first = err
			cancel()
		}
		mu.Unlock()
	}
	for w := 0; w < xferWorkers; w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for p := range ch {
				err := a.xferRequest(ctx, func(ctx context.Context) (*http.Request, error) {
					req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.relayURL(args.Relay, p.i), bytes.NewReader(p.data))
					if err == nil {
						req.Header.Set("Content-Type", "application/octet-stream")
						if p.last {
							req.Header.Set("X-Last", "1")
						}
					}
					return req, err
				}, func(*http.Response) error { return nil })
				if err != nil {
					setErr(err)
				}
			}
		}()
	}
	for i := 0; ; i++ {
		buf := make([]byte, args.Chunk)
		n, err := io.ReadFull(pr, buf)
		last := err == io.EOF || err == io.ErrUnexpectedEOF
		if err != nil && !last {
			setErr(err)
			break
		}
		select {
		case ch <- piece{i, buf[:n], last}:
		case <-ctx.Done():
		}
		if last || ctx.Err() != nil {
			break
		}
	}
	close(ch)
	wg.Wait()
	if first != nil {
		var gone errXferGone
		if errors.As(first, &gone) {
			return errors.New(string(gone))
		}
		return first
	}
	return ctx.Err()
}

// doShipIn takes the tar stream from the relay and unpacks it into Dir.
func (a *Agent) doShipIn(ctx context.Context, env proto.TaskEnvelope) proto.TaskResult {
	var args proto.FSShipArgs
	if err := json.Unmarshal(env.Args, &args); 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
	if args.Conn != "" {
		// Into an SSH host this agent holds a connection to.
		c, err := sshfsGet(args.Conn)
		if err != nil {
			return proto.TaskResult{Error: err.Error()}
		}
		conn, dir = c, path.Clean(args.Dir)
	} else {
		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()}
		}
	}
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	pr, pw := io.Pipe()
	go func() {
		for i := 0; ; i++ {
			var data []byte
			var last bool
			err := a.xferRequest(ctx, func(ctx context.Context) (*http.Request, error) {
				return http.NewRequestWithContext(ctx, http.MethodGet, a.relayURL(args.Relay, i), nil)
			}, func(resp *http.Response) error {
				b, err := io.ReadAll(io.LimitReader(resp.Body, int64(args.Chunk)+1))
				data, last = b, resp.Header.Get("X-Last") != ""
				return err
			})
			if err == nil {
				_, err = pw.Write(data)
			}
			if err != nil {
				pw.CloseWithError(err)
				return
			}
			if last {
				pw.Close()
				return
			}
		}
	}()
	if conn != nil {
		skipped, err := a.unpackRemote(conn, pr, dir, names, args.Conflict)
		if err == nil {
			_, _ = io.Copy(io.Discard, pr)
		}
		pr.CloseWithError(errors.New("распаковка закончена"))
		if err != nil {
			return proto.TaskResult{Error: err.Error()}
		}
		res := map[string]int{}
		if skipped > 0 {
			res["skipped"] = skipped
		}
		data, _ := json.Marshal(res)
		return proto.TaskResult{OK: true, Data: data}
	}
	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 {
			if st, ok := info.Sys().(*syscall.Stat_t); ok {
				u.uid, u.gid, u.chown = int(st.Uid), int(st.Gid), true
			}
		}
	}
	err := u.unpack(tar.NewReader(pr))
	if err == nil {
		// Whatever follows the archive's end is read off, so the sender
		// finishes.
		_, err = io.Copy(io.Discard, pr)
	}
	pr.CloseWithError(errors.New("распаковка закончена"))
	if err != nil {
		var gone errXferGone
		if errors.As(err, &gone) {
			err = errors.New(string(gone))
		}
		return proto.TaskResult{Error: err.Error()}
	}
	data, _ := json.Marshal(map[string]int{"done": u.done, "skipped": u.skipped})
	return proto.TaskResult{OK: true, Data: data}
}

// unpacker writes a tar stream into dir, every entry under name, following
// the conflict rule for what is there already: folders merge, and files are
// replaced or kept.
type unpacker struct {
	dir, conflict string
	names         map[string]bool // the top entries the stream may carry
	uid, gid      int
	chown         bool
	done, skipped int
}

func (u *unpacker) unpack(tr *tar.Reader) error {
	var skipUnder string
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		clean := path.Clean(hdr.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))
		if skipUnder != "" && strings.HasPrefix(target, skipUnder) {
			continue
		}
		have, err := os.Lstat(target)
		exists := err == nil
		if exists && !(hdr.Typeflag == tar.TypeDir && have.IsDir()) {
			switch u.conflict {
			case "skip":
				u.skipped++
				if hdr.Typeflag == tar.TypeDir {
					skipUnder = target + string(filepath.Separator)
				}
				continue
			case "replace":
				if have.IsDir() || hdr.Typeflag != tar.TypeReg {
					if err := os.RemoveAll(target); err != nil {
						return err
					}
				}
			default:
				return fmt.Errorf("%s уже существует", target)
			}
		}
		mode := os.FileMode(hdr.Mode).Perm()
		switch hdr.Typeflag {
		case tar.TypeDir:
			if !exists || !have.IsDir() {
				if err := os.MkdirAll(target, mode|0o700); err != nil {
					return err
				}
				_ = os.Chmod(target, mode)
				u.own(target)
			}
		case tar.TypeReg:
			if err := u.writeFile(target, tr, mode, hdr.ModTime); err != nil {
				return err
			}
		case tar.TypeSymlink:
			if err := os.Symlink(hdr.Linkname, target); err != nil {
				return err
			}
			u.own(target)
		default:
			continue
		}
		u.done++
	}
}

// writeFile writes next to the target first and puts it in place when all of
// it is there.
func (u *unpacker) writeFile(target string, r io.Reader, mode os.FileMode, mtime time.Time) error {
	tmp, err := os.CreateTemp(filepath.Dir(target), "."+filepath.Base(target)+".lpra-*")
	if err != nil {
		return err
	}
	if _, err := io.Copy(tmp, r); err != nil {
		tmp.Close()
		os.Remove(tmp.Name())
		return err
	}
	if err := tmp.Close(); err != nil {
		os.Remove(tmp.Name())
		return err
	}
	_ = os.Chmod(tmp.Name(), mode)
	_ = os.Chtimes(tmp.Name(), mtime, mtime)
	u.own(tmp.Name())
	if err := os.Rename(tmp.Name(), target); err != nil {
		os.Remove(tmp.Name())
		return err
	}
	return nil
}

func (u *unpacker) own(p string) {
	if u.chown {
		_ = os.Lchown(p, u.uid, u.gid)
	}
}

// shq quotes s for a POSIX shell.
func shq(s string) string { return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'" }

// posix wraps a command for the host's sh: the user's login shell gets one
// word to hand on, so fish or csh there does not matter.
func posix(cmd string) string {
	// csh reads "!" as history even inside single quotes; '\!' is a plain
	// "!" for it and for sh alike.
	return "sh -c " + strings.ReplaceAll(shq(cmd), "!", `'\!'`)
}

// shipSSH writes the tar stream into tar on the SSH host, in Dir.
func (a *Agent) shipSSH(ctx context.Context, args proto.FSShipArgs, src string, m *shipMeter) error {
	client, _, err := a.sshConnect(ctx, *args.SSH, func(_, note string) { m.note.Store(note) })
	if err != nil {
		return err
	}
	defer client.Close()
	go func() { <-ctx.Done(); client.Close() }()
	dir := args.Dir
	if dir == "" {
		dir = "."
	}
	target := strings.TrimSuffix(dir, "/") + "/" + args.Name

	// What tar is there, and whether the target is.
	m.note.Store("проверка на той стороне")
	out, err := sshRun(client, "command -v tar >/dev/null 2>&1 || echo LPRA_NOTAR; tar --version 2>/dev/null | head -n 1; test -e "+shq(target)+" -o -L "+shq(target)+" && echo LPRA_EXISTS")
	if err != nil {
		return err
	}
	if strings.Contains(out, "LPRA_NOTAR") {
		return errors.New("на той стороне нет tar")
	}
	if strings.Contains(out, "LPRA_EXISTS") && args.Conflict != "replace" && args.Conflict != "skip" {
		return errors.New(target + " там уже существует")
	}
	gnu := strings.Contains(out, "GNU tar")
	flags := ""
	switch args.Conflict {
	case "skip":
		if gnu {
			flags = " --skip-old-files"
		} else {
			flags = " -k"
		}
	case "replace":
		if gnu {
			flags = " --overwrite"
		}
	}

	sess, err := client.NewSession()
	if err != nil {
		return err
	}
	defer sess.Close()
	var stderr bytes.Buffer
	sess.Stderr = &limitBuf{b: &stderr, max: 4096}
	stdin, err := sess.StdinPipe()
	if err != nil {
		return err
	}
	cmd := "mkdir -p -- " + shq(dir) + " && cd -- " + shq(dir) + " && tar -xf -" + flags
	if err := sess.Start(posix(cmd)); err != nil {
		return err
	}
	m.note.Store("")
	packErr := packTar(ctx, stdin, src, args.Name, m)
	stdin.Close()
	waitErr := sess.Wait()
	msg := strings.TrimSpace(stderr.String())
	if packErr != nil {
		if msg != "" {
			return fmt.Errorf("%v (%s)", packErr, msg)
		}
		return packErr
	}
	if waitErr != nil {
		var ee *ssh.ExitError
		if errors.As(waitErr, &ee) && msg != "" {
			return errors.New("tar на той стороне: " + msg)
		}
		return waitErr
	}
	return nil
}

// sshRun runs a command and returns its output.
func sshRun(c *ssh.Client, cmd string) (string, error) {
	s, err := c.NewSession()
	if err != nil {
		return "", err
	}
	defer s.Close()
	out, err := s.Output(posix(cmd))
	var ee *ssh.ExitError
	if errors.As(err, &ee) {
		err = nil // the checks may exit non-zero; what they print is the answer
	}
	return string(out), err
}

type limitBuf struct {
	b   *bytes.Buffer
	max int
}

func (l *limitBuf) Write(p []byte) (int, error) {
	if room := l.max - l.b.Len(); room > 0 {
		l.b.Write(p[:min(len(p), room)])
	}
	return len(p), nil
}