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/netip"
 	"syscall"
 	"time"
 )
@@ -443,6 +444,114 @@ 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 netip.AddrPort) (Conn, error) {
+	if _, ok := d.LocalAddr.(*TCPAddr); !ok && d.LocalAddr != nil {
+		return nil, &AddrError{Err: "mismatched local address type", Addr: d.LocalAddr.String()}
+	}
+	var addr *TCPAddr
+	if address.IsValid() {
+		addr = &TCPAddr{
+			IP:   address.Addr().AsSlice(),
+			Port: int(address.Port()),
+			Zone: address.Addr().Zone(),
+		}
+	}
+	switch network {
+	case "tcp", "tcp4", "tcp6":
+	default:
+		return nil, &OpError{Op: "dial", Net: network, Source: d.LocalAddr, Addr: addr, Err: UnknownNetworkError(network)}
+	}
+	c, err := d.dialAddr(ctx, network, addr)
+	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 netip.Addr) (Conn, error) {
+	if _, ok := d.LocalAddr.(*IPAddr); !ok && d.LocalAddr != nil {
+		return nil, &AddrError{Err: "mismatched local address type", Addr: d.LocalAddr.String()}
+	}
+	addr := &IPAddr{
+		IP:   address.AsSlice(),
+		Zone: address.Zone(),
+	}
+	return d.dialAddr(ctx, network, addr)
+}
+
+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 netip.AddrPort) (Conn, error) {
+	if _, ok := d.LocalAddr.(*UDPAddr); !ok && d.LocalAddr != nil {
+		return nil, &AddrError{Err: "mismatched local address type", Addr: d.LocalAddr.String()}
+	}
+	addr := &UDPAddr{
+		IP:   address.Addr().AsSlice(),
+		Zone: address.Addr().Zone(),
+		Port: int(address.Port()),
+	}
+	switch network {
+	case "udp", "udp4", "udp6":
+	default:
+		return nil, &OpError{Op: "dial", Net: network, Source: d.LocalAddr, Addr: addr, Err: UnknownNetworkError(network)}
+	}
+	return d.dialAddr(ctx, network, addr)
+}
+
 // 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

@@ -6,6 +6,7 @@ package net
 
 import (
 	"context"
+	"net/netip"
 	"syscall"
 )
 
@@ -220,6 +221,42 @@ func DialIP(network string, laddr, raddr *IPAddr) (*IPConn, error) {
 	return c, nil
 }
 
+// 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 netip.Addr) (*IPConn, error) {
+	if ctx == nil {
+		panic("nil context")
+	}
+	if !raddr.IsValid() {
+		return nil, &OpErrorNetIP{Op: "dial", Net: network, Source: netip.AddrPortFrom(laddr, 0), Err: errMissingAddress}
+	}
+	sd := &sysDialer{network: network, address: raddr.String()}
+	var src *IPAddr
+	if laddr.IsValid() {
+		src = &IPAddr{
+			IP:   laddr.AsSlice(),
+			Zone: laddr.Zone(),
+		}
+	}
+	dst := &IPAddr{
+		IP:   raddr.AsSlice(),
+		Zone: raddr.Zone(),
+	}
+	c, err := sd.dialIP(ctx, src, dst)
+	if err != nil {
+		return nil, &OpErrorNetIP{Op: "dial", Net: network, Source: netip.AddrPortFrom(laddr, 0), Addr: netip.AddrPortFrom(raddr, 0), Err: err}
+	}
+	return c, nil
+}
+
 // ListenIP acts like ListenPacket for IP networks.
 //
 // The network must be an IP network name; see func Dial for details.

src/net/net.go

@@ -83,6 +83,7 @@ import (
 	"errors"
 	"internal/poll"
 	"io"
+	"net/netip"
 	"os"
 	"sync"
 	"syscall"
@@ -488,6 +489,61 @@ func (e *OpError) Error() string {
 	return s
 }
 
+// OpErrorNetIP is the error type usually returned by functions in the net
+// package. It describes the operation, network type, and address of
+// an error.
+type OpErrorNetIP struct {
+	// Op is the operation which caused the error, such as
+	// "read" or "write".
+	Op string
+
+	// Net is the network type on which this error occurred,
+	// such as "tcp" or "udp6".
+	Net string
+
+	// For operations involving a remote network connection, like
+	// Dial, Read, or Write, Source is the corresponding local
+	// network address.
+	Source netip.AddrPort
+
+	// Addr is the network address for which this error occurred.
+	// For local operations, like Listen or SetDeadline, Addr is
+	// the address of the local endpoint being manipulated.
+	// For operations involving a remote network connection, like
+	// Dial, Read, or Write, Addr is the remote address of that
+	// connection.
+	Addr netip.AddrPort
+
+	// Err is the error that occurred during the operation.
+	// The Error method panics if the error is nil.
+	Err error
+}
+
+func (e *OpErrorNetIP) Unwrap() error { return e.Err }
+
+func (e *OpErrorNetIP) Error() string {
+	if e == nil {
+		return "<nil>"
+	}
+	s := e.Op
+	if e.Net != "" {
+		s += " " + e.Net
+	}
+	if e.Source.IsValid() {
+		s += " " + e.Source.String()
+	}
+	if e.Addr.IsValid() {
+		if e.Source.IsValid() {
+			s += "->"
+		} else {
+			s += " "
+		}
+		s += e.Addr.String()
+	}
+	s += ": " + e.Err.Error()
+	return s
+}
+
 var (
 	// aLongTimeAgo is a non-zero time, far in the past, used for
 	// immediate cancellation of dials.

src/net/tcpsock.go

@@ -247,6 +247,50 @@ func DialTCP(network string, laddr, raddr *TCPAddr) (*TCPConn, error) {
 	return c, nil
 }
 
+// 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 netip.AddrPort) (*TCPConn, error) {
+	if ctx == nil {
+		panic("nil context")
+	}
+	switch network {
+	case "tcp", "tcp4", "tcp6":
+	default:
+		return nil, &OpErrorNetIP{Op: "dial", Net: network, Source: laddr, Addr: raddr, Err: UnknownNetworkError(network)}
+	}
+	if !raddr.IsValid() {
+		return nil, &OpErrorNetIP{Op: "dial", Net: network, Source: laddr, Err: errMissingAddress}
+	}
+
+	var src *TCPAddr
+	if laddr.IsValid() {
+		src = &TCPAddr{
+			IP:   laddr.Addr().AsSlice(),
+			Zone: laddr.Addr().Zone(),
+			Port: int(laddr.Port()),
+		}
+	}
+	dst := &TCPAddr{
+		IP:   raddr.Addr().AsSlice(),
+		Zone: raddr.Addr().Zone(),
+		Port: int(raddr.Port()),
+	}
+	sd := &sysDialer{network: network, address: raddr.String()}
+	c, err := sd.dialTCP(ctx, src, dst)
+	if err != nil {
+		return nil, &OpErrorNetIP{Op: "dial", Net: network, Source: laddr, Addr: raddr, Err: err}
+	}
+	return c, nil
+}
+
 // TCPListener is a TCP network listener. Clients should typically
 // use variables of type Listener instead of assuming TCP.
 type TCPListener struct {

src/net/udpsock.go

@@ -300,6 +300,50 @@ func DialUDP(network string, laddr, raddr *UDPAddr) (*UDPConn, error) {
 	return c, nil
 }
 
+// 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 netip.AddrPort) (*UDPConn, error) {
+	if ctx == nil {
+		panic("nil context")
+	}
+	switch network {
+	case "udp", "udp4", "udp6":
+	default:
+		return nil, &OpErrorNetIP{Op: "dial", Net: network, Source: laddr, Addr: raddr, Err: UnknownNetworkError(network)}
+	}
+	if !raddr.IsValid() {
+		return nil, &OpErrorNetIP{Op: "dial", Net: network, Source: laddr, Err: errMissingAddress}
+	}
+	sd := &sysDialer{network: network, address: raddr.String()}
+
+	var src *UDPAddr
+	if laddr.IsValid() {
+		src = &UDPAddr{
+			IP:   laddr.Addr().AsSlice(),
+			Zone: laddr.Addr().Zone(),
+			Port: int(laddr.Port()),
+		}
+	}
+	dst := &UDPAddr{
+		IP:   raddr.Addr().AsSlice(),
+		Zone: raddr.Addr().Zone(),
+		Port: int(raddr.Port()),
+	}
+	c, err := sd.dialUDP(ctx, src, dst)
+	if err != nil {
+		return nil, &OpErrorNetIP{Op: "dial", Net: network, Source: laddr, Addr: raddr, Err: err}
+	}
+	return c, nil
+}
+
 // ListenUDP acts like ListenPacket for UDP networks.
 //
 // The network must be a UDP network name; see func Dial for details.

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}
 	}