From af62b5017e7f2219656f755569975e5b3f0daec6 Mon Sep 17 00:00:00 2001
From: Thijs Schreijer <thijs@thijsschreijer.nl>
Date: Fri, 1 May 2020 12:33:59 +0200
Subject: [PATCH] change(balancer) getPeer returns the hostheader instead of
 the added hostname

BREAKING: mostly the effect is the since `useSRVname` hasn't been released
yet. The breaking part is that in case an IP address is added through
`addHost`, the returned hostname by `getPeer` will now be `nil` where
it previously was the IP address.

Semantically the return value changed from the 'host' to the 'host-header'
but as stated; changes only if it was an IP address to begin with.
---
 spec/balancer/generic_spec.lua  |  4 ++--
 src/resty/dns/balancer/base.lua | 41 ++++++++++++++++++++++-----------
 2 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/spec/balancer/generic_spec.lua b/spec/balancer/generic_spec.lua
index f3623c2..e3d5952 100644
--- a/spec/balancer/generic_spec.lua
+++ b/spec/balancer/generic_spec.lua
@@ -1581,7 +1581,7 @@ for algorithm, balancer_module in helpers.balancer_types() do
         local ip, port, hostname, handle = b:getPeer()
         assert.equal("4.3.2.1", ip)
         assert.equal(8000, port)
-        assert.equal("4.3.2.1", hostname)
+        assert.equal(nil, hostname)
         assert.equal("userdata", type(handle.__udata))
       end)
 
@@ -1591,7 +1591,7 @@ for algorithm, balancer_module in helpers.balancer_types() do
         local ip, port, hostname, handle = b:getPeer()
         assert.equal("[::1]", ip)
         assert.equal(8000, port)
-        assert.equal("::1", hostname)
+        assert.equal(nil, hostname)
         assert.equal("userdata", type(handle.__udata))
       end)
 
diff --git a/src/resty/dns/balancer/base.lua b/src/resty/dns/balancer/base.lua
index 2d5680f..1e8ea30 100644
--- a/src/resty/dns/balancer/base.lua
+++ b/src/resty/dns/balancer/base.lua
@@ -201,7 +201,7 @@ local mt_objBalancer = { __index = objBalancer }
 -- ===========================================================================
 
 -- Returns the peer info.
--- @return ip-address, port and hostname of the target, or nil+err if unavailable
+-- @return ip-address, port and hostheader for the target, or nil+err if unavailable
 -- or lookup error
 function objAddr:getPeer(cacheOnly)
   if not self.available then
@@ -219,17 +219,13 @@ function objAddr:getPeer(cacheOnly)
     local ip, port, try_list = self.host.balancer.dns.toip(self.ip, self.port, cacheOnly)
     if not ip then
       port = tostring(port) .. ". Tried: " .. tostring(try_list)
+      return ip, port
     end
-    if self.useSRVname then
-      -- return the nested SRV name as the hostname
-      return ip, port, self.ip
-    end
-    -- use the hostname as it was added to the balancer
-    return ip, port, self.host.hostname
-  else
-    -- just an IP address
-    return self.ip, self.port, self.host.hostname
+
+    return ip, port, self.hostHeader
   end
+
+  return self.ip, self.port, self.hostHeader
 end
 
 -- disables an address object from the balancer.
@@ -344,6 +340,24 @@ function objBalancer:newAddress(addr)
 
   addr.host:addWeight(addr.weight)
 
+  if addr.host.nameType ~= "name" then
+    -- hostname is an IP address
+    addr.hostHeader = nil
+  else
+    -- hostname is an actual name
+    if addr.ipType ~= "name" then
+      -- the address is an ip, so use the hostname as header value
+      addr.hostHeader = addr.host.hostname
+    else
+      -- the address itself is a nested name (SRV)
+      if addr.useSRVname then
+        addr.hostHeader = addr.ip
+      else
+        addr.hostHeader = addr.host.hostname
+      end
+    end
+  end
+
   ngx_log(ngx_DEBUG, addr.host.log_prefix, "new address for host '", addr.host.hostname,
           "' created: ", addr.ip, ":", addr.port, " (weight ", addr.weight,")")
 
@@ -815,6 +829,7 @@ function objBalancer:newHost(host)
   host.lastSorted = nil      -- last successful dns query, sorted for comparison
   host.addresses = {}        -- list of addresses (address objects) this host resolves to
   host.expire = nil          -- time when the dns query this host is based upon expires
+  host.nameType = dns_utils.hostnameType(host.hostname)  -- 'ipv4', 'ipv6' or 'name'
 
 
   -- insert into our parent balancer before recalculating (in queryDns)
@@ -1065,16 +1080,16 @@ end
 -- retain some state over retries. See also `setAddressStatus`.
 -- @param hashValue (optional) number for consistent hashing, if supported by
 -- the algorithm. The hashValue must be an (evenly distributed) `integer >= 0`.
--- @return `ip + port + hostname` + `handle`, or `nil+error`
+-- @return `ip + port + hostheader` + `handle`, or `nil+error`
 -- @within User properties
 -- @usage
 -- -- get an IP address
--- local ip, port, hostname, handle = b:getPeer()
+-- local ip, port, hostheader, handle = b:getPeer()
 --
 -- -- go do the connection stuff here...
 --
 -- -- on a retry do:
--- ip, port, hostname, handle = b:getPeer(true, handle)  -- pass in previous 'handle'
+-- ip, port, hostheader, handle = b:getPeer(true, handle)  -- pass in previous 'handle'
 --
 -- -- go try again
 --