imap-copier/internal/imapx/copy.go

package imapx

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"strings"
	"time"

	"github.com/emersion/go-imap/v2"
	"github.com/emersion/go-imap/v2/imapclient"
)

// CopyDeps injects the dedup/progress hooks used by CopyFolder. APPEND to
// dst always happens before MarkMigrated is called, so a crash between the
// two only ever causes a message to be re-copied (never lost) on the next
// run.
type CopyDeps struct {
	IsMigrated   func(key string) (bool, error)
	MarkMigrated func(folder, key string) error
	OnProgress   func(copied, skipped int)
	// OnFolder is called once per folder right after EXAMINE, before the
	// (potentially long) envelope fetch, with the message count in the source
	// folder — for progress visibility.
	OnFolder func(srcFolder, dstFolder string, total int64)
	// OnScan is called during the streaming metadata pass with how many of the
	// folder's messages have been examined so far — so the UI shows movement
	// while dedup decisions are made, before bodies start copying.
	OnScan func(scanned, total int64)
	// OnError is called for each message-level error, with a message reference
	// ("UID N: subject", empty when the envelope is unavailable) and the error
	// text — so the orchestrator can persist individual errors for the
	// per-account error modal. Folder-level errors are reported by the caller.
	OnError func(ref, msg string)
	// OnActivity is called repeatedly WHILE a single message body is streamed
	// (each FETCH read chunk and each APPEND write chunk). Copying one large
	// message can take longer than the orchestrator's stall timeout; without an
	// in-body signal the watchdog can't tell a slow-but-live transfer from a
	// wedged connection and cancels a healthy copy. May be nil.
	OnActivity func()
}

// CopyResult summarizes the outcome of one CopyFolder run.
type CopyResult struct {
	Copied  int
	Skipped int
	Errors  int
}

// metaScanBatch bounds how many messages one Pass-1 metadata FETCH covers. A
// single unbounded FETCH 1:* over a large mailbox keeps one command open for
// the entire scan; under parallel load the server can stop responding and,
// since go-imap has no per-command deadline, the worker wedges forever. Short
// windows keep each command brief so the server stays responsive and ctx is
// checked between windows.
const metaScanBatch = 1000

// msgRef builds a human-readable reference for a message error: its UID plus
// subject when known, e.g. "UID 42: Invoice". Falls back to just the UID.
func msgRef(uid imap.UID, subject string) string {
	if subject == "" {
		return fmt.Sprintf("UID %d", uid)
	}
	return fmt.Sprintf("UID %d: %s", uid, subject)
}

// metaBatches tiles 1..total into contiguous, non-overlapping windows of at
// most batchSize, covering every sequence number exactly once.
func metaBatches(total, batchSize uint32) []imap.SeqRange {
	if total == 0 || batchSize == 0 {
		return nil
	}
	var out []imap.SeqRange
	for start := uint32(1); start <= total; start += batchSize {
		stop := start + batchSize - 1
		if stop > total {
			stop = total
		}
		out = append(out, imap.SeqRange{Start: start, Stop: stop})
		if stop == total {
			break // guard against uint32 overflow when total is near max
		}
	}
	return out
}

// CopyFolder streams messages from srcFolder on src to dstFolder on dst.
//
// The source folder is opened read-only (EXAMINE) and is never mutated:
// no \Deleted flags are set and no EXPUNGE is issued. Each message body is
// held in memory only for the duration of a single FETCH->APPEND and is
// never written to disk. Messages already migrated (per deps.IsMigrated)
// are skipped without re-fetching their bodies.
func CopyFolder(ctx context.Context, src, dst *imapclient.Client, srcFolder, dstFolder string, deps CopyDeps) (CopyResult, error) {
	var res CopyResult

	sel, err := src.Select(srcFolder, &imap.SelectOptions{ReadOnly: true}).Wait()
	if err != nil {
		return res, fmt.Errorf("examine src %q: %w", srcFolder, err)
	}
	total := int64(sel.NumMessages)
	if deps.OnFolder != nil {
		deps.OnFolder(srcFolder, dstFolder, total)
	}
	if total == 0 {
		return res, nil
	}

	// dst folder must exist (idempotent create; ignore "already exists").
	_ = dst.Create(dstFolder, nil).Wait()

	// Pass 1: STREAM metadata (no bodies) via Next(), dedup as we go, and queue
	// only the new messages. Streaming (not Collect) means progress shows during
	// the scan and memory stays flat — we hold small meta for new messages only.
	type queued struct {
		uid          imap.UID
		key          string
		subject      string
		flags        []imap.Flag
		internalDate time.Time
	}
	var todo []queued
	var scanned int64
	reportErr := func(ref, msg string) {
		if deps.OnError != nil {
			deps.OnError(ref, msg)
		}
	}
	// Scan metadata in bounded windows instead of one FETCH 1:*, so each
	// command is short (the server stays responsive) and ctx is checked on
	// every window boundary — not just between messages of one giant command.
	for _, win := range metaBatches(sel.NumMessages, metaScanBatch) {
		if err := ctx.Err(); err != nil {
			return res, err
		}
		fc := src.Fetch(imap.SeqSet{win}, &imap.FetchOptions{
			UID: true, Envelope: true, RFC822Size: true, Flags: true, InternalDate: true,
		})
		for {
			if err := ctx.Err(); err != nil {
				_ = fc.Close()
				return res, err
			}
			msg := fc.Next()
			if msg == nil {
				break
			}
			buf, err := msg.Collect()
			if err != nil {
				res.Errors++
				reportErr("", "read message metadata: "+err.Error())
				continue
			}
			scanned++
			key := MessageKey(buf.Envelope, buf.RFC822Size)
			subject := ""
			if buf.Envelope != nil {
				subject = buf.Envelope.Subject
			}
			already, err := deps.IsMigrated(key)
			if err != nil {
				res.Errors++
				reportErr(msgRef(buf.UID, subject), "dedup lookup: "+err.Error())
			} else if already {
				res.Skipped++
				if deps.OnProgress != nil {
					deps.OnProgress(res.Copied, res.Skipped)
				}
			} else {
				todo = append(todo, queued{uid: buf.UID, key: key, subject: subject, flags: buf.Flags, internalDate: buf.InternalDate})
			}
			if deps.OnScan != nil {
				deps.OnScan(scanned, total)
			}
		}
		if err := fc.Close(); err != nil {
			return res, fmt.Errorf("fetch meta %q: %w", srcFolder, err)
		}
	}

	// Pass 2: fetch bodies for the queued (new) messages, one at a time.
	for _, q := range todo {
		if err := ctx.Err(); err != nil {
			return res, err
		}
		if err := streamOne(src, dst, dstFolder, q.uid, q.flags, q.internalDate, deps.OnActivity); err != nil {
			res.Errors++
			reportErr(msgRef(q.uid, q.subject), "copy message: "+err.Error())
			// A closed/broken connection won't recover: every remaining APPEND
			// would fail identically. Abort the folder instead of logging
			// thousands of the same error; a re-run resumes via dedup.
			if errors.Is(err, net.ErrClosed) || strings.Contains(err.Error(), "use of closed network connection") {
				return res, fmt.Errorf("dst connection lost in %q: %w", dstFolder, err)
			}
			continue
		}
		if err := deps.MarkMigrated(dstFolder, q.key); err != nil {
			res.Errors++
			reportErr(msgRef(q.uid, q.subject), "mark migrated: "+err.Error())
			continue
		}
		res.Copied++
		if deps.OnProgress != nil {
			deps.OnProgress(res.Copied, res.Skipped)
		}
	}
	return res, nil
}

// slowMessage marks how long one message's FETCH or APPEND phase may take
// before it is logged as anomalous. Well below the orchestrator's 3-minute
// stall timeout, so a message that trips the watchdog always leaves a log line
// naming the phase (FETCH vs APPEND) and size — turning a silent stall into
// evidence of which side and which message is the culprit.
const slowMessage = 20 * time.Second

// touchReader wraps a body Read stream and pings onActivity on every non-empty
// read, so a long FETCH keeps the stall watchdog fed byte-by-byte.
type touchReader struct {
	r  io.Reader
	on func()
}

func (t touchReader) Read(p []byte) (int, error) {
	n, err := t.r.Read(p)
	if n > 0 && t.on != nil {
		t.on()
	}
	return n, err
}

// touchWriter wraps the APPEND write stream and pings onActivity on every
// non-empty write, so a long upload keeps the stall watchdog fed byte-by-byte.
type touchWriter struct {
	w  io.Writer
	on func()
}

func (t touchWriter) Write(p []byte) (int, error) {
	n, err := t.w.Write(p)
	if n > 0 && t.on != nil {
		t.on()
	}
	return n, err
}

// streamOne FETCHes BODY[] for one message and APPENDs it into dst without
// spooling to disk. The body is buffered in RAM only for the duration of
// this single FETCH->APPEND round trip. onActivity (may be nil) fires as bytes
// move in either direction, feeding the orchestrator's stall watchdog.
func streamOne(src, dst *imapclient.Client, dstFolder string, uid imap.UID, flags []imap.Flag, internalDate time.Time, onActivity func()) error {
	bodySection := &imap.FetchItemBodySection{}
	fetchStart := time.Now()
	fetchCmd := src.Fetch(imap.UIDSetNum(uid), &imap.FetchOptions{
		BodySection: []*imap.FetchItemBodySection{bodySection},
	})
	defer fetchCmd.Close()

	msg := fetchCmd.Next()
	if msg == nil {
		return fmt.Errorf("no message for uid %v", uid)
	}
	var body []byte
	for {
		item := msg.Next()
		if item == nil {
			break
		}
		if d, ok := item.(imapclient.FetchItemDataBodySection); ok {
			b, err := io.ReadAll(touchReader{r: d.Literal, on: onActivity})
			if err != nil {
				return err
			}
			body = b
		}
	}
	if err := fetchCmd.Close(); err != nil {
		return err
	}
	if body == nil {
		return fmt.Errorf("empty body uid %v", uid)
	}
	fetchDur := time.Since(fetchStart)

	appendStart := time.Now()
	appendCmd := dst.Append(dstFolder, int64(len(body)), &imap.AppendOptions{Flags: keepFlags(flags), Time: internalDate})
	// Append acquires go-imap's per-client encoder mutex and holds it until
	// Close() calls enc.end(). Close() MUST run on every path: if io.Copy
	// fails mid-write (server stall, idle timeout), returning without Close()
	// leaks the mutex and the NEXT Append on this client deadlocks forever on
	// beginCommand. Close() is idempotent and always releases the lock.
	_, copyErr := io.Copy(touchWriter{w: appendCmd, on: onActivity}, bytes.NewReader(body))
	closeErr := appendCmd.Close()
	if copyErr != nil {
		return fmt.Errorf("append body uid %v: %w", uid, copyErr)
	}
	if closeErr != nil {
		return closeErr
	}
	if _, err := appendCmd.Wait(); err != nil {
		return err
	}
	appendDur := time.Since(appendStart)

	// One message that individually eats a large slice of the stall budget is
	// the prime suspect behind a "no progress" cancel; name it, its size, and
	// which phase was slow so the culprit is visible in the logs.
	if fetchDur > slowMessage || appendDur > slowMessage {
		slog.Warn("slow message copy", "uid", uid, "bytes", len(body),
			"fetch", fetchDur.Round(time.Millisecond), "append", appendDur.Round(time.Millisecond))
	}
	return nil
}

// keepFlags drops \Recent: it cannot be set via APPEND. go-imap v2 beta.8
// no longer defines an imap.FlagRecent constant (RFC 9051 dropped \Recent
// from IMAP4rev2), so match it by its literal wire form instead.
func keepFlags(flags []imap.Flag) []imap.Flag {
	out := make([]imap.Flag, 0, len(flags))
	for _, f := range flags {
		if f == "\\Recent" {
			continue
		}
		out = append(out, f)
	}
	return out
}