well-goknown/vendor/github.com/valyala/fasthttp/tcpdialer.go

498 lines
14 KiB
Go
Raw Permalink Normal View History

package fasthttp
import (
"context"
"errors"
2024-09-26 01:59:44 +00:00
"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
2024-09-26 01:59:44 +00:00
// 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
2024-09-26 01:59:44 +00:00
// 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 {
// 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
2024-09-26 01:59:44 +00:00
// LocalAddr is the local address to use when dialing an
// address.
// If nil, a local address is automatically chosen.
LocalAddr *net.TCPAddr
concurrencyCh chan struct{}
tcpAddrsMap sync.Map
2024-09-26 01:59:44 +00:00
// 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
// DNSCacheDuration may be used to override the default DNS cache duration (DefaultDNSCacheDuration)
DNSCacheDuration time.Duration
once sync.Once
2024-09-26 01:59:44 +00:00
// DisableDNSResolution may be used to disable DNS resolution
DisableDNSResolution bool
}
// 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
2024-09-26 01:59:44 +00:00
// 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
2024-09-26 01:59:44 +00:00
// 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
}
2024-09-26 01:59:44 +00:00
if !d.DisableDNSResolution {
go d.tcpAddrsClean()
}
})
deadline := time.Now().Add(timeout)
network := "tcp4"
if dualStack {
network = "tcp"
}
2024-09-26 01:59:44 +00:00
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
2024-09-26 01:59:44 +00:00
n := uint32(len(addrs)) // #nosec G115
for n > 0 {
2024-09-26 01:59:44 +00:00
conn, err = d.tryDial(network, addrs[idx%n].String(), deadline, d.concurrencyCh)
if err == nil {
return conn, nil
}
2024-09-26 01:59:44 +00:00
if errors.Is(err, ErrDialTimeout) {
return nil, err
}
idx++
n--
}
return nil, err
}
2024-09-26 01:59:44 +00:00
func (d *TCPDialer) tryDial(
network string, addr string, deadline time.Time, concurrencyCh chan struct{},
) (net.Conn, error) {
timeout := time.Until(deadline)
if timeout <= 0 {
2024-09-26 01:59:44 +00:00
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 {
2024-09-26 01:59:44 +00:00
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()
2024-09-26 01:59:44 +00:00
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)
}
2024-09-26 01:59:44 +00:00
return conn, nil
}
// ErrDialTimeout is returned when TCP dialing is timed out.
var ErrDialTimeout = errors.New("dialing to the given TCP address timed out")
2024-09-26 01:59:44 +00:00
// 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 {
wrapErr error
Upstream string
}
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 {
resolveTime time.Time
2024-09-26 01:59:44 +00:00
addrs []net.TCPAddr
addrsIdx uint32
pending int32
}
// 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()
2024-09-26 01:59:44 +00:00
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")