.

vendor/github.com/valyala/fasthttp/tcpdialer.go

@@ -0,0 +1,496 @@
+package fasthttp
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"strconv"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// Dial dials the given TCP addr using tcp4.
+//
+// This function has the following additional features comparing to net.Dial:
+//
+//   - It reduces load on DNS resolver by caching resolved TCP addressed
+//     for DNSCacheDuration.
+//   - It dials all the resolved TCP addresses in round-robin manner until
+//     connection is established. This may be useful if certain addresses
+//     are temporarily unreachable.
+//   - It returns ErrDialTimeout if connection cannot be established during
+//     DefaultDialTimeout seconds. Use DialTimeout for customizing dial timeout.
+//
+// This dialer is intended for custom code wrapping before passing
+// to Client.Dial or HostClient.Dial.
+//
+// For instance, per-host counters and/or limits may be implemented
+// by such wrappers.
+//
+// The addr passed to the function must contain port. Example addr values:
+//
+//   - foobar.baz:443
+//   - foo.bar:80
+//   - aaa.com:8080
+func Dial(addr string) (net.Conn, error) {
+	return defaultDialer.Dial(addr)
+}
+
+// DialTimeout dials the given TCP addr using tcp4 using the given timeout.
+//
+// This function has the following additional features comparing to net.Dial:
+//
+//   - It reduces load on DNS resolver by caching resolved TCP addressed
+//     for DNSCacheDuration.
+//   - It dials all the resolved TCP addresses in round-robin manner until
+//     connection is established. This may be useful if certain addresses
+//     are temporarily unreachable.
+//
+// This dialer is intended for custom code wrapping before passing
+// to Client.DialTimeout or HostClient.DialTimeout.
+//
+// For instance, per-host counters and/or limits may be implemented
+// by such wrappers.
+//
+// The addr passed to the function must contain port. Example addr values:
+//
+//   - foobar.baz:443
+//   - foo.bar:80
+//   - aaa.com:8080
+func DialTimeout(addr string, timeout time.Duration) (net.Conn, error) {
+	return defaultDialer.DialTimeout(addr, timeout)
+}
+
+// DialDualStack dials the given TCP addr using both tcp4 and tcp6.
+//
+// This function has the following additional features comparing to net.Dial:
+//
+//   - It reduces load on DNS resolver by caching resolved TCP addressed
+//     for DNSCacheDuration.
+//   - It dials all the resolved TCP addresses in round-robin manner until
+//     connection is established. This may be useful if certain addresses
+//     are temporarily unreachable.
+//   - It returns ErrDialTimeout if connection cannot be established during
+//     DefaultDialTimeout seconds. Use DialDualStackTimeout for custom dial
+//     timeout.
+//
+// This dialer is intended for custom code wrapping before passing
+// to Client.Dial or HostClient.Dial.
+//
+// For instance, per-host counters and/or limits may be implemented
+// by such wrappers.
+//
+// The addr passed to the function must contain port. Example addr values:
+//
+//   - foobar.baz:443
+//   - foo.bar:80
+//   - aaa.com:8080
+func DialDualStack(addr string) (net.Conn, error) {
+	return defaultDialer.DialDualStack(addr)
+}
+
+// DialDualStackTimeout dials the given TCP addr using both tcp4 and tcp6
+// using the given timeout.
+//
+// This function has the following additional features comparing to net.Dial:
+//
+//   - It reduces load on DNS resolver by caching resolved TCP addressed
+//     for DNSCacheDuration.
+//   - It dials all the resolved TCP addresses in round-robin manner until
+//     connection is established. This may be useful if certain addresses
+//     are temporarily unreachable.
+//
+// This dialer is intended for custom code wrapping before passing
+// to Client.DialTimeout or HostClient.DialTimeout.
+//
+// For instance, per-host counters and/or limits may be implemented
+// by such wrappers.
+//
+// The addr passed to the function must contain port. Example addr values:
+//
+//   - foobar.baz:443
+//   - foo.bar:80
+//   - aaa.com:8080
+func DialDualStackTimeout(addr string, timeout time.Duration) (net.Conn, error) {
+	return defaultDialer.DialDualStackTimeout(addr, timeout)
+}
+
+var defaultDialer = &TCPDialer{Concurrency: 1000}
+
+// Resolver represents interface of the tcp resolver.
+type Resolver interface {
+	LookupIPAddr(context.Context, string) (names []net.IPAddr, err error)
+}
+
+// TCPDialer contains options to control a group of Dial calls.
+type TCPDialer struct {
+	// Concurrency controls the maximum number of concurrent Dials
+	// that can be performed using this object.
+	// Setting this to 0 means unlimited.
+	//
+	// WARNING: This can only be changed before the first Dial.
+	// Changes made after the first Dial will not affect anything.
+	Concurrency int
+
+	// LocalAddr is the local address to use when dialing an
+	// address.
+	// If nil, a local address is automatically chosen.
+	LocalAddr *net.TCPAddr
+
+	// This may be used to override DNS resolving policy, like this:
+	// var dialer = &fasthttp.TCPDialer{
+	// 	Resolver: &net.Resolver{
+	// 		PreferGo:     true,
+	// 		StrictErrors: false,
+	// 		Dial: func (ctx context.Context, network, address string) (net.Conn, error) {
+	// 			d := net.Dialer{}
+	// 			return d.DialContext(ctx, "udp", "8.8.8.8:53")
+	// 		},
+	// 	},
+	// }
+	Resolver Resolver
+
+	// DisableDNSResolution may be used to disable DNS resolution
+	DisableDNSResolution bool
+	// DNSCacheDuration may be used to override the default DNS cache duration (DefaultDNSCacheDuration)
+	DNSCacheDuration time.Duration
+
+	tcpAddrsMap sync.Map
+
+	concurrencyCh chan struct{}
+
+	once sync.Once
+}
+
+// Dial dials the given TCP addr using tcp4.
+//
+// This function has the following additional features comparing to net.Dial:
+//
+//   - It reduces load on DNS resolver by caching resolved TCP addressed
+//     for DNSCacheDuration.
+//   - It dials all the resolved TCP addresses in round-robin manner until
+//     connection is established. This may be useful if certain addresses
+//     are temporarily unreachable.
+//   - It returns ErrDialTimeout if connection cannot be established during
+//     DefaultDialTimeout seconds. Use DialTimeout for customizing dial timeout.
+//
+// This dialer is intended for custom code wrapping before passing
+// to Client.Dial or HostClient.Dial.
+//
+// For instance, per-host counters and/or limits may be implemented
+// by such wrappers.
+//
+// The addr passed to the function must contain port. Example addr values:
+//
+//   - foobar.baz:443
+//   - foo.bar:80
+//   - aaa.com:8080
+func (d *TCPDialer) Dial(addr string) (net.Conn, error) {
+	return d.dial(addr, false, DefaultDialTimeout)
+}
+
+// DialTimeout dials the given TCP addr using tcp4 using the given timeout.
+//
+// This function has the following additional features comparing to net.Dial:
+//
+//   - It reduces load on DNS resolver by caching resolved TCP addressed
+//     for DNSCacheDuration.
+//   - It dials all the resolved TCP addresses in round-robin manner until
+//     connection is established. This may be useful if certain addresses
+//     are temporarily unreachable.
+//
+// This dialer is intended for custom code wrapping before passing
+// to Client.DialTimeout or HostClient.DialTimeout.
+//
+// For instance, per-host counters and/or limits may be implemented
+// by such wrappers.
+//
+// The addr passed to the function must contain port. Example addr values:
+//
+//   - foobar.baz:443
+//   - foo.bar:80
+//   - aaa.com:8080
+func (d *TCPDialer) DialTimeout(addr string, timeout time.Duration) (net.Conn, error) {
+	return d.dial(addr, false, timeout)
+}
+
+// DialDualStack dials the given TCP addr using both tcp4 and tcp6.
+//
+// This function has the following additional features comparing to net.Dial:
+//
+//   - It reduces load on DNS resolver by caching resolved TCP addressed
+//     for DNSCacheDuration.
+//   - It dials all the resolved TCP addresses in round-robin manner until
+//     connection is established. This may be useful if certain addresses
+//     are temporarily unreachable.
+//   - It returns ErrDialTimeout if connection cannot be established during
+//     DefaultDialTimeout seconds. Use DialDualStackTimeout for custom dial
+//     timeout.
+//
+// This dialer is intended for custom code wrapping before passing
+// to Client.Dial or HostClient.Dial.
+//
+// For instance, per-host counters and/or limits may be implemented
+// by such wrappers.
+//
+// The addr passed to the function must contain port. Example addr values:
+//
+//   - foobar.baz:443
+//   - foo.bar:80
+//   - aaa.com:8080
+func (d *TCPDialer) DialDualStack(addr string) (net.Conn, error) {
+	return d.dial(addr, true, DefaultDialTimeout)
+}
+
+// DialDualStackTimeout dials the given TCP addr using both tcp4 and tcp6
+// using the given timeout.
+//
+// This function has the following additional features comparing to net.Dial:
+//
+//   - It reduces load on DNS resolver by caching resolved TCP addressed
+//     for DNSCacheDuration.
+//   - It dials all the resolved TCP addresses in round-robin manner until
+//     connection is established. This may be useful if certain addresses
+//     are temporarily unreachable.
+//
+// This dialer is intended for custom code wrapping before passing
+// to Client.DialTimeout or HostClient.DialTimeout.
+//
+// For instance, per-host counters and/or limits may be implemented
+// by such wrappers.
+//
+// The addr passed to the function must contain port. Example addr values:
+//
+//   - foobar.baz:443
+//   - foo.bar:80
+//   - aaa.com:8080
+func (d *TCPDialer) DialDualStackTimeout(addr string, timeout time.Duration) (net.Conn, error) {
+	return d.dial(addr, true, timeout)
+}
+
+func (d *TCPDialer) dial(addr string, dualStack bool, timeout time.Duration) (net.Conn, error) {
+	d.once.Do(func() {
+		if d.Concurrency > 0 {
+			d.concurrencyCh = make(chan struct{}, d.Concurrency)
+		}
+
+		if d.DNSCacheDuration == 0 {
+			d.DNSCacheDuration = DefaultDNSCacheDuration
+		}
+
+		if !d.DisableDNSResolution {
+			go d.tcpAddrsClean()
+		}
+	})
+	deadline := time.Now().Add(timeout)
+	network := "tcp4"
+	if dualStack {
+		network = "tcp"
+	}
+	if d.DisableDNSResolution {
+		return d.tryDial(network, addr, deadline, d.concurrencyCh)
+	}
+	addrs, idx, err := d.getTCPAddrs(addr, dualStack, deadline)
+	if err != nil {
+		return nil, err
+	}
+	var conn net.Conn
+	n := uint32(len(addrs))
+	for n > 0 {
+		conn, err = d.tryDial(network, addrs[idx%n].String(), deadline, d.concurrencyCh)
+		if err == nil {
+			return conn, nil
+		}
+		if errors.Is(err, ErrDialTimeout) {
+			return nil, err
+		}
+		idx++
+		n--
+	}
+	return nil, err
+}
+
+func (d *TCPDialer) tryDial(
+	network string, addr string, deadline time.Time, concurrencyCh chan struct{},
+) (net.Conn, error) {
+	timeout := time.Until(deadline)
+	if timeout <= 0 {
+		return nil, wrapDialWithUpstream(ErrDialTimeout, addr)
+	}
+
+	if concurrencyCh != nil {
+		select {
+		case concurrencyCh <- struct{}{}:
+		default:
+			tc := AcquireTimer(timeout)
+			isTimeout := false
+			select {
+			case concurrencyCh <- struct{}{}:
+			case <-tc.C:
+				isTimeout = true
+			}
+			ReleaseTimer(tc)
+			if isTimeout {
+				return nil, wrapDialWithUpstream(ErrDialTimeout, addr)
+			}
+		}
+		defer func() { <-concurrencyCh }()
+	}
+
+	dialer := net.Dialer{}
+	if d.LocalAddr != nil {
+		dialer.LocalAddr = d.LocalAddr
+	}
+
+	ctx, cancelCtx := context.WithDeadline(context.Background(), deadline)
+	defer cancelCtx()
+	conn, err := dialer.DialContext(ctx, network, addr)
+	if err != nil {
+		if ctx.Err() == context.DeadlineExceeded {
+			return nil, wrapDialWithUpstream(ErrDialTimeout, addr)
+		}
+		return nil, wrapDialWithUpstream(err, addr)
+	}
+	return conn, nil
+}
+
+// ErrDialTimeout is returned when TCP dialing is timed out.
+var ErrDialTimeout = errors.New("dialing to the given TCP address timed out")
+
+// ErrDialWithUpstream wraps dial error with upstream info.
+//
+// Should use errors.As to get upstream information from error:
+//
+//	hc := fasthttp.HostClient{Addr: "foo.com,bar.com"}
+//	err := hc.Do(req, res)
+//
+//	var dialErr *fasthttp.ErrDialWithUpstream
+//	if errors.As(err, &dialErr) {
+//		upstream = dialErr.Upstream // 34.206.39.153:80
+//	}
+type ErrDialWithUpstream struct {
+	Upstream string
+	wrapErr  error
+}
+
+func (e *ErrDialWithUpstream) Error() string {
+	return fmt.Sprintf("error when dialing %s: %s", e.Upstream, e.wrapErr.Error())
+}
+
+func (e *ErrDialWithUpstream) Unwrap() error {
+	return e.wrapErr
+}
+
+func wrapDialWithUpstream(err error, upstream string) error {
+	return &ErrDialWithUpstream{
+		Upstream: upstream,
+		wrapErr:  err,
+	}
+}
+
+// DefaultDialTimeout is timeout used by Dial and DialDualStack
+// for establishing TCP connections.
+const DefaultDialTimeout = 3 * time.Second
+
+type tcpAddrEntry struct {
+	addrs    []net.TCPAddr
+	addrsIdx uint32
+
+	pending     int32
+	resolveTime time.Time
+}
+
+// DefaultDNSCacheDuration is the duration for caching resolved TCP addresses
+// by Dial* functions.
+const DefaultDNSCacheDuration = time.Minute
+
+func (d *TCPDialer) tcpAddrsClean() {
+	expireDuration := 2 * d.DNSCacheDuration
+	for {
+		time.Sleep(time.Second)
+		t := time.Now()
+		d.tcpAddrsMap.Range(func(k, v any) bool {
+			if e, ok := v.(*tcpAddrEntry); ok && t.Sub(e.resolveTime) > expireDuration {
+				d.tcpAddrsMap.Delete(k)
+			}
+			return true
+		})
+	}
+}
+
+func (d *TCPDialer) getTCPAddrs(addr string, dualStack bool, deadline time.Time) ([]net.TCPAddr, uint32, error) {
+	item, exist := d.tcpAddrsMap.Load(addr)
+	e, ok := item.(*tcpAddrEntry)
+	if exist && ok && e != nil && time.Since(e.resolveTime) > d.DNSCacheDuration {
+		// Only let one goroutine re-resolve at a time.
+		if atomic.SwapInt32(&e.pending, 1) == 0 {
+			e = nil
+		}
+	}
+
+	if e == nil {
+		addrs, err := resolveTCPAddrs(addr, dualStack, d.Resolver, deadline)
+		if err != nil {
+			item, exist := d.tcpAddrsMap.Load(addr)
+			e, ok = item.(*tcpAddrEntry)
+			if exist && ok && e != nil {
+				// Set pending to 0 so another goroutine can retry.
+				atomic.StoreInt32(&e.pending, 0)
+			}
+			return nil, 0, err
+		}
+
+		e = &tcpAddrEntry{
+			addrs:       addrs,
+			resolveTime: time.Now(),
+		}
+		d.tcpAddrsMap.Store(addr, e)
+	}
+
+	idx := atomic.AddUint32(&e.addrsIdx, 1)
+	return e.addrs, idx, nil
+}
+
+func resolveTCPAddrs(addr string, dualStack bool, resolver Resolver, deadline time.Time) ([]net.TCPAddr, error) {
+	host, portS, err := net.SplitHostPort(addr)
+	if err != nil {
+		return nil, err
+	}
+	port, err := strconv.Atoi(portS)
+	if err != nil {
+		return nil, err
+	}
+
+	if resolver == nil {
+		resolver = net.DefaultResolver
+	}
+
+	ctx, cancel := context.WithDeadline(context.Background(), deadline)
+	defer cancel()
+	ipaddrs, err := resolver.LookupIPAddr(ctx, host)
+	if err != nil {
+		return nil, err
+	}
+
+	n := len(ipaddrs)
+	addrs := make([]net.TCPAddr, 0, n)
+	for i := 0; i < n; i++ {
+		ip := ipaddrs[i]
+		if !dualStack && ip.IP.To4() == nil {
+			continue
+		}
+		addrs = append(addrs, net.TCPAddr{
+			IP:   ip.IP,
+			Port: port,
+			Zone: ip.Zone,
+		})
+	}
+	if len(addrs) == 0 {
+		return nil, errNoDNSEntries
+	}
+	return addrs, nil
+}
+
+var errNoDNSEntries = errors.New("couldn't find DNS entries for the given domain. Try using DialDualStack")