claw-email/internal/mail/imap_test.go

package mail

import (
	"context"
	"errors"
	"io"
	"net"
	"strings"
	"testing"
	"time"

	"github.com/emersion/go-imap/v2/imapclient"
	"golang.org/x/net/proxy"
	"thuanle.me/claw-email-bridge/internal/config"
)

func TestConnectAndWatch_UsesDirectDial_WhenProxyUnset(t *testing.T) {
	watcher := &IMAPWatcher{
		cfg: &config.Config{
			IMAPHost: "imap.example.com",
			IMAPPort: "993",
			IMAPUser: "u",
			IMAPPass: "p",
		},
	}

	var directCalled bool
	watcher.dialIMAP = func(addr string, _ *imapclient.Options) (*imapclient.Client, error) {
		directCalled = true
		return nil, errors.New("stop")
	}
	watcher.dialIMAPViaProxy = func(addr, proxyURL string, _ *imapclient.Options) (*imapclient.Client, error) {
		t.Fatalf("did not expect proxy dial, got addr=%s proxy=%s", addr, proxyURL)
		return nil, nil
	}

	err := watcher.connectAndWatch(context.Background())
	if err == nil || !strings.Contains(err.Error(), "stop") {
		t.Fatalf("expected stop error, got %v", err)
	}
	if !directCalled {
		t.Fatal("expected direct dial path")
	}
}

func TestConnectAndWatch_UsesProxyDial_WhenProxySet(t *testing.T) {
	watcher := &IMAPWatcher{
		cfg: &config.Config{
			IMAPHost:     "imap.example.com",
			IMAPPort:     "993",
			IMAPUser:     "u",
			IMAPPass:     "p",
			IMAPProxyURL: "socks5://127.0.0.1:1080",
		},
	}

	var proxyCalled bool
	watcher.dialIMAP = func(addr string, _ *imapclient.Options) (*imapclient.Client, error) {
		t.Fatalf("did not expect direct dial, got addr=%s", addr)
		return nil, nil
	}
	watcher.dialIMAPViaProxy = func(addr, proxyURL string, _ *imapclient.Options) (*imapclient.Client, error) {
		proxyCalled = true
		return nil, errors.New("stop")
	}

	err := watcher.connectAndWatch(context.Background())
	if err == nil || !strings.Contains(err.Error(), "stop") {
		t.Fatalf("expected stop error, got %v", err)
	}
	if !proxyCalled {
		t.Fatal("expected proxy dial path")
	}
}

func TestDialTLSViaSOCKS5_RejectsNonSocks5Scheme(t *testing.T) {
	_, err := dialTLSViaSOCKS5("imap.example.com:993", "http://proxy:8080", nil)
	if err == nil || !strings.Contains(err.Error(), "unsupported proxy scheme") {
		t.Fatalf("expected unsupported scheme error, got %v", err)
	}
}

type fakeContextDialer struct{}

func (fakeContextDialer) Dial(network, address string) (net.Conn, error) {
	return nil, errors.New("used dial")
}

func (fakeContextDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	if _, ok := ctx.Deadline(); !ok {
		return nil, errors.New("missing deadline")
	}
	return nil, errors.New("used dialcontext")
}

func TestDialTLSViaSOCKS5_UsesDialContextWithTimeout(t *testing.T) {
	orig := socks5DialerFactory
	t.Cleanup(func() { socks5DialerFactory = orig })

	socks5DialerFactory = func(network, address string, auth *proxy.Auth, forward proxy.Dialer) (proxy.Dialer, error) {
		if _, ok := forward.(proxy.ContextDialer); !ok {
			t.Fatalf("expected ContextDialer, got %T", forward)
		}
		return fakeContextDialer{}, nil
	}

	_, err := dialTLSViaSOCKS5("imap.example.com:993", "socks5://127.0.0.1:1080", nil)
	if err == nil || !strings.Contains(err.Error(), "proxy dial: used dialcontext") {
		t.Fatalf("expected DialContext path, got %v", err)
	}
}

type trackingConn struct {
	deadline time.Time
}

func (c *trackingConn) Read(b []byte) (int, error)  { return 0, io.EOF }
func (c *trackingConn) Write(b []byte) (int, error) { return 0, io.EOF }
func (c *trackingConn) Close() error                { return nil }
func (c *trackingConn) LocalAddr() net.Addr         { return dummyAddr("local") }
func (c *trackingConn) RemoteAddr() net.Addr        { return dummyAddr("remote") }
func (c *trackingConn) SetDeadline(t time.Time) error {
	c.deadline = t
	return nil
}
func (c *trackingConn) SetReadDeadline(t time.Time) error  { return nil }
func (c *trackingConn) SetWriteDeadline(t time.Time) error { return nil }

type dummyAddr string

func (a dummyAddr) Network() string { return "tcp" }
func (a dummyAddr) String() string  { return string(a) }

type contextConnDialer struct {
	conn *trackingConn
}

func (d contextConnDialer) Dial(network, address string) (net.Conn, error) {
	return nil, errors.New("used dial")
}

func (d contextConnDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	return d.conn, nil
}

func TestDialTLSViaSOCKS5_SetsDeadlineForTLSHandshake(t *testing.T) {
	orig := socks5DialerFactory
	t.Cleanup(func() { socks5DialerFactory = orig })

	conn := &trackingConn{}
	socks5DialerFactory = func(network, address string, auth *proxy.Auth, forward proxy.Dialer) (proxy.Dialer, error) {
		return contextConnDialer{conn: conn}, nil
	}

	_, err := dialTLSViaSOCKS5("imap.example.com:993", "socks5://127.0.0.1:1080", nil)
	if err == nil || !strings.Contains(err.Error(), "tls handshake") {
		t.Fatalf("expected tls handshake error, got %v", err)
	}
	if conn.deadline.IsZero() {
		t.Fatal("expected TLS handshake deadline to be set")
	}
}