net: add context-aware DialTCPContext, DialUDPContext, etc

Fixes #49097

Author: moredure

src/net/dial.go

@@ -7,6 +7,7 @@ package net
 import (
 	"context"
 	"internal/nettrace"
+	"net"
 	"syscall"
 	"time"
 )
@@ -443,6 +444,97 @@ func (d *Dialer) DialContext(ctx context.Context, network, address string) (Conn
 	return c, nil
 }
 
+func (d *Dialer) dialAddr(ctx context.Context, network string, address Addr) (Conn, error) {
+	if ctx == nil {
+		panic("nil context")
+	}
+	if address == nil {
+		return nil, &OpError{Op: "dial", Net: network, Source: d.LocalAddr, Addr: nil, Err: errMissingAddress}
+	}
+	deadline := d.deadline(ctx, time.Now())
+	if !deadline.IsZero() {
+		if d, ok := ctx.Deadline(); !ok || deadline.Before(d) {
+			subCtx, cancel := context.WithDeadline(ctx, deadline)
+			defer cancel()
+			ctx = subCtx
+		}
+	}
+	if oldCancel := d.Cancel; oldCancel != nil {
+		subCtx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		go func() {
+			select {
+			case <-oldCancel:
+				cancel()
+			case <-subCtx.Done():
+			}
+		}()
+		ctx = subCtx
+	}
+	sd := &sysDialer{
+		Dialer:  *d,
+		network: network,
+		address: address.String(),
+	}
+	return sd.dialSingle(ctx, address)
+}
+
+func (d *Dialer) DialTCP(ctx context.Context, network string, address *TCPAddr) (Conn, error) {
+	if _, ok := d.LocalAddr.(*TCPAddr); !ok && d.LocalAddr != nil {
+		return nil, &AddrError{Err: "mismatched local address type", Addr: d.LocalAddr.String()}
+	}
+	switch network {
+	case "tcp", "tcp4", "tcp6":
+	default:
+		return nil, &OpError{Op: "dial", Net: network, Source: d.LocalAddr, Addr: address, Err: UnknownNetworkError(network)}
+	}
+	c, err := d.dialAddr(ctx, network, address)
+	if err != nil {
+		return nil, err
+	}
+	if tc := c.(*TCPConn); d.KeepAlive >= 0 {
+		setKeepAlive(tc.fd, true)
+		ka := d.KeepAlive
+		if d.KeepAlive == 0 {
+			ka = defaultTCPKeepAlive
+		}
+		setKeepAlivePeriod(tc.fd, ka)
+		testHookSetKeepAlive(ka)
+	}
+	return c, nil
+}
+
+func (d *Dialer) DialIP(ctx context.Context, network string, address *IPAddr) (Conn, error) {
+	if _, ok := d.LocalAddr.(*IPAddr); !ok && d.LocalAddr != nil {
+		return nil, &AddrError{Err: "mismatched local address type", Addr: d.LocalAddr.String()}
+	}
+	return d.dialAddr(ctx, network, address)
+}
+
+func (d *Dialer) DialUnix(ctx context.Context, network string, address *UnixAddr) (Conn, error) {
+	if _, ok := d.LocalAddr.(*UnixAddr); !ok && d.LocalAddr != nil {
+		return nil, &AddrError{Err: "mismatched local address type", Addr: d.LocalAddr.String()}
+	}
+	switch network {
+	case "unix", "unixgram", "unixpacket":
+	default:
+		return nil, &OpError{Op: "dial", Net: network, Source: d.LocalAddr, Addr: address, Err: UnknownNetworkError(network)}
+	}
+	return d.dialAddr(ctx, network, address)
+}
+
+func (d *Dialer) DialUDP(ctx context.Context, network string, address *UDPAddr) (Conn, error) {
+	if _, ok := d.LocalAddr.(*UDPAddr); !ok && d.LocalAddr != nil {
+		return nil, &AddrError{Err: "mismatched local address type", Addr: d.LocalAddr.String()}
+	}
+	switch network {
+	case "udp", "udp4", "udp6":
+	default:
+		return nil, &OpError{Op: "dial", Net: network, Source: d.LocalAddr, Addr: address, Err: UnknownNetworkError(network)}
+	}
+	return d.dialAddr(ctx, network, address)
+}
+
 // dialParallel races two copies of dialSerial, giving the first a
 // head start. It returns the first established connection and
 // closes the others. Otherwise it returns an error from the first

src/net/iprawsock.go

@@ -209,11 +209,28 @@ func newIPConn(fd *netFD) *IPConn { return &IPConn{conn{fd}} }
 // If the IP field of raddr is nil or an unspecified IP address, the
 // local system is assumed.
 func DialIP(network string, laddr, raddr *IPAddr) (*IPConn, error) {
+	return DialIPContext(context.Background(), network, laddr, raddr)
+}
+
+// DialIPContext acts like DialIP but connects using
+// the provided context.
+//
+// The provided Context must be non-nil.
+//
+// The network must be an IP network name; see func Dial for details.
+//
+// If laddr is nil, a local address is automatically chosen.
+// If the IP field of raddr is nil or an unspecified IP address, the
+// local system is assumed.
+func DialIPContext(ctx context.Context, network string, laddr, raddr *IPAddr) (*IPConn, error) {
+	if ctx == nil {
+		panic("nil context")
+	}
 	if raddr == nil {
 		return nil, &OpError{Op: "dial", Net: network, Source: laddr.opAddr(), Addr: nil, Err: errMissingAddress}
 	}
 	sd := &sysDialer{network: network, address: raddr.String()}
-	c, err := sd.dialIP(context.Background(), laddr, raddr)
+	c, err := sd.dialIP(ctx, laddr, raddr)
 	if err != nil {
 		return nil, &OpError{Op: "dial", Net: network, Source: laddr.opAddr(), Addr: raddr.opAddr(), Err: err}
 	}

src/net/tcpsock.go

@@ -231,6 +231,23 @@ func newTCPConn(fd *netFD) *TCPConn {
 // If the IP field of raddr is nil or an unspecified IP address, the
 // local system is assumed.
 func DialTCP(network string, laddr, raddr *TCPAddr) (*TCPConn, error) {
+	return DialTCPContext(context.Background(), network, laddr, raddr)
+}
+
+// DialTCPContext acts like DialTCP but connects using
+// the provided context.
+//
+// The provided Context must be non-nil.
+//
+// The network must be a TCP network name; see func Dial for details.
+//
+// If laddr is nil, a local address is automatically chosen.
+// If the IP field of raddr is nil or an unspecified IP address, the
+// local system is assumed.
+func DialTCPContext(ctx context.Context, network string, laddr, raddr *TCPAddr) (*TCPConn, error) {
+	if ctx == nil {
+		panic("nil context")
+	}
 	switch network {
 	case "tcp", "tcp4", "tcp6":
 	default:
@@ -240,7 +257,7 @@ func DialTCP(network string, laddr, raddr *TCPAddr) (*TCPConn, error) {
 		return nil, &OpError{Op: "dial", Net: network, Source: laddr.opAddr(), Addr: nil, Err: errMissingAddress}
 	}
 	sd := &sysDialer{network: network, address: raddr.String()}
-	c, err := sd.dialTCP(context.Background(), laddr, raddr)
+	c, err := sd.dialTCP(ctx, laddr, raddr)
 	if err != nil {
 		return nil, &OpError{Op: "dial", Net: network, Source: laddr.opAddr(), Addr: raddr.opAddr(), Err: err}
 	}

src/net/udpsock.go

@@ -284,6 +284,23 @@ func newUDPConn(fd *netFD) *UDPConn { return &UDPConn{conn{fd}} }
 // If the IP field of raddr is nil or an unspecified IP address, the
 // local system is assumed.
 func DialUDP(network string, laddr, raddr *UDPAddr) (*UDPConn, error) {
+	return DialUDPContext(context.Background(), network, laddr, raddr)
+}
+
+// DialUDPContext acts like DialUDP but connects using
+// the provided context.
+//
+// The provided Context must be non-nil.
+//
+// The network must be a UDP network name; see func Dial for details.
+//
+// If laddr is nil, a local address is automatically chosen.
+// If the IP field of raddr is nil or an unspecified IP address, the
+// local system is assumed.
+func DialUDPContext(ctx context.Context, network string, laddr, raddr *UDPAddr) (*UDPConn, error) {
+	if ctx == nil {
+		panic("nil context")
+	}
 	switch network {
 	case "udp", "udp4", "udp6":
 	default:
@@ -293,7 +310,7 @@ func DialUDP(network string, laddr, raddr *UDPAddr) (*UDPConn, error) {
 		return nil, &OpError{Op: "dial", Net: network, Source: laddr.opAddr(), Addr: nil, Err: errMissingAddress}
 	}
 	sd := &sysDialer{network: network, address: raddr.String()}
-	c, err := sd.dialUDP(context.Background(), laddr, raddr)
+	c, err := sd.dialUDP(ctx, laddr, raddr)
 	if err != nil {
 		return nil, &OpError{Op: "dial", Net: network, Source: laddr.opAddr(), Addr: raddr.opAddr(), Err: err}
 	}

src/net/unixsock.go

@@ -201,13 +201,29 @@ func newUnixConn(fd *netFD) *UnixConn { return &UnixConn{conn{fd}} }
 // If laddr is non-nil, it is used as the local address for the
 // connection.
 func DialUnix(network string, laddr, raddr *UnixAddr) (*UnixConn, error) {
+	return DialUnixContext(context.Background(), network, laddr, raddr)
+}
+
+// DialUnixContext acts like DialUnix but connects using
+// the provided context.
+//
+// The provided Context must be non-nil.
+//
+// The network must be a Unix network name; see func Dial for details.
+//
+// If laddr is non-nil, it is used as the local address for the
+// connection.
+func DialUnixContext(ctx context.Context, network string, laddr, raddr *UnixAddr) (*UnixConn, error) {
+	if ctx == nil {
+		panic("nil context")
+	}
 	switch network {
 	case "unix", "unixgram", "unixpacket":
 	default:
 		return nil, &OpError{Op: "dial", Net: network, Source: laddr.opAddr(), Addr: raddr.opAddr(), Err: UnknownNetworkError(network)}
 	}
 	sd := &sysDialer{network: network, address: raddr.String()}
-	c, err := sd.dialUnix(context.Background(), laddr, raddr)
+	c, err := sd.dialUnix(ctx, laddr, raddr)
 	if err != nil {
 		return nil, &OpError{Op: "dial", Net: network, Source: laddr.opAddr(), Addr: raddr.opAddr(), Err: err}
 	}