imap-copier/internal/imapx/underflow_test.go

package imapx

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync"
	"testing"
	"time"

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

// underflowServer is a minimal IMAP server that reproduces the amega.kz bug: it
// answers a BODY[] FETCH by announcing a literal LARGER than the bytes it then
// sends, and afterwards goes silent — exactly what makes go-imap's strict
// literal reader block forever. LOGIN/EXAMINE and the Pass-1 metadata FETCH are
// answered normally so a full CopyFolder can reach the poisoned message.
//
// It serves one connection per accept and keeps accepting, so a reconnect gets
// a fresh, well-behaved session (its second EXAMINE reports zero messages, so
// the resumed folder simply finishes).
type underflowServer struct {
	ln      net.Listener
	mu      sync.Mutex
	accepts int // how many connections have been accepted
	stop    chan struct{}
}

func newUnderflowServer(t *testing.T) *underflowServer {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	s := &underflowServer{ln: ln, stop: make(chan struct{})}
	go s.serve()
	return s
}

func (s *underflowServer) addr() Endpoint {
	a := s.ln.Addr().(*net.TCPAddr)
	return Endpoint{Host: "127.0.0.1", Port: a.Port, TLSMode: "plain"}
}

func (s *underflowServer) close() { close(s.stop); _ = s.ln.Close() }

func (s *underflowServer) serve() {
	for {
		conn, err := s.ln.Accept()
		if err != nil {
			return
		}
		s.mu.Lock()
		s.accepts++
		first := s.accepts == 1
		s.mu.Unlock()
		go s.handle(conn, first)
	}
}

// handle drives one connection. On the FIRST connection the mailbox reports one
// message and its BODY[] fetch under-delivers; on any later connection (i.e.
// after a reconnect) the mailbox is empty so the resumed folder completes.
func (s *underflowServer) handle(conn net.Conn, first bool) {
	defer func() { _ = conn.Close() }()
	br := bufio.NewReader(conn)
	fmt.Fprint(conn, "* OK IMAP4rev1 ready\r\n")
	for {
		line, err := br.ReadString('\n')
		if err != nil {
			return
		}
		fields := strings.Fields(line)
		if len(fields) == 0 {
			continue
		}
		tag := fields[0]
		up := strings.ToUpper(line)
		switch {
		case strings.Contains(up, "LOGIN"):
			fmt.Fprintf(conn, "%s OK LOGIN completed\r\n", tag)
		case strings.Contains(up, "EXAMINE"), strings.Contains(up, "SELECT"):
			n := 0
			if first {
				n = 1
			}
			fmt.Fprintf(conn, "* %d EXISTS\r\n", n)
			fmt.Fprint(conn, "* OK [UIDVALIDITY 1] ok\r\n")
			fmt.Fprintf(conn, "%s OK [READ-ONLY] EXAMINE completed\r\n", tag)
		case strings.Contains(up, "BODY["), strings.Contains(up, "BODY.PEEK"):
			// Poison: announce 100000 bytes, send 10, then stall until shutdown.
			fmt.Fprint(conn, "* 1 FETCH (UID 1 BODY[] {100000}\r\n")
			fmt.Fprint(conn, "0123456789")
			<-s.stop
			return
		case strings.Contains(up, "FETCH"):
			// Pass-1 metadata fetch for the single message.
			fmt.Fprint(conn, "* 1 FETCH (UID 1 RFC822.SIZE 100 FLAGS () "+
				"INTERNALDATE \"01-Jan-2020 00:00:00 +0000\" "+
				"ENVELOPE (\"Wed, 01 Jan 2020 00:00:00 +0000\" \"poison\" NIL NIL NIL NIL NIL NIL NIL \"<poison@x>\"))\r\n")
			fmt.Fprintf(conn, "%s OK FETCH completed\r\n", tag)
		case strings.Contains(up, "LOGOUT"):
			fmt.Fprintf(conn, "* BYE\r\n%s OK LOGOUT completed\r\n", tag)
			return
		case strings.Contains(up, "CREATE"), strings.Contains(up, "NOOP"):
			fmt.Fprintf(conn, "%s OK completed\r\n", tag)
		default:
			fmt.Fprintf(conn, "%s OK completed\r\n", tag)
		}
	}
}

// TestStreamOneTimesOutOnLiteralUnderflow proves a server that announces a
// larger BODY[] literal than it sends makes streamOne return ErrBodyTimeout
// (bounded by bodyIdleTimeout) instead of hanging forever.
func TestStreamOneTimesOutOnLiteralUnderflow(t *testing.T) {
	restore := shortenBodyIdle(200 * time.Millisecond)
	defer restore()

	srv := newUnderflowServer(t)
	defer srv.close()
	ctx := context.Background()

	src, err := Connect(ctx, srv.addr())
	if err != nil {
		t.Fatalf("connect: %v", err)
	}
	defer func() { _ = src.Close() }()
	if err := src.Login("u", "p").Wait(); err != nil {
		t.Fatalf("login: %v", err)
	}

	done := make(chan error, 1)
	go func() {
		done <- streamOne(src, src, "INBOX", imap.UID(1), nil, time.Time{}, nil)
	}()

	select {
	case err := <-done:
		if !errors.Is(err, ErrBodyTimeout) {
			t.Fatalf("want ErrBodyTimeout, got %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("streamOne hung on literal underflow instead of timing out")
	}
}

// TestCopyFolderSkipsAndReconnectsOnUnderflow proves CopyFolder does not hang on
// a poisoned message: it marks the message migrated (so future runs skip it),
// invokes ReconnectSrc, and returns with the error counted — the folder is not
// wedged.
func TestCopyFolderSkipsAndReconnectsOnUnderflow(t *testing.T) {
	restore := shortenBodyIdle(200 * time.Millisecond)
	defer restore()

	srv := newUnderflowServer(t)
	defer srv.close()
	ctx := context.Background()

	src, err := Connect(ctx, srv.addr())
	if err != nil {
		t.Fatalf("connect: %v", err)
	}
	defer func() { _ = src.Close() }()
	if err := src.Login("u", "p").Wait(); err != nil {
		t.Fatalf("login: %v", err)
	}

	var marked []string
	var reconnected bool
	deps := CopyDeps{
		IsMigrated:   func(string) (bool, error) { return false, nil },
		MarkMigrated: func(_, k string) error { marked = append(marked, k); return nil },
		OnProgress:   func(_, _ int) {},
		ReconnectSrc: func() (*Client, error) {
			reconnected = true
			nc, derr := Connect(ctx, srv.addr())
			if derr != nil {
				return nil, derr
			}
			if lerr := nc.Login("u", "p").Wait(); lerr != nil {
				return nil, lerr
			}
			return nc, nil
		},
	}

	done := make(chan CopyResult, 1)
	go func() {
		r, _ := CopyFolder(ctx, src, src, "INBOX", "INBOX", deps)
		done <- r
	}()

	select {
	case r := <-done:
		if r.Errors == 0 {
			t.Fatalf("expected the poisoned message counted as an error, got %+v", r)
		}
		if !reconnected {
			t.Fatal("ReconnectSrc was never called")
		}
		if len(marked) != 1 {
			t.Fatalf("poisoned message should be marked migrated once, got %v", marked)
		}
	case <-time.After(8 * time.Second):
		t.Fatal("CopyFolder hung on a poisoned message instead of skipping it")
	}
}

func shortenBodyIdle(d time.Duration) func() {
	old := bodyIdleTimeout
	bodyIdleTimeout = d
	return func() { bodyIdleTimeout = old }
}