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

package syncengine

import (
	"context"
	"encoding/base64"
	"fmt"
	"io"
	"os/exec"
	"strings"
)

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) {
	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[1:]...)
	}

	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
}

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

	flags := strings.Fields(pair.RsyncFlags)
	args = append(args, flags...)

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

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

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

	return args
}

type RemoteMachine struct {
	Host        string
	Port        int
	SSHUser     string
	PrivKey     string
	DestPrivKey string
}

func (r *RsyncRunner) RunRemote(ctx context.Context, pair *SyncPairConfig, remote RemoteMachine, onLine func(stream string, line string)) (*RsyncResult, error) {
	args := r.buildArgs(pair)

	var remoteCmd string
	if remote.DestPrivKey != "" {
		encodedKey := base64.StdEncoding.EncodeToString([]byte(remote.DestPrivKey))
		innerSSH := fmt.Sprintf(`ssh -i /tmp/syncserver-dest-key -o StrictHostKeyChecking=accept-new -o UserKnownHostsFile=/dev/null`)
		rsyncPart := fmt.Sprintf("rsync %s -e %q %s %s",
			strings.Join(args, " "), innerSSH, pair.Source, pair.Dest)
		remoteCmd = fmt.Sprintf(
			`echo '%s' | base64 -d > /tmp/syncserver-dest-key && chmod 600 /tmp/syncserver-dest-key && %s; STATUS=$?; rm -f /tmp/syncserver-dest-key; exit $STATUS`,
			encodedKey, rsyncPart)
	} else {
		remoteCmd = "rsync " + strings.Join(args, " ")
	}

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

	cmd := exec.CommandContext(ctx, "ssh", sshArgs...)

	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 ssh: %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
}