imap-copier/internal/imapx/dial.go

package imapx

import (
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"time"

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

type Endpoint struct {
	Host    string
	Port    int
	TLSMode string // ssl | starttls | plain
}

func (e Endpoint) addr() string { return fmt.Sprintf("%s:%d", e.Host, e.Port) }

// Client wraps an imapclient.Client together with the raw network connection it
// runs over. The embedded *imapclient.Client provides the full IMAP API; the
// retained conn lets callers impose a read deadline on the socket for the
// duration of a body transfer.
//
// This defends against servers that announce a BODY[] literal LARGER than the
// bytes they actually send (a protocol violation observed on some webmail
// servers). go-imap reads a literal strictly by its announced size, so a short
// literal makes it block forever waiting for bytes that never arrive. A
// deadline around the body read turns that infinite hang into a timeout the
// copier can recover from.
type Client struct {
	*imapclient.Client
	conn net.Conn
}

// SetReadDeadline sets (or, with the zero time, clears) a deadline on the
// underlying socket. Used to bound a single body read; always cleared again
// once the read completes so it never affects idle periods.
func (c *Client) SetReadDeadline(t time.Time) error {
	return c.conn.SetReadDeadline(t)
}

// dialTimeout bounds establishing the TCP connection.
const dialTimeout = 30 * time.Second

// dialOnce establishes one connection and returns a ready *Client. ctx bounds
// the TCP dial. We deliberately do NOT impose a socket-level read deadline:
// a blanket read deadline can't tell an idle connection (e.g. dst sitting idle
// during a long src scan) from one stuck mid-response, and would wrongly close
// idle connections. Stall detection is done at the orchestrator level via a
// progress watchdog; go-imap's own per-command timeouts bound active commands.
func dialOnce(ctx context.Context, ep Endpoint) (*Client, error) {
	d := &net.Dialer{Timeout: dialTimeout}
	raw, err := d.DialContext(ctx, "tcp", ep.addr())
	if err != nil {
		return nil, err
	}

	switch ep.TLSMode {
	case "ssl":
		// NextProtos mirrors imapclient.DialTLS's ALPN advertisement.
		tlsConn := tls.Client(raw, &tls.Config{ServerName: ep.Host, NextProtos: []string{"imap"}})
		if err := tlsConn.HandshakeContext(ctx); err != nil {
			_ = raw.Close()
			return nil, err
		}
		c, err := waitGreeting(imapclient.New(tlsConn, nil))
		if err != nil {
			return nil, err
		}
		return &Client{Client: c, conn: tlsConn}, nil
	case "starttls":
		// Deadline goes on the raw TCP conn: it sits beneath the TLS layer that
		// NewStartTLS negotiates, and a TCP read deadline still interrupts the
		// TLS read above it.
		opts := &imapclient.Options{TLSConfig: &tls.Config{ServerName: ep.Host}}
		c, err := imapclient.NewStartTLS(raw, opts)
		if err != nil {
			return nil, err
		}
		return &Client{Client: c, conn: raw}, nil
	case "plain":
		c, err := waitGreeting(imapclient.New(raw, nil))
		if err != nil {
			return nil, err
		}
		return &Client{Client: c, conn: raw}, nil
	default:
		_ = raw.Close()
		return nil, fmt.Errorf("unknown tls_mode %q", ep.TLSMode)
	}
}

// waitGreeting blocks for the server's initial greeting so a mute server is
// caught at connect time rather than at the first command. NewStartTLS already
// awaits the greeting during its STARTTLS upgrade.
func waitGreeting(c *imapclient.Client) (*imapclient.Client, error) {
	if err := c.WaitGreeting(); err != nil {
		_ = c.Close()
		return nil, err
	}
	return c, nil
}

func Connect(ctx context.Context, ep Endpoint) (*Client, error) {
	const attempts = 3
	var lastErr error
	for i := 0; i < attempts; i++ {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		c, err := dialOnce(ctx, ep)
		if err == nil {
			return c, nil
		}
		lastErr = err
		if i < attempts-1 {
			backoff := time.Duration(200*(i+1)) * time.Millisecond
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(backoff):
			}
		}
	}
	return nil, lastErr
}

func TestEndpoint(ctx context.Context, ep Endpoint) error {
	c, err := Connect(ctx, ep)
	if err != nil {
		return err
	}
	return c.Logout().Wait()
}