clonetool/control.go

package main

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

// 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
}

func startController(spec Spec, tag string, cfg *SyncConfig) (*Controller, error) {
	var cmd *exec.Cmd
	if spec.IsLocal() {
		cmd = localAgentCommand([]string{"agent", "--role", "control"})
	} else {
		remoteArgs := []string{cfg.RemoteBin, "agent", "--role", "control"}
		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)}, 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{}, fmt.Errorf("%s: stat %s: %s", c.tag, path, resp.Message)
	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 fmt.Errorf("%s: prepare %s: %s", c.tag, path, resp.Message)
	default:
		return fmt.Errorf("%s: unexpected response %q to prepare", c.tag, resp.Type)
	}
}

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

// 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, hashes [][32]byte, cb transferCallbacks) (ok bool, reason string, err error) {
	req.Type = msgConnectPush
	return c.connectAndPump(req, hashes, 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, hashes [][32]byte, cb transferCallbacks) (ok bool, reason string, err error) {
	req.Type = msgConnectPull
	return c.connectAndPump(req, hashes, cb, msgPullOK, msgPullFailed)
}

func (c *Controller) connectAndPump(req CtrlMsg, hashes [][32]byte, 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)
	}
	if err := c.fw.WriteFrame(frameHashTable, flattenHashes(hashes)); err != nil {
		return false, "", fmt.Errorf("%s: send hash table: %w", c.tag, 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 msgBlockDoneBatch:
			if cb.onBlockDone != nil {
				for _, e := range m.Entries {
					cb.onBlockDone(e)
				}
			}
		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
}

func flattenHashes(hashes [][32]byte) []byte {
	buf := make([]byte, len(hashes)*32)
	for i, h := range hashes {
		copy(buf[i*32:], h[:])
	}
	return buf
}

func unflattenHashes(b []byte) ([][32]byte, error) {
	if len(b)%32 != 0 {
		return nil, fmt.Errorf("hash table: %d bytes is not a multiple of 32", len(b))
	}
	out := make([][32]byte, len(b)/32)
	for i := range out {
		copy(out[i][:], b[i*32:i*32+32])
	}
	return out, nil
}