package fasthttp import ( "context" "errors" "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.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 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.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 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 // 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.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) 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.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) 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 } go d.tcpAddrsClean() }) deadline := time.Now().Add(timeout) addrs, idx, err := d.getTCPAddrs(addr, dualStack, deadline) if err != nil { return nil, err } network := "tcp4" if dualStack { network = "tcp" } var conn net.Conn n := uint32(len(addrs)) for n > 0 { conn, err = d.tryDial(network, &addrs[idx%n], deadline, d.concurrencyCh) if err == nil { return conn, nil } if err == ErrDialTimeout { return nil, err } idx++ n-- } return nil, err } func (d *TCPDialer) tryDial(network string, addr *net.TCPAddr, deadline time.Time, concurrencyCh chan struct{}) (net.Conn, error) { timeout := time.Until(deadline) if timeout <= 0 { return nil, ErrDialTimeout } 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, ErrDialTimeout } } 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.String()) if err != nil && ctx.Err() == context.DeadlineExceeded { return nil, ErrDialTimeout } return conn, err } // ErrDialTimeout is returned when TCP dialing is timed out. var ErrDialTimeout = errors.New("dialing to the given TCP address timed out") // 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 interface{}) 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")