email_agent/internal/mailfetch/imap.go

// Package mailfetch reads mail from Gmail over IMAP and turns it into stored
// messages.
//
// The IMAP client is intentionally dependency-free: it implements only the
// subset the Email Agent needs (LOGIN, SELECT, UID SEARCH, UID FETCH) and it
// always dials IPv4. IPv6 to Gmail is broken from .13 and previously caused a
// socket EOF (AGENTS.md §9).
package mailfetch

import (
	"bufio"
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"regexp"
	"strconv"
	"strings"
	"time"
)

// DefaultTimeout bounds a single IMAP read/write.
const DefaultTimeout = 60 * time.Second

// Client is a minimal read-only IMAP client.
type Client struct {
	conn    net.Conn
	br      *bufio.Reader
	tagSeq  int
	timeout time.Duration
	preAuth bool
}

// response is one IMAP response: the assembled text (literals removed) plus any
// literal payloads in order of appearance.
type response struct {
	line     string
	literals [][]byte
}

// ServerError is a tagged NO/BAD completion from the IMAP server.
type ServerError struct {
	Command  string
	Response string
}

func (e *ServerError) Error() string {
	return fmt.Sprintf("imap %s failed: %s", e.Command, e.Response)
}

// Dial opens an implicit-TLS IMAP connection to host:port, forcing IPv4.
func Dial(ctx context.Context, host string, port int, timeout time.Duration) (*Client, error) {
	if timeout <= 0 {
		timeout = DefaultTimeout
	}
	addr := net.JoinHostPort(host, strconv.Itoa(port))
	dialer := &net.Dialer{Timeout: timeout, KeepAlive: 30 * time.Second}
	raw, err := dialer.DialContext(ctx, "tcp4", addr)
	if err != nil {
		return nil, fmt.Errorf("dial %s over IPv4: %w", addr, err)
	}

	tlsConn := tls.Client(raw, &tls.Config{
		ServerName: host,
		MinVersion: tls.VersionTLS12,
	})
	hsCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	if err := tlsConn.HandshakeContext(hsCtx); err != nil {
		_ = raw.Close()
		return nil, fmt.Errorf("tls handshake with %s: %w", addr, err)
	}

	c := &Client{conn: tlsConn, br: bufio.NewReader(tlsConn), timeout: timeout}

	greeting, err := c.readResponse()
	if err != nil {
		_ = tlsConn.Close()
		return nil, fmt.Errorf("read IMAP greeting: %w", err)
	}
	upper := strings.ToUpper(greeting.line)
	if strings.HasPrefix(upper, "* PREAUTH") {
		c.preAuth = true
		return c, nil
	}
	if !strings.HasPrefix(upper, "* OK") {
		_ = tlsConn.Close()
		return nil, fmt.Errorf("unexpected IMAP greeting: %s", greeting.line)
	}
	return c, nil
}

// Login authenticates with an app password.
func (c *Client) Login(user, pass string) error {
	if c.preAuth {
		return nil
	}
	return c.exec("LOGIN "+imapQuote(user)+" "+imapQuote(pass), nil)
}

// Select opens a mailbox and returns its UIDVALIDITY.
func (c *Client) Select(mailbox string) (uint32, error) {
	var uidValidity uint32
	err := c.exec("SELECT "+imapQuote(mailbox), func(r *response) error {
		if v, ok := parseUIDValidity(r.line); ok {
			uidValidity = v
		}
		return nil
	})
	return uidValidity, err
}

// SearchSince returns the UIDs of messages with an internal date on or after t.
func (c *Client) SearchSince(t time.Time) ([]uint32, error) {
	date := t.UTC().Format("02-Jan-2006")
	var uids []uint32
	err := c.exec("UID SEARCH SINCE "+date, func(r *response) error {
		if !strings.HasPrefix(strings.ToUpper(r.line), "* SEARCH") {
			return nil
		}
		fields := strings.Fields(r.line)
		for _, f := range fields[2:] {
			if n, err := strconv.ParseUint(f, 10, 32); err == nil {
				uids = append(uids, uint32(n))
			}
		}
		return nil
	})
	return uids, err
}

// FetchResult is one message returned by Fetch.
type FetchResult struct {
	UID           uint32
	GmailMsgID    string
	GmailThreadID string
	Raw           []byte
}

// fetchChunk bounds how many UIDs go into a single FETCH command.
const fetchChunk = 20

// Fetch retrieves the full RFC822 source for each UID, in chunks.
func (c *Client) Fetch(uids []uint32) ([]FetchResult, error) {
	var out []FetchResult
	for start := 0; start < len(uids); start += fetchChunk {
		end := start + fetchChunk
		if end > len(uids) {
			end = len(uids)
		}
		set := uidSet(uids[start:end])
		var batch []FetchResult
		err := c.exec("UID FETCH "+set+" (UID X-GM-MSGID X-GM-THRID BODY.PEEK[])", func(r *response) error {
			if fr, ok := parseFetch(r); ok {
				batch = append(batch, fr)
			}
			return nil
		})
		if err != nil {
			return out, err
		}
		out = append(out, batch...)
	}
	return out, nil
}

// Logout sends LOGOUT and closes the connection.
func (c *Client) Logout() error {
	err := c.exec("LOGOUT", nil)
	_ = c.conn.Close()
	return err
}

// Close closes the connection without a graceful logout.
func (c *Client) Close() error {
	if c == nil || c.conn == nil {
		return nil
	}
	return c.conn.Close()
}

// --- protocol plumbing ------------------------------------------------------

func (c *Client) exec(command string, onUntagged func(*response) error) error {
	c.tagSeq++
	tag := fmt.Sprintf("A%04d", c.tagSeq)
	if err := c.writeLine(tag + " " + command); err != nil {
		return err
	}
	for {
		resp, err := c.readResponse()
		if err != nil {
			return err
		}
		switch {
		case strings.HasPrefix(resp.line, tag+" "):
			status := strings.TrimSpace(strings.TrimPrefix(resp.line, tag+" "))
			if strings.HasPrefix(strings.ToUpper(status), "OK") {
				return nil
			}
			return &ServerError{Command: command, Response: status}
		case strings.HasPrefix(resp.line, "* "):
			if onUntagged != nil {
				if err := onUntagged(resp); err != nil {
					return err
				}
			}
		case strings.HasPrefix(resp.line, "+ "):
			return fmt.Errorf("imap server requested a continuation, which this client does not support: %s", resp.line)
		default:
			// Ignore unexpected non-tagged lines rather than desynchronising.
		}
	}
}

func (c *Client) writeLine(s string) error {
	if err := c.conn.SetWriteDeadline(time.Now().Add(c.timeout)); err != nil {
		return err
	}
	if _, err := c.conn.Write([]byte(s + "\r\n")); err != nil {
		return fmt.Errorf("imap write: %w", err)
	}
	return nil
}

// readResponse reads one logical IMAP response, resolving literals.
func (c *Client) readResponse() (*response, error) {
	var sb strings.Builder
	var lits [][]byte
	for {
		line, err := c.readLine()
		if err != nil {
			return nil, err
		}
		if n, ok := literalLen(line); ok {
			sb.WriteString(line[:strings.LastIndexByte(line, '{')])
			buf := make([]byte, n)
			if _, err := io.ReadFull(c.br, buf); err != nil {
				return nil, fmt.Errorf("read literal: %w", err)
			}
			lits = append(lits, buf)
			continue
		}
		sb.WriteString(line)
		return &response{line: sb.String(), literals: lits}, nil
	}
}

func (c *Client) readLine() (string, error) {
	if err := c.conn.SetReadDeadline(time.Now().Add(c.timeout)); err != nil {
		return "", err
	}
	line, err := c.br.ReadString('\n')
	if err != nil && !(err == io.EOF && line != "") {
		return "", fmt.Errorf("imap read: %w", err)
	}
	return strings.TrimRight(line, "\r\n"), nil
}

// literalLen reports the byte count of a trailing `{n}` literal marker.
func literalLen(line string) (int, bool) {
	if !strings.HasSuffix(line, "}") {
		return 0, false
	}
	i := strings.LastIndexByte(line, '{')
	if i < 0 {
		return 0, false
	}
	n, err := strconv.Atoi(line[i+1 : len(line)-1])
	if err != nil || n < 0 {
		return 0, false
	}
	return n, true
}

var (
	uidValidityRE = regexp.MustCompile(`(?i)\[UIDVALIDITY (\d+)\]`)
	fetchUIDRE    = regexp.MustCompile(`(?i)(?:^|[\s(])UID (\d+)`)
	fetchMsgIDRE  = regexp.MustCompile(`(?i)(?:^|[\s(])X-GM-MSGID (\d+)`)
	fetchThrIDRE  = regexp.MustCompile(`(?i)(?:^|[\s(])X-GM-THRID (\d+)`)
)

func parseUIDValidity(line string) (uint32, bool) {
	m := uidValidityRE.FindStringSubmatch(line)
	if m == nil {
		return 0, false
	}
	n, err := strconv.ParseUint(m[1], 10, 32)
	if err != nil {
		return 0, false
	}
	return uint32(n), true
}

func parseFetch(r *response) (FetchResult, bool) {
	if len(r.literals) == 0 || !strings.Contains(strings.ToUpper(r.line), "FETCH") {
		return FetchResult{}, false
	}
	fr := FetchResult{Raw: r.literals[0]}
	if m := fetchUIDRE.FindStringSubmatch(r.line); m != nil {
		if n, err := strconv.ParseUint(m[1], 10, 32); err == nil {
			fr.UID = uint32(n)
		}
	}
	if m := fetchMsgIDRE.FindStringSubmatch(r.line); m != nil {
		fr.GmailMsgID = m[1]
	}
	if m := fetchThrIDRE.FindStringSubmatch(r.line); m != nil {
		fr.GmailThreadID = m[1]
	}
	return fr, true
}

func uidSet(uids []uint32) string {
	parts := make([]string, len(uids))
	for i, u := range uids {
		parts[i] = strconv.FormatUint(uint64(u), 10)
	}
	return strings.Join(parts, ",")
}

// imapQuote quotes a string as an IMAP quoted-string.
func imapQuote(s string) string {
	s = strings.ReplaceAll(s, `\`, `\\`)
	s = strings.ReplaceAll(s, `"`, `\"`)
	return `"` + s + `"`
}