Improve host name resolution for TCP proxy

This commit is contained in:
HMH 2021-03-23 10:24:03 +00:00 committed by GitHub
parent 29908098e4
commit 8947f85ddd
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 30 additions and 59 deletions

View file

@ -31,7 +31,7 @@ func NewProxy(address string, terminationDelay time.Duration, proxyProtocol *dyn
return nil, fmt.Errorf("unknown proxyProtocol version: %d", proxyProtocol.Version)
}
// enable the refresh of the target only if the address in an IP
// 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
@ -48,23 +48,14 @@ func NewProxy(address string, terminationDelay time.Duration, proxyProtocol *dyn
// ServeTCP forwards the connection to a service.
func (p *Proxy) ServeTCP(conn WriteCloser) {
log.Debugf("Handling connection from %s", conn.RemoteAddr())
log.WithoutContext().Debugf("Handling connection from %s", conn.RemoteAddr())
// needed because of e.g. server.trackedConnection
defer conn.Close()
if p.refreshTarget {
tcpAddr, err := net.ResolveTCPAddr("tcp", p.address)
connBackend, err := p.dialBackend()
if err != nil {
log.Errorf("Error resolving tcp address: %v", err)
return
}
p.target = tcpAddr
}
connBackend, err := net.DialTCP("tcp", nil, p.target)
if err != nil {
log.Errorf("Error while connection to backend: %v", err)
log.WithoutContext().Errorf("Error while connecting to backend: %v", err)
return
}
@ -91,6 +82,19 @@ func (p *Proxy) ServeTCP(conn WriteCloser) {
<-errChan
}
func (p Proxy) dialBackend() (*net.TCPConn, error) {
if !p.refreshTarget {
return net.DialTCP("tcp", nil, p.target)
}
conn, err := net.Dial("tcp", p.address)
if err != nil {
return nil, err
}
return conn.(*net.TCPConn), nil
}
func (p Proxy) connCopy(dst, src WriteCloser, errCh chan error) {
_, err := io.Copy(dst, src)
errCh <- err

View file

@ -6,7 +6,6 @@ import (
"fmt"
"io"
"net"
"sync"
"testing"
"time"
@ -177,17 +176,16 @@ func TestLookupAddress(t *testing.T) {
testCases := []struct {
desc string
address string
expectSame assert.ComparisonAssertionFunc
expectRefresh bool
}{
{
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,
expectRefresh: true,
},
}
@ -201,44 +199,13 @@ func TestLookupAddress(t *testing.T) {
require.NotNil(t, proxy.target)
proxyListener, err := net.Listen("tcp", ":0")
conn, err := proxy.dialBackend()
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
if test.expectRefresh {
assert.NotEqual(t, test.address, conn.RemoteAddr().String())
} else {
assert.Equal(t, test.address, conn.RemoteAddr().String())
}
})
}