497 lines
14 KiB
Go
497 lines
14 KiB
Go
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 {
|
|
// 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
|
|
|
|
// 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
|
|
|
|
// 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
|
|
|
|
// 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
|
|
// 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)) // #nosec G115
|
|
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 {
|
|
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
|
|
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()
|
|
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")
|