clonetool/control.go

package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
)

// errNeedPriv is wrapped into the error from a control-agent call that
// failed only because the agent lacked permission to open a device.
// bringUpController watches for it to decide whether to retry under sudo.
var errNeedPriv = errors.New("permission denied opening device")

// Controller is the manager's handle on one long-lived control agent
// (spawned locally or over ssh), used for stat/prepare/connect_push/
// connect_pull/close. The manager never does any source/dest I/O itself —
// every byte of the file/device it's syncing is read or written by an
// agent process, either this one or the sink/source-stream it spawns.
type Controller struct {
	tag  string
	cmd  *exec.Cmd
	in   io.WriteCloser
	fw   *FrameWriter
	fr   *FrameReader
	sudo bool
}

func startController(spec Spec, tag string, cfg *SyncConfig, remoteBin string, sudo bool) (*Controller, error) {
	agentArgs := []string{"agent", "--role", "control"}
	var cmd *exec.Cmd
	if spec.IsLocal() {
		if sudo {
			cmd = sudoLocalCommand(agentArgs)
		} else {
			cmd = localAgentCommand(agentArgs)
		}
	} else {
		remoteArgs := append([]string{remoteBin}, agentArgs...)
		if sudo {
			remoteArgs = append([]string{"sudo", "-n", "--"}, remoteArgs...)
		}
		cmd = sshCommand(cfg.SSHBin, cfg.SSHOpts, false, cfg.ConnectTimeoutSec, spec.User, spec.Host, remoteArgs)
	}
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, err
	}
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("start %s control agent: %w", tag, err)
	}
	go relayPrefixed(stderr, tag)

	return &Controller{tag: tag, cmd: cmd, in: stdin, fw: NewFrameWriter(stdin), fr: NewFrameReader(stdout), sudo: sudo}, nil
}

func relayPrefixed(r io.Reader, tag string) {
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 64*1024), 1024*1024)
	for sc.Scan() {
		fmt.Fprintf(os.Stderr, "[%s] %s\n", tag, sc.Text())
	}
}

func (c *Controller) readOne() (CtrlMsg, error) {
	typ, payload, err := c.fr.ReadFrame()
	if err != nil {
		return CtrlMsg{}, fmt.Errorf("%s: read control frame: %w", c.tag, err)
	}
	if typ != frameCtrlJSON {
		return CtrlMsg{}, fmt.Errorf("%s: unexpected frame type %d on control channel", c.tag, typ)
	}
	var m CtrlMsg
	if err := json.Unmarshal(payload, &m); err != nil {
		return CtrlMsg{}, fmt.Errorf("%s: decode control message: %w", c.tag, err)
	}
	return m, nil
}

func (c *Controller) call(req CtrlMsg) (CtrlMsg, error) {
	if err := c.fw.WriteJSON(req); err != nil {
		return CtrlMsg{}, fmt.Errorf("%s: send %s: %w", c.tag, req.Type, err)
	}
	return c.readOne()
}

func (c *Controller) Stat(path string) (PathInfo, error) {
	resp, err := c.call(CtrlMsg{Type: msgStat, Path: path})
	if err != nil {
		return PathInfo{}, err
	}
	switch resp.Type {
	case msgStatOK:
		return PathInfo{Exists: resp.Exists, IsDevice: resp.IsDevice, Size: resp.Size}, nil
	case msgError:
		return PathInfo{}, c.agentErr("stat", path, resp)
	default:
		return PathInfo{}, fmt.Errorf("%s: unexpected response %q to stat", c.tag, resp.Type)
	}
}

func (c *Controller) Prepare(path string, size int64) error {
	resp, err := c.call(CtrlMsg{Type: msgPrepare, Path: path, Size: size})
	if err != nil {
		return err
	}
	switch resp.Type {
	case msgPrepareOK:
		return nil
	case msgError:
		return c.agentErr("prepare", path, resp)
	default:
		return fmt.Errorf("%s: unexpected response %q to prepare", c.tag, resp.Type)
	}
}

// agentErr turns an agent's error reply into an error, tagging it with
// errNeedPriv when the agent said the cause was a permission problem on a
// device (so the caller can retry the whole agent under sudo).
func (c *Controller) agentErr(op, path string, resp CtrlMsg) error {
	if resp.NeedPriv {
		return fmt.Errorf("%s: %s %s: %s: %w", c.tag, op, path, resp.Message, errNeedPriv)
	}
	return fmt.Errorf("%s: %s %s: %s", c.tag, op, path, resp.Message)
}

// ProbeDisk asks the agent to enumerate a whole disk (partition table, boot
// regions, per-partition filesystem facts, tool availability).
func (c *Controller) ProbeDisk(path string) (*DiskLayout, error) {
	resp, err := c.call(CtrlMsg{Type: msgProbeDisk, Path: path})
	if err != nil {
		return nil, err
	}
	switch resp.Type {
	case msgProbeDiskOK:
		return parseDiskLayoutJSON(resp.LayoutJSON)
	case msgError:
		return nil, c.agentErr("probe", path, resp)
	default:
		return nil, fmt.Errorf("%s: unexpected response %q to probe_disk", c.tag, resp.Type)
	}
}

// BuildLayout asks the agent to create/size an image and write the partition
// table (optsJSON is a buildOpts blob).
func (c *Controller) BuildLayout(path, optsJSON string) error {
	resp, err := c.call(CtrlMsg{Type: msgBuildLayout, Path: path, Options: optsJSON})
	if err != nil {
		return err
	}
	switch resp.Type {
	case msgBuildLayoutOK:
		return nil
	case msgError:
		return c.agentErr("build_layout", path, resp)
	default:
		return fmt.Errorf("%s: unexpected response %q to build_layout", c.tag, resp.Type)
	}
}

// ClonePartition drives one fs-image or file-level partition clone. dir is
// "push" (send to source agent) or "pull" (send to dest agent).
func (c *Controller) ClonePartition(req CtrlMsg, cb transferCallbacks) (ok bool, reason string, err error) {
	req.Type = msgClonePartition
	okType, failType := msgPushOK, msgPushFailed
	if req.Reason == "pull" {
		okType, failType = msgPullOK, msgPullFailed
	}
	return c.connectAndPump(req, cb, okType, failType)
}

// ReinstallBoot asks the (dest) agent to run grub-install/bcdboot; the
// returned string is a human summary. Best-effort — never returns an error
// for a boot-repair failure, only for a transport failure.
func (c *Controller) ReinstallBoot(path, optsJSON string) (string, error) {
	resp, err := c.call(CtrlMsg{Type: msgReinstallBoot, Path: path, Options: optsJSON})
	if err != nil {
		return "", err
	}
	switch resp.Type {
	case msgReinstallBootOK:
		return resp.Message, nil
	case msgError:
		return "", c.agentErr("reinstall_boot", path, resp)
	default:
		return "", fmt.Errorf("%s: unexpected response %q to reinstall_boot", c.tag, resp.Type)
	}
}

// transferCallbacks receives streaming updates while a connect_push or
// connect_pull is in flight.
type transferCallbacks struct {
	onProgress func(CtrlMsg)
}

// ConnectPush asks the source control agent to try connecting straight out
// to the destination host and driving the whole transfer itself. ok=false
// with a non-empty reason means the SSH handshake didn't succeed (caller
// should try ConnectPull instead); a non-nil err means something failed
// after the transfer was already committed.
func (c *Controller) ConnectPush(req CtrlMsg, cb transferCallbacks) (ok bool, reason string, err error) {
	req.Type = msgConnectPush
	return c.connectAndPump(req, cb, msgPushOK, msgPushFailed)
}

// ConnectPull asks the destination control agent to try connecting out to
// the source host and pulling the transfer.
func (c *Controller) ConnectPull(req CtrlMsg, cb transferCallbacks) (ok bool, reason string, err error) {
	req.Type = msgConnectPull
	return c.connectAndPump(req, cb, msgPullOK, msgPullFailed)
}

func (c *Controller) connectAndPump(req CtrlMsg, cb transferCallbacks, okType, failedType string) (bool, string, error) {
	if err := c.fw.WriteJSON(req); err != nil {
		return false, "", fmt.Errorf("%s: send %s: %w", c.tag, req.Type, err)
	}
	for {
		m, err := c.readOne()
		if err != nil {
			return false, "", err
		}
		switch m.Type {
		case msgProgress:
			if cb.onProgress != nil {
				cb.onProgress(m)
			}
		case okType:
			return true, "", nil
		case failedType:
			return false, m.Reason, nil
		case msgError:
			return false, "", fmt.Errorf("%s: %s", c.tag, m.Message)
		default:
			return false, "", fmt.Errorf("%s: unexpected message %q during transfer", c.tag, m.Type)
		}
	}
}

func (c *Controller) Close() error {
	_ = c.fw.WriteJSON(CtrlMsg{Type: msgClose})
	_ = c.in.Close()
	err := c.cmd.Wait()
	if err != nil {
		var exitErr *exec.ExitError
		if errors.As(err, &exitErr) {
			return nil // agent exiting after we closed its stdin is expected
		}
		return err
	}
	return nil
}