Adds Wilcard check for hostname. Adds wildcard hostname tests.

rootfs/etc/nginx/lua/balancer/sticky.lua

@@ -105,6 +105,16 @@ end
 local function should_set_cookie(self)
   if self.cookie_session_affinity.locations then
     local locs = self.cookie_session_affinity.locations[ngx.var.host]
+    if locs == nil then
+      -- Based off of wildcard hostname in ../certificate.lua
+      local wildcardHostname, _, err = ngx.re.sub(ngx.var.host, "^[^\\.]+\\.", "*.", "jo")
+      if err then
+        ngx.log(ngx.ERR, "error: ", err);
+      elseif wildcardHostname then
+        locs = self.cookie_session_affinity.locations[wildcardHostname]
+      end
+    end
+
     if locs ~= nil then
       for _, path in pairs(locs) do
         if ngx.var.location_path == path then

rootfs/etc/nginx/lua/test/balancer/sticky_test.lua

@@ -263,3 +263,209 @@ describe("Sticky", function()
     end)
   end)
 end)
+
+describe("StickyWildcard", function()
+  before_each(function()
+    mock_ngx({ var = { location_path = "/", host = "dev.test.com" } })
+  end)
+
+  after_each(function()
+    reset_ngx()
+  end)
+
+  local test_backend = get_test_backend()
+  local test_backend_endpoint= test_backend.endpoints[1].address .. ":" .. test_backend.endpoints[1].port
+
+  describe("balance()", function()
+    local mocked_cookie_new = cookie.new
+
+    before_each(function()
+      package.loaded["balancer.sticky"] = nil
+      sticky = require("balancer.sticky")
+    end)
+
+    after_each(function()
+      cookie.new = mocked_cookie_new
+    end)
+
+    context("when client doesn't have a cookie set and location is in cookie_locations", function()
+      it("picks an endpoint for the client", function()
+        local sticky_balancer_instance = sticky:new(test_backend)
+        local peer = sticky_balancer_instance:balance()
+        assert.equal(peer, test_backend_endpoint)
+      end)
+
+      it("sets a cookie on the client", function()
+        local s = {}
+        cookie.new = function(self)
+          local cookie_instance = {
+            set = function(self, payload)
+              assert.equal(payload.key, test_backend.sessionAffinityConfig.cookieSessionAffinity.name)
+              assert.equal(payload.path, ngx.var.location_path)
+              assert.equal(payload.domain, nil)
+              assert.equal(payload.httponly, true)
+              assert.equal(payload.secure, false)
+              return true, nil
+            end,
+            get = function(k) return false end,
+          }
+          s = spy.on(cookie_instance, "set")
+          return cookie_instance, false
+        end
+        local b = get_test_backend()
+        b.sessionAffinityConfig.cookieSessionAffinity.locations = {}
+        b.sessionAffinityConfig.cookieSessionAffinity.locations["*.test.com"] = {"/"}
+        local sticky_balancer_instance = sticky:new(b)
+        assert.has_no.errors(function() sticky_balancer_instance:balance() end)
+        assert.spy(s).was_called()
+      end)
+
+      it("sets a secure cookie on the client when being in ssl mode", function()
+        ngx.var.https = "on"
+        local s = {}
+        cookie.new = function(self)
+          local cookie_instance = {
+            set = function(self, payload)
+              assert.equal(payload.key, test_backend.sessionAffinityConfig.cookieSessionAffinity.name)
+              assert.equal(payload.path, ngx.var.location_path)
+              assert.equal(payload.domain, nil)
+              assert.equal(payload.httponly, true)
+              assert.equal(payload.secure, true)
+              return true, nil
+            end,
+            get = function(k) return false end,
+          }
+          s = spy.on(cookie_instance, "set")
+          return cookie_instance, false
+        end
+        local b = get_test_backend()
+        b.sessionAffinityConfig.cookieSessionAffinity.locations = {}
+        b.sessionAffinityConfig.cookieSessionAffinity.locations["*.test.com"] = {"/"}
+        local sticky_balancer_instance = sticky:new(b)
+        assert.has_no.errors(function() sticky_balancer_instance:balance() end)
+        assert.spy(s).was_called()
+      end)
+    end)
+
+    context("when client doesn't have a cookie set and location not in cookie_locations", function()
+      it("picks an endpoint for the client", function()
+        local sticky_balancer_instance = sticky:new(test_backend)
+        local peer = sticky_balancer_instance:balance()
+        assert.equal(peer, test_backend_endpoint)
+      end)
+
+      it("does not set a cookie on the client", function()
+        local s = {}
+        cookie.new = function(self)
+          local cookie_instance = {
+            set = function(self, payload)
+              assert.equal(payload.key, test_backend.sessionAffinityConfig.cookieSessionAffinity.name)
+              assert.equal(payload.path, ngx.var.location_path)
+              assert.equal(payload.domain, ngx.var.host)
+              assert.equal(payload.httponly, true)
+              return true, nil
+            end,
+            get = function(k) return false end,
+          }
+          s = spy.on(cookie_instance, "set")
+          return cookie_instance, false
+        end
+        local sticky_balancer_instance = sticky:new(get_test_backend())
+        assert.has_no.errors(function() sticky_balancer_instance:balance() end)
+        assert.spy(s).was_not_called()
+      end)
+    end)
+
+    context("when client has a cookie set", function()
+      it("does not set a cookie", function()
+        local s = {}
+        cookie.new = function(self)
+          local return_obj = {
+            set = function(v) return false, nil end,
+            get = function(k) return test_backend_endpoint end,
+          }
+          s = spy.on(return_obj, "set")
+          return return_obj, false
+        end
+        local sticky_balancer_instance = sticky:new(test_backend)
+        assert.has_no.errors(function() sticky_balancer_instance:balance() end)
+        assert.spy(s).was_not_called()
+      end)
+
+      it("returns the correct endpoint for the client", function()
+        local sticky_balancer_instance = sticky:new(test_backend)
+        local peer = sticky_balancer_instance:balance()
+        assert.equal(peer, test_backend_endpoint)
+      end)
+    end)
+  end)
+
+  local function get_several_test_backends(change_on_failure)
+    return {
+      name = "access-router-production-web-80",
+      endpoints = {
+        { address = "10.184.7.40", port = "8080", maxFails = 0, failTimeout = 0 },
+        { address = "10.184.7.41", port = "8080", maxFails = 0, failTimeout = 0 },
+      },
+      sessionAffinityConfig = {
+        name = "cookie",
+        cookieSessionAffinity = { name = "test_name", hash = "sha1", change_on_failure = change_on_failure }
+      },
+    }
+  end
+
+  describe("balance() after error", function()
+    local mocked_cookie_new = cookie.new
+
+    before_each(function()
+      package.loaded["balancer.sticky"] = nil
+      sticky = require("balancer.sticky")
+    end)
+
+    after_each(function()
+      cookie.new = mocked_cookie_new
+    end)
+
+    context("when request to upstream fails", function()
+      it("changes upstream when change_on_failure option is true", function()
+        -- create sticky cookie
+        cookie.new = function(self)
+          local return_obj = {
+            set = function(v) return false, nil end,
+            get = function(k) return "" end,
+          }
+          return return_obj, false
+        end
+
+        local options = {false, true}
+
+        for _, option in ipairs(options) do
+          local sticky_balancer_instance = sticky:new(get_several_test_backends(option))
+
+          local old_upstream = sticky_balancer_instance:balance()
+          for _ = 1, 100 do
+            -- make sure upstream doesn't change on subsequent calls of balance()
+            assert.equal(old_upstream, sticky_balancer_instance:balance())
+          end
+
+          -- simulate request failure
+          sticky_balancer_instance.get_last_failure = function()
+            return "failed"
+          end
+          _G.ngx.var = { upstream_addr = old_upstream }
+
+          for _ = 1, 100 do
+            local new_upstream = sticky_balancer_instance:balance()
+            if option == false then
+              -- upstream should be the same inspite of error, if change_on_failure option is false
+              assert.equal(new_upstream, old_upstream)
+            else
+              -- upstream should change after error, if change_on_failure option is true
+              assert.not_equal(new_upstream, old_upstream)
+            end
+          end
+        end
+      end)
+    end)
+  end)
+end)