clonetool/syncside.go

package main

import (
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"time"
)

// readBlockAt reads exactly min(len(buf), totalSize-offset) bytes at offset,
// treating a fully-satisfied read as success even if the underlying
// implementation also reports io.EOF for it.
func readBlockAt(f *os.File, buf []byte, offset, totalSize int64) (int, error) {
	remaining := totalSize - offset
	if remaining <= 0 {
		return 0, io.EOF
	}
	want := int64(len(buf))
	if remaining < want {
		want = remaining
	}
	n, err := f.ReadAt(buf[:want], offset)
	if err != nil && !(err == io.EOF && int64(n) == want) {
		return n, err
	}
	return n, nil
}

// hashFileBlocks reads f in blockSize-byte blocks up to size and returns the
// SHA-256 of each. This is how a destination fingerprints its *current*
// content at the start of every sync — clonetool keeps no hash state of its
// own between runs. onProgress, if non-nil, is called with (blocksHashed,
// totalBlocks) before the first block and after each one.
func hashFileBlocks(f *os.File, size, blockSize int64, onProgress func(done, total int64)) ([][32]byte, error) {
	blockCount := (size + blockSize - 1) / blockSize
	out := make([][32]byte, blockCount)
	buf := make([]byte, blockSize)
	if onProgress != nil {
		onProgress(0, blockCount)
	}
	for i := int64(0); i < blockCount; i++ {
		n, err := readBlockAt(f, buf, i*blockSize, size)
		if err != nil {
			return nil, fmt.Errorf("hash %s at block %d: %w", f.Name(), i, err)
		}
		out[i] = sha256.Sum256(buf[:n])
		if onProgress != nil {
			onProgress(i+1, blockCount)
		}
	}
	return out, nil
}

// scanProgressEmitter returns an onProgress callback for hashFileBlocks that
// forwards "scan" phase progress onto out, throttled to ~3/second.
func scanProgressEmitter(out *FrameWriter) func(done, total int64) {
	var last time.Time
	return func(done, total int64) {
		if done != 0 && done != total && time.Since(last) < 350*time.Millisecond {
			return
		}
		last = time.Now()
		_ = out.WriteJSON(CtrlMsg{Type: msgProgress, Phase: "scan", Copied: done, TotalBlocks: total})
	}
}

// sourceLoopParams drives the single read -> hash -> compare -> maybe-send
// pass over a source path. It is used both by the standalone
// "source-stream" role (writing to its own stdout) and by a source control
// agent's push driver (writing into a spawned ssh subprocess's stdin).
type sourceLoopParams struct {
	File      *os.File // already opened for reading by the caller
	Size      int64
	BlockSize int64
	Hashes    [][32]byte // previous known hashes, len == block count
	Out       *FrameWriter
	OnSkip    func(index uint64)
	OnSend    func(index uint64, hash [32]byte) // called after the DATA frame is written
}

func runSourceLoop(p sourceLoopParams) error {
	f := p.File
	buf := make([]byte, p.BlockSize)
	blockCount := uint64(len(p.Hashes))
	for index := uint64(0); index < blockCount; index++ {
		offset := int64(index) * p.BlockSize
		n, err := readBlockAt(f, buf, offset, p.Size)
		if err != nil {
			return fmt.Errorf("read %s at block %d: %w", f.Name(), index, err)
		}
		hash := sha256.Sum256(buf[:n])
		if hash == p.Hashes[index] {
			if p.OnSkip != nil {
				p.OnSkip(index)
			}
			continue
		}
		if err := p.Out.WriteFrame(frameData, encodeDataFrame(index, hash, buf[:n])); err != nil {
			return fmt.Errorf("send block %d: %w", index, err)
		}
		if p.OnSend != nil {
			p.OnSend(index, hash)
		}
	}
	return p.Out.WriteFrame(frameDone, nil)
}

// destLoopParams drives the receive -> verify -> pwrite pass on the
// destination path. It is used both by the standalone "sink" role (reading
// from its own stdin, and required to ack/err back over AckOut) and by a
// dest control agent's pull driver (reading from a spawned ssh subprocess's
// stdout; no AckOut needed since the write is confirmed locally, in the
// same process, before OnWritten is called).
type destLoopParams struct {
	File      *os.File // already opened for read/write by the caller
	BlockSize int64
	In        *FrameReader
	AckOut    *FrameWriter // optional
	// OnCtrlMsg handles an interleaved frameCtrlJSON frame (used only in
	// pull mode, where the remote source-stream has no other channel back
	// to the manager for progress updates). Sink-role callers leave it nil.
	OnCtrlMsg func(CtrlMsg)
}

func runDestLoop(p destLoopParams) error {
	f := p.File
	for {
		typ, payload, err := p.In.ReadFrame()
		if err != nil {
			return fmt.Errorf("read frame: %w", err)
		}
		switch typ {
		case frameDone:
			return f.Sync()
		case frameData:
			index, hash, block, err := decodeDataFrame(payload)
			if err != nil {
				return err
			}
			if sha256.Sum256(block) != hash {
				if p.AckOut != nil {
					_ = p.AckOut.WriteFrame(frameErr, encodeErrFrame(index, "hash mismatch after transfer"))
				}
				return fmt.Errorf("block %d: hash mismatch after transfer", index)
			}
			if _, err := f.WriteAt(block, int64(index)*p.BlockSize); err != nil {
				if p.AckOut != nil {
					_ = p.AckOut.WriteFrame(frameErr, encodeErrFrame(index, err.Error()))
				}
				return fmt.Errorf("write block %d: %w", index, err)
			}
			if p.AckOut != nil {
				if err := p.AckOut.WriteFrame(frameAck, encodeIndexFrame(index)); err != nil {
					return err
				}
			}
		case frameCtrlJSON:
			if p.OnCtrlMsg != nil {
				var m CtrlMsg
				if err := json.Unmarshal(payload, &m); err == nil {
					p.OnCtrlMsg(m)
				}
			}
		default:
			return fmt.Errorf("unexpected frame type %d on data channel", typ)
		}
	}
}