internal/mail/imap.go

@@ -39,12 +39,18 @@ type IMAPWatcher struct {
 
 const imapDialTimeout = 30 * time.Second
 
-type timeoutDialer struct {
+type timeoutContextDialer struct {
 	timeout time.Duration
 }
 
-func (d timeoutDialer) Dial(network, address string) (net.Conn, error) {
-	return (&net.Dialer{Timeout: d.timeout}).Dial(network, address)
+func (d timeoutContextDialer) Dial(network, address string) (net.Conn, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), d.timeout)
+	defer cancel()
+	return d.DialContext(ctx, network, address)
+}
+
+func (d timeoutContextDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	return (&net.Dialer{Timeout: d.timeout}).DialContext(ctx, network, address)
 }
 
 var socks5DialerFactory = proxy.SOCKS5
@@ -189,12 +195,20 @@ func dialTLSViaSOCKS5(addr, proxyURL string, options *imapclient.Options) (*imap
 		auth = &proxy.Auth{User: u.User.Username(), Password: pw}
 	}
 
-	dialer, err := socks5DialerFactory("tcp", u.Host, auth, timeoutDialer{timeout: imapDialTimeout})
+	dialer, err := socks5DialerFactory("tcp", u.Host, auth, timeoutContextDialer{timeout: imapDialTimeout})
 	if err != nil {
 		return nil, fmt.Errorf("create socks5 dialer: %w", err)
 	}
 
-	conn, err := dialer.Dial("tcp", addr)
+	dialCtx, dialCancel := context.WithTimeout(context.Background(), imapDialTimeout)
+	defer dialCancel()
+
+	var conn net.Conn
+	if cd, ok := dialer.(proxy.ContextDialer); ok {
+		conn, err = cd.DialContext(dialCtx, "tcp", addr)
+	} else {
+		conn, err = dialer.Dial("tcp", addr)
+	}
 	if err != nil {
 		return nil, fmt.Errorf("proxy dial: %w", err)
 	}
@@ -204,10 +218,14 @@ func dialTLSViaSOCKS5(addr, proxyURL string, options *imapclient.Options) (*imap
 		host = addr
 	}
 	tlsConn := tls.Client(conn, &tls.Config{ServerName: host})
-	if err := tlsConn.Handshake(); err != nil {
+	_ = tlsConn.SetDeadline(time.Now().Add(imapDialTimeout))
+	handshakeCtx, handshakeCancel := context.WithTimeout(context.Background(), imapDialTimeout)
+	defer handshakeCancel()
+	if err := tlsConn.HandshakeContext(handshakeCtx); err != nil {
 		_ = tlsConn.Close()
 		return nil, fmt.Errorf("tls handshake: %w", err)
 	}
+	_ = tlsConn.SetDeadline(time.Time{})
 
 	return imapclient.New(tlsConn, options), nil
 }

internal/mail/imap_test.go

@@ -3,6 +3,7 @@ package mail
 import (
 	"context"
 	"errors"
+	"io"
 	"net"
 	"strings"
 	"testing"
@@ -79,29 +80,83 @@ func TestDialTLSViaSOCKS5_RejectsNonSocks5Scheme(t *testing.T) {
 	}
 }
 
-type errDialer struct{}
+type fakeContextDialer struct{}
 
-func (errDialer) Dial(network, address string) (net.Conn, error) {
-	return nil, errors.New("stop")
+func (fakeContextDialer) Dial(network, address string) (net.Conn, error) {
+	return nil, errors.New("used dial")
 }
 
-func TestDialTLSViaSOCKS5_UsesTimeoutForwardDialer(t *testing.T) {
+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) {
-		td, ok := forward.(timeoutDialer)
-		if !ok {
-			t.Fatalf("expected timeoutDialer, got %T", forward)
+		if _, ok := forward.(proxy.ContextDialer); !ok {
+			t.Fatalf("expected ContextDialer, got %T", forward)
 		}
-		if td.timeout != 30*time.Second {
-			t.Fatalf("expected timeout 30s, got %s", td.timeout)
-		}
-		return errDialer{}, nil
+		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: stop") {
-		t.Fatalf("expected proxy dial stop error, got %v", err)
+	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")
 	}
 }