move-data-nas/internal/syncengine/rsync_runner.go

package syncengine

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
)

var allowedRsyncFlags = map[string]bool{
	"-v":              true,
	"-vv":             true,
	"-q":              true,
	"-h":              true,
	"-P":              true,
	"-n":              true,
	"-z":              true,
	"-c":              true,
	"-u":              true,
	"-W":              true,
	"-i":              true,
	"-a":              true,
	"-r":              true,
	"-l":              true,
	"-t":              true,
	"-p":              true,
	"-g":              true,
	"-o":              true,
	"-D":              true,
	"--verbose":       true,
	"--quiet":         true,
	"--help":          true,
	"--partial":       true,
	"--partial-dir":   true,
	"--delay-updates": true,
	"--delete":        true,
	"--delete-before": true,
	"--delete-after":  true,
	"--delete-excluded": true,
	"--exclude":       true,
	"--exclude-from":  true,
	"--dry-run":       true,
	"--compress":      true,
	"--skip-compress": true,
	"--whole-file":    true,
	"--checksum":     true,
	"--update":        true,
	"--existing":      true,
	"--ignore-existing": true,
	"--remove-source-files": true,
	"--chmod":         true,
	"--owner":         true,
	"--group":         true,
	"--perms":         true,
	"--executability":  true,
	"--acls":          true,
	"--xattrs":        true,
	"--numeric-ids":   true,
	"--fake-super":    true,
	"--bwlimit":       true,
	"--max-size":      true,
	"--min-size":      true,
	"--append":        true,
	"--append-verify": true,
	"--itemize-changes": true,
}

var blockedRsyncFlags = map[string]bool{
	"--rsync-path":   true,
	"-e":            true,
	"--files-from":  true,
	"--read-batch":  true,
	"--write-batch": true,
	"--log-file":    true,
}

func isSafeRsyncFlag(flag string) bool {
	if allowedRsyncFlags[flag] {
		return true
	}
	safePrefixes := []string{
		"-a", "-v", "-z", "-P", "-n", "-c", "-u", "-W", "-i",
		"--exclude=", "--chmod=", "--bwlimit=", "--max-size=", "--min-size=",
		"--partial-dir=", "--skip-compress=",
	}
	for _, p := range safePrefixes {
		if strings.HasPrefix(flag, p) {
			return true
		}
	}
	return false
}

type RsyncResult struct {
	ExitCode   int
	Stdout     string
	Stderr     string
	Stats      *RsyncStats
}

type RsyncRunner struct {
	sshDir  string
	privKey string
}

type SyncPairConfig struct {
	ID              int64
	Name            string
	SourceMachineID *int64
	SourcePath      string
	DestMachineID   *int64
	DestPath        string
	Direction       string
	RsyncFlags      string
	ExcludePatterns []string
	Source          string
	Dest            string
}

func NewRsyncRunner(sshDir, privKey string) *RsyncRunner {
	return &RsyncRunner{sshDir: sshDir, privKey: privKey}
}

func (r *RsyncRunner) Run(ctx context.Context, pair *SyncPairConfig, onLine func(stream string, line string)) (*RsyncResult, error) {
	cmd := r.buildRsyncCmd(ctx, pair)
	return r.runCmd(ctx, cmd, onLine)
}

func (r *RsyncRunner) buildRsyncCmd(ctx context.Context, pair *SyncPairConfig) *exec.Cmd {
	args := r.buildArgs(pair)
	cmd := exec.CommandContext(ctx, "rsync", args...)
	if r.privKey != "" {
		sshCmd := fmt.Sprintf("ssh -i %s -o StrictHostKeyChecking=accept-new -o UserKnownHostsFile=%s",
			r.privKey, strings.TrimRight(r.sshDir, "/")+"/known_hosts")
		cmd.Args = append([]string{"rsync", "-e", sshCmd}, args...)
	}
	return cmd
}

func (r *RsyncRunner) buildArgs(pair *SyncPairConfig) []string {
	var args []string

	flags := strings.Fields(pair.RsyncFlags)
	for _, flag := range flags {
		if strings.HasPrefix(flag, "-") {
			if blockedRsyncFlags[flag] {
				slog.Warn("blocked dangerous rsync flag", "flag", flag)
				continue
			}
			if !isSafeRsyncFlag(flag) {
				slog.Warn("disallowed rsync flag", "flag", flag)
				continue
			}
		}
		args = append(args, flag)
	}

	for _, pattern := range pair.ExcludePatterns {
		args = append(args, "--exclude="+pattern)
	}

	if pair.Direction == "mirror" {
		args = append(args, "--delete")
	}

	args = append(args, "--")
	src := ensureDirSlash(pair.Source)
	if pair.Direction == "pull" {
		args = append(args, pair.Dest, src)
	} else {
		args = append(args, src, pair.Dest)
	}

	return args
}

// ensureDirSlash guarantees the source path is treated by rsync as a
// directory whose contents are copied, regardless of whether the user
// supplied a trailing slash. This avoids the common foot-gun where
// "rsync host:/path/series /dest/" creates /dest/series/<contents> nested
// inside an extra "series" subdirectory.
func ensureDirSlash(p string) string {
	if strings.HasSuffix(p, "/") {
		return p
	}
	return p + "/"
}

type MachineKeys struct {
	Host    string
	Port    int
	SSHUser string
	PrivKey string
}

func (r *RsyncRunner) RunRemote(ctx context.Context, pair *SyncPairConfig, src *MachineKeys, dst *MachineKeys, destKey string, onLine func(stream string, line string)) (*RsyncResult, error) {
	if src.Port == 0 {
		src.Port = 22
	}
	if src.PrivKey == "" {
		src.PrivKey = filepath.Join(r.sshDir, "id_ed25519")
	}
	if destKey == "" {
		destKey = filepath.Join(r.sshDir, "id_ed25519")
	}

	args := r.buildArgs(pair)

	innerSSH := fmt.Sprintf("ssh -i %s -o StrictHostKeyChecking=accept-new -o UserKnownHostsFile=%s",
		destKey, filepath.Join(r.sshDir, "known_hosts"))
	rsyncFlags := args[:len(args)-2]
	sourcePath := args[len(args)-2]
	destPath := args[len(args)-1]

	var rsyncCmd []string
	rsyncCmd = append(rsyncCmd, "rsync")
	rsyncCmd = append(rsyncCmd, "-e")
	rsyncCmd = append(rsyncCmd, innerSSH)
	rsyncCmd = append(rsyncCmd, rsyncFlags...)
	rsyncCmd = append(rsyncCmd, sourcePath, destPath)

	remoteCmd := "rsync"
	for _, arg := range rsyncFlags {
		remoteCmd += " " + strconv.Quote(arg)
	}
	remoteCmd += " -e " + strconv.Quote(innerSSH) + " " + strconv.Quote(sourcePath) + " " + strconv.Quote(destPath)
	remoteCmd = "sh -c " + strconv.Quote(remoteCmd)

	sshArgs := []string{
		"-i", src.PrivKey,
		"-o", "StrictHostKeyChecking=accept-new",
		"-o", "UserKnownHostsFile=" + filepath.Join(r.sshDir, "known_hosts"),
		"-p", fmt.Sprintf("%d", src.Port),
		fmt.Sprintf("%s@%s", src.SSHUser, src.Host),
	}
	sshArgs = append(sshArgs, remoteCmd)

	cmd := exec.CommandContext(ctx, "ssh", sshArgs...)
	return r.runCmd(ctx, cmd, onLine)
}

func (r *RsyncRunner) runCmd(ctx context.Context, cmd *exec.Cmd, onLine func(stream string, line string)) (*RsyncResult, error) {
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("stdout pipe: %w", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("stderr pipe: %w", err)
	}

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("starting rsync: %w", err)
	}

	var outLines, errLines []string
	done := make(chan struct{})
	go func() {
		br := io.Reader(stdout)
		buf := make([]byte, 4096)
		for {
			n, err := br.Read(buf)
			if n > 0 {
				line := strings.TrimRight(string(buf[:n]), "\r\n")
				if line != "" {
					outLines = append(outLines, line)
					if onLine != nil {
						onLine("stdout", line)
					}
				}
			}
			if err != nil {
				break
			}
		}
		select {
		case <-done:
		default:
			close(done)
		}
	}()

	go func() {
		br := io.Reader(stderr)
		buf := make([]byte, 4096)
		for {
			n, err := br.Read(buf)
			if n > 0 {
				line := strings.TrimRight(string(buf[:n]), "\r\n")
				if line != "" {
					errLines = append(errLines, line)
					if onLine != nil {
						onLine("stderr", line)
					}
				}
			}
			if err != nil {
				break
			}
		}
		select {
		case <-done:
		default:
			close(done)
		}
	}()

	err = cmd.Wait()
	<-done

	result := &RsyncResult{
		ExitCode: 0,
		Stdout:   strings.Join(outLines, "\n"),
		Stderr:   strings.Join(errLines, "\n"),
		Stats:    ParseFinalStats(strings.Join(outLines, "\n")),
	}
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			result.ExitCode = exitErr.ExitCode()
		} else {
			result.ExitCode = -1
		}
	}

	return result, nil
}