Improve service name lookup on TCP routers

pkg/tcp/proxy_test.go

@@ -5,9 +5,11 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"sync"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -79,3 +81,74 @@ func TestCloseWrite(t *testing.T) {
 	require.Equal(t, int64(4), n)
 	require.Equal(t, "PONG", buffer.String())
 }
+
+func TestLookupAddress(t *testing.T) {
+	testCases := []struct {
+		desc       string
+		address    string
+		expectSame assert.ComparisonAssertionFunc
+	}{
+		{
+			desc:       "IP doesn't need refresh",
+			address:    "8.8.4.4:53",
+			expectSame: assert.Same,
+		},
+		{
+			desc:       "Hostname needs refresh",
+			address:    "dns.google:53",
+			expectSame: assert.NotSame,
+		},
+	}
+
+	for _, test := range testCases {
+		test := test
+		t.Run(test.desc, func(t *testing.T) {
+			t.Parallel()
+
+			proxy, err := NewProxy(test.address, 10*time.Millisecond)
+			require.NoError(t, err)
+
+			require.NotNil(t, proxy.target)
+
+			proxyListener, err := net.Listen("tcp", ":0")
+			require.NoError(t, err)
+
+			var wg sync.WaitGroup
+			go func(wg *sync.WaitGroup) {
+				for {
+					conn, err := proxyListener.Accept()
+					require.NoError(t, err)
+
+					proxy.ServeTCP(conn.(*net.TCPConn))
+
+					wg.Done()
+				}
+			}(&wg)
+
+			var lastTarget *net.TCPAddr
+
+			for i := 0; i < 3; i++ {
+				wg.Add(1)
+
+				conn, err := net.Dial("tcp", proxyListener.Addr().String())
+				require.NoError(t, err)
+
+				_, err = conn.Write([]byte("ping\n"))
+				require.NoError(t, err)
+
+				err = conn.Close()
+				require.NoError(t, err)
+
+				wg.Wait()
+
+				assert.NotNil(t, proxy.target)
+
+				if lastTarget != nil {
+					test.expectSame(t, lastTarget, proxy.target)
+				}
+
+				lastTarget = proxy.target
+			}
+		})
+	}
+}