From 575d4ab431784bf69ade5f0b890b2dddef39e7ea Mon Sep 17 00:00:00 2001
From: Douglas De Toni Machado <douglas.machado@traefik.io>
Date: Thu, 19 May 2022 11:40:09 -0300
Subject: [PATCH] Fix initial tcp lookup when address is not available

---
 pkg/tcp/proxy.go      | 32 +++++++++++++++++---------------
 pkg/tcp/proxy_test.go | 20 ++++++++++----------
 2 files changed, 27 insertions(+), 25 deletions(-)

diff --git a/pkg/tcp/proxy.go b/pkg/tcp/proxy.go
index 06dc2884a..22e35e191 100644
--- a/pkg/tcp/proxy.go
+++ b/pkg/tcp/proxy.go
@@ -14,33 +14,31 @@ import (
 // Proxy forwards a TCP request to a TCP service.
 type Proxy struct {
 	address          string
-	target           *net.TCPAddr
+	tcpAddr          *net.TCPAddr
 	terminationDelay time.Duration
 	proxyProtocol    *dynamic.ProxyProtocol
-	refreshTarget    bool
 }
 
 // NewProxy creates a new Proxy.
 func NewProxy(address string, terminationDelay time.Duration, proxyProtocol *dynamic.ProxyProtocol) (*Proxy, error) {
-	tcpAddr, err := net.ResolveTCPAddr("tcp", address)
-	if err != nil {
-		return nil, err
-	}
-
 	if proxyProtocol != nil && (proxyProtocol.Version < 1 || proxyProtocol.Version > 2) {
 		return nil, fmt.Errorf("unknown proxyProtocol version: %d", proxyProtocol.Version)
 	}
 
-	// enable the refresh of the target only if the address in not an IP
-	refreshTarget := false
-	if host, _, err := net.SplitHostPort(address); err == nil && net.ParseIP(host) == nil {
-		refreshTarget = true
+	// Creates the tcpAddr only for IP based addresses,
+	// because there is no need to resolve the name on every new connection,
+	// and building it should happen once.
+	var tcpAddr *net.TCPAddr
+	if host, _, err := net.SplitHostPort(address); err == nil && net.ParseIP(host) != nil {
+		tcpAddr, err = net.ResolveTCPAddr("tcp", address)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	return &Proxy{
 		address:          address,
-		target:           tcpAddr,
-		refreshTarget:    refreshTarget,
+		tcpAddr:          tcpAddr,
 		terminationDelay: terminationDelay,
 		proxyProtocol:    proxyProtocol,
 	}, nil
@@ -83,10 +81,14 @@ func (p *Proxy) ServeTCP(conn WriteCloser) {
 }
 
 func (p Proxy) dialBackend() (*net.TCPConn, error) {
-	if !p.refreshTarget {
-		return net.DialTCP("tcp", nil, p.target)
+	// Dial using directly the TCPAddr for IP based addresses.
+	if p.tcpAddr != nil {
+		return net.DialTCP("tcp", nil, p.tcpAddr)
 	}
 
+	log.WithoutContext().Debugf("Dial with lookup to address %s", p.address)
+
+	// Dial with DNS lookup for host based addresses.
 	conn, err := net.Dial("tcp", p.address)
 	if err != nil {
 		return nil, err
diff --git a/pkg/tcp/proxy_test.go b/pkg/tcp/proxy_test.go
index e70b36796..fb8cad1dd 100644
--- a/pkg/tcp/proxy_test.go
+++ b/pkg/tcp/proxy_test.go
@@ -176,16 +176,20 @@ func TestLookupAddress(t *testing.T) {
 	testCases := []struct {
 		desc          string
 		address       string
-		expectRefresh bool
+		expectAddr    assert.ComparisonAssertionFunc
+		expectRefresh assert.ValueAssertionFunc
 	}{
 		{
-			desc:    "IP doesn't need refresh",
-			address: "8.8.4.4:53",
+			desc:          "IP doesn't need refresh",
+			address:       "8.8.4.4:53",
+			expectAddr:    assert.Equal,
+			expectRefresh: assert.NotNil,
 		},
 		{
 			desc:          "Hostname needs refresh",
 			address:       "dns.google:53",
-			expectRefresh: true,
+			expectAddr:    assert.NotEqual,
+			expectRefresh: assert.Nil,
 		},
 	}
 
@@ -197,16 +201,12 @@ func TestLookupAddress(t *testing.T) {
 			proxy, err := NewProxy(test.address, 10*time.Millisecond, nil)
 			require.NoError(t, err)
 
-			require.NotNil(t, proxy.target)
+			test.expectRefresh(t, proxy.tcpAddr)
 
 			conn, err := proxy.dialBackend()
 			require.NoError(t, err)
 
-			if test.expectRefresh {
-				assert.NotEqual(t, test.address, conn.RemoteAddr().String())
-			} else {
-				assert.Equal(t, test.address, conn.RemoteAddr().String())
-			}
+			test.expectAddr(t, test.address, conn.RemoteAddr().String())
 		})
 	}
 }