Apply address translator right after hosts resolved

In order for DNS feature to work properly with AddressTranslator
we need to address being translated right after DNS is resolved.

Author: dkropachev

conn_test.go

@@ -1257,7 +1257,7 @@ func (srv *TestServer) session() (*Session, error) {
 }
 
 func (srv *TestServer) host() *HostInfo {
-	hosts, err := hostInfo(nil, srv.Address, 9042)
+	hosts, err := hostInfo(nil, nil, srv.Address, 9042)
 	if err != nil {
 		srv.t.Fatal(err)
 	}

control.go

@@ -141,7 +141,7 @@ func (c *controlConn) heartBeat() {
 	}
 }
 
-func hostInfo(resolver DNSResolver, addr string, defaultPort int) ([]*HostInfo, error) {
+func hostInfo(resolver DNSResolver, translateAddressPort func(addr net.IP, port int) (net.IP, int), addr string, defaultPort int) ([]*HostInfo, error) {
 	var port int
 	host, portStr, err := net.SplitHostPort(addr)
 	if err != nil {
@@ -174,7 +174,12 @@ func hostInfo(resolver DNSResolver, addr string, defaultPort int) ([]*HostInfo,
 
 	for _, ip := range ips {
 		if validIpAddr(ip) {
-			hosts = append(hosts, &HostInfo{hostname: host, connectAddress: ip, port: port})
+			hh := &HostInfo{hostname: host, connectAddress: ip, port: port}
+			hh.untranslatedConnectAddress = ip
+			if translateAddressPort != nil {
+				hh.connectAddress, hh.port = translateAddressPort(ip, port)
+			}
+			hosts = append(hosts, hh)
 		}
 	}
 
@@ -424,7 +429,7 @@ func (c *controlConn) attemptReconnect() error {
 	c.session.logger.Printf("gocql: control falling back to initial contact points.\n")
 	// Fallback to initial contact points, as it may be the case that all known initialHosts
 	// changed their IPs while keeping the same hostname(s).
-	initialHosts, resolvErr := addrsToHosts(c.session.cfg.DNSResolver, c.session.cfg.Hosts, c.session.cfg.Port, c.session.logger)
+	initialHosts, resolvErr := addrsToHosts(c.session.cfg.DNSResolver, c.session.cfg.translateAddressPort, c.session.cfg.Hosts, c.session.cfg.Port, c.session.logger)
 	if resolvErr != nil {
 		return fmt.Errorf("resolve contact points' hostnames: %v", resolvErr)
 	}

control_test.go

@@ -44,7 +44,7 @@ func TestHostInfo_Lookup(t *testing.T) {
 	}
 
 	for i, test := range tests {
-		hosts, err := hostInfo(resolver, test.addr, 1)
+		hosts, err := hostInfo(resolver, nil, test.addr, 1)
 		if err != nil {
 			t.Errorf("%d: %v", i, err)
 			continue

exec.go

@@ -69,7 +69,7 @@ func NewSingleHostQueryExecutor(cfg *ClusterConfig) (e SingleHostQueryExecutor,
 	}
 
 	var hosts []*HostInfo
-	if hosts, err = addrsToHosts(c.DNSResolver, c.Hosts, c.Port, c.Logger); err != nil {
+	if hosts, err = addrsToHosts(c.DNSResolver, c.translateAddressPort, c.Hosts, c.Port, c.Logger); err != nil {
 		err = fmt.Errorf("addrs to hosts: %w", err)
 		return
 	}

filters.go

@@ -70,7 +70,7 @@ func DataCentreHostFilter(dataCenter string) HostFilter {
 // WhiteListHostFilter filters incoming hosts by checking that their address is
 // in the initial hosts whitelist.
 func WhiteListHostFilter(hosts ...string) HostFilter {
-	hostInfos, err := addrsToHosts(defaultDnsResolver, hosts, 9042, nopLogger{})
+	hostInfos, err := addrsToHosts(defaultDnsResolver, nil, hosts, 9042, nopLogger{})
 	if err != nil {
 		// dont want to panic here, but rather not break the API
 		panic(fmt.Errorf("unable to lookup host info from address: %v", err))

session.go

@@ -116,10 +116,10 @@ var queryPool = &sync.Pool{
 	},
 }
 
-func addrsToHosts(resolver DNSResolver, addrs []string, defaultPort int, logger StdLogger) ([]*HostInfo, error) {
+func addrsToHosts(resolver DNSResolver, translateAddressPort func(addr net.IP, port int) (net.IP, int), addrs []string, defaultPort int, logger StdLogger) ([]*HostInfo, error) {
 	var hosts []*HostInfo
 	for _, hostaddr := range addrs {
-		resolvedHosts, err := hostInfo(resolver, hostaddr, defaultPort)
+		resolvedHosts, err := hostInfo(resolver, translateAddressPort, hostaddr, defaultPort)
 		if err != nil {
 			// Try other hosts if unable to resolve DNS name
 			if _, ok := err.(*net.DNSError); ok {
@@ -259,7 +259,7 @@ func (s *Session) init() error {
 		return nil
 	}
 
-	hosts, err := addrsToHosts(s.cfg.DNSResolver, s.cfg.Hosts, s.cfg.Port, s.logger)
+	hosts, err := addrsToHosts(s.cfg.DNSResolver, s.cfg.translateAddressPort, s.cfg.Hosts, s.cfg.Port, s.logger)
 	if err != nil {
 		return err
 	}