clonetool/agent.go

696 lines
22 KiB
Go

package main
import (
"encoding/json"
"errors"
"flag"
"fmt"
"io"
"os"
"os/exec"
"strconv"
"sync"
"sync/atomic"
"time"
)
func cmdAgent(args []string) error {
fs := flag.NewFlagSet("agent", flag.ContinueOnError)
role := fs.String("role", "", "control|sink|source-stream|fs-send|fs-recv (internal)")
path := fs.String("path", "", "path to read/write")
size := fs.Int64("size", 0, "total sync size in bytes")
base := fs.Int64("base", 0, "byte offset the window starts at (clone-disk boot region / offset partition)")
blockSize := fs.Int64("block-size", defaultBlockSize, "block size in bytes")
fsType := fs.String("fs", "", "filesystem type (fs-send/fs-recv)")
fsTool := fs.String("fstool", "", "fs-image tool family (fs-send/fs-recv)")
peerDisk := fs.String("peerdisk", "", "this helper's whole-disk/image path (fs-send/fs-recv)")
shrinkTo := fs.Int64("shrink", 0, "shrink this fs to N bytes before sending (fs-send)")
if err := fs.Parse(args); err != nil {
return err
}
switch *role {
case "control":
return runControlAgent()
case "sink":
return runSinkRole(*path, *base, *size, *blockSize)
case "source-stream":
return runSourceStreamRole(*path, *base, *size, *blockSize)
case "fs-send":
n, _ := strconv.Atoi(*path)
return runFSSendRole(n, *fsType, *fsTool, *peerDisk, *shrinkTo)
case "fs-recv":
n, _ := strconv.Atoi(*path)
return runFSRecvRole(n, *fsType, *fsTool, *peerDisk)
default:
return fmt.Errorf("agent: unknown or missing --role %q", *role)
}
}
// ---------------------------------------------------------------------
// control role: long-lived per-side orchestration agent, driven by the
// manager over stdin/stdout with CtrlMsg frames.
// ---------------------------------------------------------------------
func runControlAgent() error {
in := NewFrameReader(os.Stdin)
out := NewFrameWriter(os.Stdout)
for {
typ, payload, err := in.ReadFrame()
if err != nil {
return nil // manager closed the pipe; nothing left to do
}
if typ != frameCtrlJSON {
return fmt.Errorf("control agent: unexpected frame type %d", typ)
}
var m CtrlMsg
if err := json.Unmarshal(payload, &m); err != nil {
return fmt.Errorf("control agent: decode message: %w", err)
}
switch m.Type {
case msgStat:
handleStat(out, m)
case msgPrepare:
handlePrepare(out, m)
case msgConnectPush:
runPushDriver(m, out)
case msgConnectPull:
runPullDriver(m, out)
case msgProbeDisk:
handleProbeDisk(out, m)
case msgBuildLayout:
handleBuildLayout(out, m)
case msgClonePartition:
handleClonePartition(out, m)
case msgReinstallBoot:
handleReinstallBoot(out, m)
case msgClose:
detachAllDisks()
_ = out.WriteJSON(CtrlMsg{Type: msgBye})
return nil
default:
_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: fmt.Sprintf("unknown command %q", m.Type)})
}
}
}
func handleStat(out *FrameWriter, m CtrlMsg) {
info, err := statPath(m.Path)
if err != nil {
_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: err.Error(), NeedPriv: isPermErr(err)})
return
}
_ = out.WriteJSON(CtrlMsg{Type: msgStatOK, Exists: info.Exists, IsDevice: info.IsDevice, Size: info.Size})
}
func handlePrepare(out *FrameWriter, m CtrlMsg) {
if err := prepareDest(m.Path, m.Size); err != nil {
_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: err.Error(), NeedPriv: isPermErr(err)})
return
}
_ = out.WriteJSON(CtrlMsg{Type: msgPrepareOK, Size: m.Size})
}
// isPermErr reports whether err (or anything it wraps, including a bare
// syscall errno from an ioctl) is a permission failure — the signal that
// retrying the agent under sudo might help.
func isPermErr(err error) bool {
return errors.Is(err, os.ErrPermission)
}
// peerAgentCommand builds the *exec.Cmd for a one-shot helper (sink /
// source-stream) on the peer — locally or over ssh — wrapping it in sudo
// when req.Sudo is set (interactive sudo locally; "sudo -n" remotely, since
// a password prompt would corrupt the binary stream on stdout).
func peerAgentCommand(req CtrlMsg, tailArgs []string) *exec.Cmd {
if req.PeerLocal {
if req.Sudo {
return sudoLocalCommand(tailArgs)
}
return localAgentCommand(tailArgs)
}
peerArgs := append([]string{req.RemoteBin}, tailArgs...)
if req.Sudo {
peerArgs = append([]string{"sudo", "-n", "--"}, peerArgs...)
}
return sshCommand(req.SSHBin, req.SSHOpts, true, req.ConnectTimeoutSec, req.PeerUser, req.PeerHost, peerArgs)
}
// ---------------------------------------------------------------------
// push driver: runs inside the SOURCE control agent. Spawns ssh straight
// to the destination host, running the "sink" role, and — once it answers
// READY — performs the whole read/hash/compare/send loop itself.
// ---------------------------------------------------------------------
func runPushDriver(req CtrlMsg, out *FrameWriter) {
tailArgs := []string{
"agent", "--role", "sink",
"--path", req.PeerPath,
"--base", strconv.FormatInt(req.PeerBase, 10),
"--size", strconv.FormatInt(req.Size, 10),
"--block-size", strconv.FormatInt(req.BlockSize, 10),
}
cmd := peerAgentCommand(req, tailArgs)
stdin, err := cmd.StdinPipe()
if err != nil {
_ = out.WriteJSON(CtrlMsg{Type: msgPushFailed, Reason: err.Error()})
return
}
stdout, err := cmd.StdoutPipe()
if err != nil {
_ = out.WriteJSON(CtrlMsg{Type: msgPushFailed, Reason: err.Error()})
return
}
stderrBuf := newLimitedBuffer(4096)
cmd.Stderr = stderrBuf
if err := cmd.Start(); err != nil {
_ = out.WriteJSON(CtrlMsg{Type: msgPushFailed, Reason: fmt.Sprintf("start ssh: %v", err)})
return
}
fw := NewFrameWriter(stdin)
fr := NewFrameReader(stdout)
timeout := time.Duration(req.ConnectTimeoutSec+2) * time.Second
if err := waitReady(fr, timeout); err != nil {
_ = cmd.Process.Kill()
_ = cmd.Wait()
_ = out.WriteJSON(CtrlMsg{Type: msgPushFailed, Reason: fmt.Sprintf("%v (remote stderr: %s)", err, stderrBuf.String())})
return
}
// Handshake succeeded: we're committed to push for this run. The sink
// streams the destination's current per-block hashes as it scans; the
// source loop consumes them in order and reads/compares its own blocks
// as they arrive, so the two scans and the transfer all overlap.
srcFile, err := os.Open(req.Path)
if err != nil {
_ = cmd.Process.Kill()
_ = cmd.Wait()
_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: err.Error(), NeedPriv: isPermErr(err)})
return
}
defer srcFile.Close()
if fatalErr := pumpPush(req, alignmentFor(req.Path), srcFile, fw, fr, out); fatalErr != nil {
_ = cmd.Process.Kill()
_ = cmd.Wait()
_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: fatalErr.Error()})
return
}
if err := cmd.Wait(); err != nil {
_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: fmt.Sprintf("sink process: %v (stderr: %s)", err, stderrBuf.String())})
return
}
_ = out.WriteJSON(CtrlMsg{Type: msgPushOK})
}
func waitReady(fr *FrameReader, timeout time.Duration) error {
type result struct {
typ frameType
err error
}
ch := make(chan result, 1)
go func() {
typ, _, err := fr.ReadFrame()
ch <- result{typ, err}
}()
select {
case r := <-ch:
if r.err != nil {
return fmt.Errorf("handshake failed: %w", r.err)
}
if r.typ != frameReady {
return fmt.Errorf("handshake failed: unexpected frame type %d", r.typ)
}
return nil
case <-time.After(timeout):
return fmt.Errorf("handshake timed out after %s", timeout)
}
}
type ackEvent struct {
index uint64
eof bool
err error
}
// pumpPush runs the source-side read/hash/compare/send loop against fw
// (the pipe to the remote sink) while a single reader goroutine over fr
// demultiplexes the three things the sink sends back on that one stream:
// the streamed destination block hashes (fed to the source loop as they
// arrive, so its own scan overlaps the sink's), the ACK/ERR frames for
// blocks it wrote, and any relayed scan-progress. It returns once every
// sent block's write has been confirmed.
func pumpPush(req CtrlMsg, align int64, srcFile *os.File, fw *FrameWriter, fr *FrameReader, out *FrameWriter) error {
blockCount := (req.Size + req.BlockSize - 1) / req.BlockSize
const maxInFlight = 32
sem := make(chan struct{}, maxInFlight)
var pendingMu sync.Mutex
pending := make(map[uint64]bool)
var copied, skipped, srcRead int64
var lastProgress time.Time
maybeProgress := func() {
if time.Since(lastProgress) < 500*time.Millisecond {
return
}
lastProgress = time.Now()
_ = out.WriteJSON(CtrlMsg{
Type: msgProgress, Phase: "xfer",
Copied: atomic.LoadInt64(&copied), Skipped: atomic.LoadInt64(&skipped), SrcRead: atomic.LoadInt64(&srcRead),
TotalBlocks: blockCount,
})
}
// lastActivityNano records the last moment either side made progress: a
// frame arriving from the sink (hashes, ACKs, relayed progress) or the
// source loop reading and comparing a block. The idle watchdog below
// treats a stall as fatal only when *neither* side has moved for a whole
// window — otherwise a check-only run, where every block matches and so no
// DATA frames and no ACKs are ever exchanged, trips the watchdog purely
// because the wire went quiet while the two scans churned through
// identical data.
var lastActivityNano int64
atomic.StoreInt64(&lastActivityNano, time.Now().UnixNano())
markActivity := func() { atomic.StoreInt64(&lastActivityNano, time.Now().UnixNano()) }
// hashCh is sized to hold every block hash so the reader below never
// blocks handing hashes off (which, since ACKs share the same stream,
// would otherwise be able to deadlock against the in-flight-send limit).
// This is the same order of memory the old whole-table transfer used.
hashBuf := blockCount
if hashBuf < 1 {
hashBuf = 1
}
hashCh := make(chan blockHash, hashBuf)
ackEvents := make(chan ackEvent, 256)
go func() {
hashClosed := false
closeHash := func() {
if !hashClosed {
close(hashCh)
hashClosed = true
}
}
for {
typ, payload, err := fr.ReadFrame()
if err != nil {
closeHash()
if err == io.EOF {
ackEvents <- ackEvent{eof: true}
} else {
ackEvents <- ackEvent{err: err}
}
return
}
markActivity()
switch typ {
case frameBlockHash:
idx, h, derr := decodeBlockHashFrame(payload)
if derr != nil {
closeHash()
ackEvents <- ackEvent{err: derr}
return
}
hashCh <- blockHash{index: idx, hash: h}
case frameHashDone:
closeHash()
case frameAck:
idx, derr := decodeIndexFrame(payload)
if derr != nil {
closeHash()
ackEvents <- ackEvent{err: derr}
return
}
ackEvents <- ackEvent{index: idx}
case frameErr:
idx, msg, _ := decodeErrFrame(payload)
closeHash()
ackEvents <- ackEvent{err: fmt.Errorf("remote reported error at block %d: %s", idx, msg)}
return
case frameCtrlJSON:
var m CtrlMsg
if json.Unmarshal(payload, &m) == nil && m.Type == msgProgress {
_ = out.WriteJSON(m)
}
default:
closeHash()
ackEvents <- ackEvent{err: fmt.Errorf("unexpected frame type %d from sink", typ)}
return
}
}
}()
sendErrCh := make(chan error, 1)
go func() {
sendErrCh <- runSourceLoop(sourceLoopParams{
File: srcFile, Size: req.Size, BlockSize: req.BlockSize, Align: align, Hashes: hashCh, Out: fw,
OnRead: func(uint64) { markActivity(); atomic.AddInt64(&srcRead, 1); maybeProgress() },
OnSkip: func(uint64) { markActivity(); atomic.AddInt64(&skipped, 1); maybeProgress() },
OnSend: func(idx uint64, _ [32]byte) {
markActivity()
sem <- struct{}{}
pendingMu.Lock()
pending[idx] = true
pendingMu.Unlock()
},
})
}()
// Wait for both: the send loop to finish (all DATA frames + DONE sent)
// and the ack stream to end. Sink closes its stdout (a clean EOF) only
// after it has acked every block it received, so an EOF while blocks
// are still unconfirmed is treated as a real failure below.
sendCh := sendErrCh
ackCh := ackEvents
var fatalErr error
// idleTimeout is how long *both* sides can be silent before the transfer
// is declared stalled. The timer may fire sooner; a fire is only fatal
// when lastActivityNano confirms nothing has moved for the whole window.
const idleTimeout = 120 * time.Second
idle := time.NewTimer(idleTimeout)
defer idle.Stop()
for sendCh != nil || ackCh != nil {
if !idle.Stop() {
select {
case <-idle.C:
default:
}
}
idle.Reset(idleTimeout)
select {
case sendErr := <-sendCh:
sendCh = nil
if sendErr != nil {
fatalErr = sendErr
}
case ev := <-ackCh:
switch {
case ev.err != nil:
fatalErr = ev.err
ackCh = nil
case ev.eof:
ackCh = nil
default:
pendingMu.Lock()
_, ok := pending[ev.index]
delete(pending, ev.index)
pendingMu.Unlock()
if ok {
atomic.AddInt64(&copied, 1)
}
<-sem
maybeProgress()
}
case <-idle.C:
idleFor := time.Since(time.Unix(0, atomic.LoadInt64(&lastActivityNano)))
if idleFor < idleTimeout {
continue // progress somewhere within the window; re-arm and wait
}
fatalErr = fmt.Errorf("timed out waiting for the sink (no activity for %s)", idleFor.Round(time.Second))
}
if fatalErr != nil {
break
}
}
if fatalErr != nil {
return fatalErr
}
pendingMu.Lock()
n := len(pending)
pendingMu.Unlock()
if n > 0 {
return fmt.Errorf("sink closed the connection with %d block write confirmation(s) still outstanding", n)
}
_ = out.WriteJSON(CtrlMsg{
Type: msgProgress, Phase: "xfer",
Copied: atomic.LoadInt64(&copied), Skipped: atomic.LoadInt64(&skipped), SrcRead: atomic.LoadInt64(&srcRead),
TotalBlocks: blockCount,
})
return nil
}
// ---------------------------------------------------------------------
// pull driver: runs inside the DEST control agent. It reads and hashes the
// local destination itself, spawns ssh to the source host running the
// "source-stream" role, feeds it that hash table, then writes whatever it
// streams back straight to the local destination — no round trip needed to
// confirm a write, since dest-agent itself performed it.
// ---------------------------------------------------------------------
func runPullDriver(req CtrlMsg, out *FrameWriter) {
tailArgs := []string{
"agent", "--role", "source-stream",
"--path", req.PeerPath,
"--base", strconv.FormatInt(req.PeerBase, 10),
"--size", strconv.FormatInt(req.Size, 10),
"--block-size", strconv.FormatInt(req.BlockSize, 10),
}
cmd := peerAgentCommand(req, tailArgs)
stdin, err := cmd.StdinPipe()
if err != nil {
_ = out.WriteJSON(CtrlMsg{Type: msgPullFailed, Reason: err.Error()})
return
}
stdout, err := cmd.StdoutPipe()
if err != nil {
_ = out.WriteJSON(CtrlMsg{Type: msgPullFailed, Reason: err.Error()})
return
}
stderrBuf := newLimitedBuffer(4096)
cmd.Stderr = stderrBuf
if err := cmd.Start(); err != nil {
_ = out.WriteJSON(CtrlMsg{Type: msgPullFailed, Reason: fmt.Sprintf("start ssh: %v", err)})
return
}
fw := NewFrameWriter(stdin)
fr := NewFrameReader(stdout)
timeout := time.Duration(req.ConnectTimeoutSec+2) * time.Second
if err := waitReady(fr, timeout); err != nil {
_ = cmd.Process.Kill()
_ = cmd.Wait()
_ = out.WriteJSON(CtrlMsg{Type: msgPullFailed, Reason: fmt.Sprintf("%v (remote stderr: %s)", err, stderrBuf.String())})
return
}
// Handshake succeeded. Open the destination and fingerprint its current
// content block by block, streaming each hash to the source stream the
// moment it is computed so it can start comparing straight away; write
// whatever it streams back into the same file concurrently.
dstFile, err := os.OpenFile(req.Path, os.O_RDWR, 0)
if err != nil {
_ = cmd.Process.Kill()
_ = cmd.Wait()
_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: err.Error(), NeedPriv: isPermErr(err)})
return
}
defer dstFile.Close()
align := alignmentFor(req.Path)
hashErrCh := make(chan error, 1)
go func() {
err := streamHashBlocks(dstFile, req.Base, req.Size, req.BlockSize, align, func(bh blockHash) error {
return fw.WriteFrame(frameBlockHash, encodeBlockHashFrame(bh.index, bh.hash))
}, scanProgressEmitter(out))
if err == nil {
err = fw.WriteFrame(frameHashDone, nil)
}
if err != nil {
// Unblock the source stream (waiting for more hashes) so the
// dest loop below can unwind instead of hanging.
_ = stdin.Close()
}
hashErrCh <- err
}()
var written int64
var lastWritten time.Time
loopErr := runDestLoop(destLoopParams{
File: dstFile, Base: req.Base, BlockSize: req.BlockSize, Align: align, In: fr,
OnCtrlMsg: func(m CtrlMsg) {
if m.Type == msgProgress {
_ = out.WriteJSON(m)
}
},
OnWrite: func(uint64) {
written++
if time.Since(lastWritten) >= 500*time.Millisecond {
lastWritten = time.Now()
_ = out.WriteJSON(CtrlMsg{Type: msgProgress, Phase: "write", Written: written})
}
},
})
_ = out.WriteJSON(CtrlMsg{Type: msgProgress, Phase: "write", Written: written})
hashErr := <-hashErrCh
if hashErr != nil {
_ = cmd.Process.Kill()
_ = cmd.Wait()
_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: fmt.Sprintf("scan destination: %v", hashErr), NeedPriv: isPermErr(hashErr)})
return
}
if loopErr != nil {
_ = cmd.Process.Kill()
_ = cmd.Wait()
_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: loopErr.Error()})
return
}
if err := cmd.Wait(); err != nil {
_ = out.WriteJSON(CtrlMsg{Type: msgError, Message: fmt.Sprintf("source-stream process: %v (stderr: %s)", err, stderrBuf.String())})
return
}
_ = out.WriteJSON(CtrlMsg{Type: msgPullOK})
}
// ---------------------------------------------------------------------
// sink role: one-shot process spawned (over ssh, in push mode) on the
// destination host. Dumb write endpoint: verify+pwrite+ack per block.
// ---------------------------------------------------------------------
func runSinkRole(path string, base, size, blockSize int64) error {
out := NewFrameWriter(os.Stdout)
in := NewFrameReader(os.Stdin)
f, err := os.OpenFile(path, os.O_RDWR, 0)
if err != nil {
fmt.Fprintf(os.Stderr, "sink: open %s: %v\n", path, err)
return err
}
defer f.Close()
align := alignmentFor(path)
// Answer the handshake immediately so the push driver's short readiness
// timeout isn't spent hashing a large destination.
if err := out.WriteFrame(frameReady, nil); err != nil {
return err
}
// Stream the destination's per-block hashes as they're computed while the
// dest loop below concurrently receives and writes changed blocks. A
// block is always hashed before a write for it can arrive (the source
// only sends after seeing that block's hash), so the two accesses to f
// never race on the same region.
hashErrCh := make(chan error, 1)
go func() {
err := streamHashBlocks(f, base, size, blockSize, align, func(bh blockHash) error {
return out.WriteFrame(frameBlockHash, encodeBlockHashFrame(bh.index, bh.hash))
}, scanProgressEmitter(out))
if err == nil {
err = out.WriteFrame(frameHashDone, nil)
} else {
fmt.Fprintf(os.Stderr, "sink: hash %s: %v\n", path, err)
_ = out.WriteFrame(frameErr, encodeErrFrame(0, err.Error()))
}
hashErrCh <- err
}()
var written int64
var lastWritten time.Time
loopErr := runDestLoop(destLoopParams{
File: f, BlockSize: blockSize, Align: align, In: in, AckOut: out,
OnWrite: func(uint64) {
written++
if time.Since(lastWritten) >= 500*time.Millisecond {
lastWritten = time.Now()
_ = out.WriteJSON(CtrlMsg{Type: msgProgress, Phase: "write", Written: written})
}
},
})
_ = out.WriteJSON(CtrlMsg{Type: msgProgress, Phase: "write", Written: written})
if hashErr := <-hashErrCh; hashErr != nil {
return hashErr
}
return loopErr
}
// ---------------------------------------------------------------------
// source-stream role: one-shot process spawned (over ssh, in pull mode) on
// the source host. Consumes the streamed destination block hashes, then
// performs the same read/hash/compare/send loop a local push driver would,
// writing straight to its own stdout.
// ---------------------------------------------------------------------
func runSourceStreamRole(path string, base, size, blockSize int64) error {
out := NewFrameWriter(os.Stdout)
in := NewFrameReader(os.Stdin)
f, err := os.Open(path)
if err != nil {
fmt.Fprintf(os.Stderr, "source-stream: open %s: %v\n", path, err)
return err
}
defer f.Close()
align := alignmentFor(path)
if err := out.WriteFrame(frameReady, nil); err != nil {
return err
}
blockCount := (size + blockSize - 1) / blockSize
hashBuf := blockCount
if hashBuf < 1 {
hashBuf = 1
}
hashCh := make(chan blockHash, hashBuf)
readErrCh := make(chan error, 1)
go func() {
for {
typ, payload, err := in.ReadFrame()
if err != nil {
readErrCh <- fmt.Errorf("source-stream: read hash stream: %w", err)
close(hashCh)
return
}
switch typ {
case frameBlockHash:
idx, h, derr := decodeBlockHashFrame(payload)
if derr != nil {
readErrCh <- derr
close(hashCh)
return
}
hashCh <- blockHash{index: idx, hash: h}
case frameHashDone:
readErrCh <- nil
close(hashCh)
return
default:
readErrCh <- fmt.Errorf("source-stream: unexpected frame type %d", typ)
close(hashCh)
return
}
}
}()
var copied, skipped, srcRead int64
var lastProgress time.Time
maybeProgress := func() {
if time.Since(lastProgress) < 500*time.Millisecond {
return
}
lastProgress = time.Now()
_ = out.WriteJSON(CtrlMsg{Type: msgProgress, Phase: "xfer", Copied: copied, Skipped: skipped, SrcRead: srcRead, TotalBlocks: blockCount})
}
loopErr := runSourceLoop(sourceLoopParams{
File: f, Size: size, BlockSize: blockSize, Align: align, Hashes: hashCh, Out: out,
OnRead: func(uint64) { srcRead++; maybeProgress() },
OnSkip: func(uint64) { skipped++; maybeProgress() },
OnSend: func(uint64, [32]byte) { copied++; maybeProgress() },
})
_ = out.WriteJSON(CtrlMsg{Type: msgProgress, Phase: "xfer", Copied: copied, Skipped: skipped, SrcRead: srcRead, TotalBlocks: blockCount})
if readErr := <-readErrCh; readErr != nil {
return readErr
}
return loopErr
}