diff --git a/internal/mail/imap.go b/internal/mail/imap.go index c59da47..8c9880d 100644 --- a/internal/mail/imap.go +++ b/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 } diff --git a/internal/mail/imap_test.go b/internal/mail/imap_test.go index 70c2f97..22a47ee 100644 --- a/internal/mail/imap_test.go +++ b/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") } }