feat(balancer) add option to load balance hashing on cookie (#3472)

This allows better distribution for sticky sessions when many consumers
may be behind the same NAT gateway (e.g. a call center). When the
request comes in, we check for the configured cookie name. If not
present in the request, a random guid is created and used for balancing
and stored in `ngx.ctx` so the cookie can be added in the response
header. Because the cookie is created if it doesn't already exist, the
`hash_fallback` option is not available when cookies are the primary
identifier.

kong-0.13.1-0.rockspec

@@ -31,6 +31,7 @@ dependencies = {
   "lua-resty-worker-events == 0.3.3",
   "lua-resty-mediador == 0.1.2",
   "lua-resty-healthcheck == 0.4.0",
+  "lua-resty-cookie == 0.1.0",
   "lua-resty-mlcache == 2.0.2",
   -- external Kong plugins
   "kong-plugin-azure-functions == 0.1.0",

kong/dao/migrations/cassandra.lua

@@ -748,5 +748,16 @@ return {
       CREATE INDEX IF NOT EXISTS ON consumers(custom_id);
       CREATE INDEX IF NOT EXISTS ON consumers(username);
     ]]
+  },
+  {
+    name = "2018-05-17-173100_hash_on_cookie",
+    up = [[
+      ALTER TABLE upstreams ADD hash_on_cookie text;
+      ALTER TABLE upstreams ADD hash_on_cookie_path text;
+    ]],
+    down = [[
+      ALTER TABLE upstreams DROP hash_on_cookie;
+      ALTER TABLE upstreams DROP hash_on_cookie_path;
+    ]]
   }
 }

kong/dao/migrations/postgres.lua

@@ -785,4 +785,15 @@ return {
     ]],
     down = nil
   },
+  {
+    name = "2018-05-17-173100_hash_on_cookie",
+    up = [[
+      ALTER TABLE upstreams ADD hash_on_cookie text;
+      ALTER TABLE upstreams ADD hash_on_cookie_path text;
+    ]],
+    down = [[
+      ALTER TABLE upstreams DROP hash_on_cookie;
+      ALTER TABLE upstreams DROP hash_on_cookie_path;
+    ]]
+  }
 }

kong/dao/schemas/upstreams.lua

@@ -168,6 +168,7 @@ return {
         "consumer",
         "ip",
         "header",
+        "cookie",
       },
     },
     hash_fallback = {
@@ -179,6 +180,7 @@ return {
         "consumer",
         "ip",
         "header",
+        "cookie",
       },
     },
     hash_on_header = {
@@ -189,6 +191,15 @@ return {
       -- header name, if `hash_fallback == "header"`
       type = "string",
     },
+    hash_on_cookie = {
+      -- cookie name, if `hash_on` or `hash_fallback` == `"cookie"`
+      type = "string",
+    },
+    hash_on_cookie_path = {
+      -- cookie path, if `hash_on` or `hash_fallback` == `"cookie"`
+      type = "string",
+      default = "/",
+    },
     slots = {
       -- the number of slots in the loadbalancer algorithm
       type = "number",
@@ -230,6 +241,20 @@ return {
       end
     end
 
+    if config.hash_on_cookie then
+      local ok, err = utils.validate_cookie_name(config.hash_on_cookie)
+      if not ok then
+        return false, Errors.schema("Cookie name: " .. err)
+      end
+    end
+
+    if config.hash_on_cookie_path then
+      local ok, err = check_http_path(config.hash_on_cookie_path)
+      if not ok then
+        return false, Errors.schema("Cookie path: " .. err)
+      end
+    end
+
     if (config.hash_on == "header"
         and not config.hash_on_header) or
        (config.hash_fallback == "header"
@@ -238,13 +263,27 @@ return {
                                   "but no header name provided")
     end
 
+    if (config.hash_on == "cookie" or config.hash_fallback == "cookie")
+        and not config.hash_on_cookie then
+      return false, Errors.schema("Hashing on 'cookie', " ..
+                                  "but no cookie name provided")
+    end
+
     if config.hash_on and config.hash_fallback then
       if config.hash_on == "none" then
         if config.hash_fallback ~= "none" then
           return false, Errors.schema("Cannot set fallback if primary " ..
                                       "'hash_on' is not set")
         end
 
+      elseif config.hash_on == "cookie" then
+        if config.hash_fallback ~= "none" then
+          -- Cookie value will be set with the server response if not present,
+          -- so the fallback value would never be evaluated.
+          return false, Errors.schema("Cannot set `hash_fallback` if primary " ..
+                                      "`hash_on` is set to cookie")
+        end
+
       else
         if config.hash_on == config.hash_fallback then
           if config.hash_on ~= "header" then

kong/runloop/balancer.lua

@@ -1,5 +1,6 @@
 local pl_tablex = require "pl.tablex"
 local singletons = require "kong.singletons"
+local utils = require "kong.tools.utils"
 
 -- due to startup/require order, cannot use the ones from 'singletons' here
 local dns_client = require "resty.dns.client"
@@ -656,6 +657,25 @@ local create_hash = function(upstream)
       if type(identifier) == "table" then
         identifier = table_concat(identifier)
       end
+
+    elseif hash_on == "cookie" then
+      identifier = ngx.var["cookie_" .. upstream.hash_on_cookie]
+
+      -- If the cookie doesn't exist, create one and store in `ctx`
+      -- to be added to the "Set-Cookie" header in the response
+      if not identifier then
+        identifier = utils.uuid()
+
+        -- TODO: This should be added the `ngx.ctx.balancer_address`
+        -- structure, where other balancer-related values are stored,
+        -- when that is renamed/ refactored.
+        ctx.balancer_hash_cookie = {
+          key = upstream.hash_on_cookie,
+          value = identifier,
+          path = upstream.hash_on_cookie_path
+        }
+      end
+
     end
 
     if identifier then

kong/runloop/handler.lua

@@ -7,6 +7,7 @@
 --
 -- In the `access_by_lua` phase, it is responsible for retrieving the route being proxied by
 -- a consumer. Then it is responsible for loading the plugins to execute on this request.
+local ck          = require "resty.cookie"
 local utils       = require "kong.tools.utils"
 local Router      = require "kong.router"
 local ApiRouter   = require "kong.api_router"
@@ -707,6 +708,18 @@ return {
             header[upstream_status_header] = matches[0]
           end
         end
+
+        local cookie_hash_data = ctx.balancer_hash_cookie
+        if cookie_hash_data then
+          local cookie = ck:new()
+          local ok, err = cookie:set(cookie_hash_data)
+
+          if not ok then
+            log(ngx.WARN, "failed to set the cookie for hash-based load balancing: ", err,
+                          " (key=", cookie_hash_data.key,
+                          ", path=", cookie_hash_data.path, ")")
+          end
+        end
       end
     end,
     after = function(ctx)

kong/tools/utils.lua

@@ -810,4 +810,22 @@ _M.validate_header_name = function(name)
               "', allowed characters are A-Z, a-z, 0-9, '_', and '-'"
 end
 
+--- Validates a cookie name.
+-- Checks characters used in a cookie name to be valid
+-- a-z, A-Z, 0-9, '_' and '-' are allowed.
+-- @param name (string) the cookie name to verify
+-- @return the valid cookie name, or `nil+error`
+_M.validate_cookie_name = function(name)
+  if name == nil or name == "" then
+    return nil, "no cookie name provided"
+  end
+
+  if re_match(name, "^[a-zA-Z0-9-_]+$", "jo") then
+    return name
+  end
+
+  return nil, "bad cookie name '" .. name ..
+              "', allowed characters are A-Z, a-z, 0-9, '_', and '-'"
+end
+
 return _M

spec/01-unit/004-utils_spec.lua

@@ -525,6 +525,21 @@ describe("Utils", function()
       end
     end
   end)
+  it("validate_cookie_name() validates cookie names", function()
+    local header_chars = [[_-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz]]
+
+    for i = 1, 255 do
+      local c = string.char(i)
+
+      if string.find(header_chars, c, nil, true) then
+        assert(utils.validate_cookie_name(c) == c,
+          "ascii character '" .. c .. "' (" .. i .. ") should have been allowed")
+      else
+        assert(utils.validate_cookie_name(c) == nil,
+          "ascii character " .. i .. " should not have been allowed")
+      end
+    end
+  end)
   it("pack() stores results, including nils, properly", function()
     assert.same({ n = 0 }, utils.pack())
     assert.same({ n = 1 }, utils.pack(nil))

spec/01-unit/011-balancer_spec.lua

@@ -506,6 +506,27 @@ describe("Balancer", function()
       })
       assert.are.same(crc32(value), hash)
     end)
+    describe("cookie", function()
+      it("uses the cookie when present in the request", function()
+        local value = "some cookie value"
+        ngx.var.cookie_Foo = value
+        local hash = balancer._create_hash({
+          hash_on = "cookie",
+          hash_on_cookie = "Foo",
+        })
+        assert.are.same(crc32(value), hash)
+        assert.is_nil(ngx.ctx.balancer_hash_cookie)
+      end)
+      it("creates the cookie when not present in the request", function()
+        balancer._create_hash({
+          hash_on = "cookie",
+          hash_on_cookie = "Foo",
+          hash_on_cookie_path = "/",
+        })
+        assert.are.same(ngx.ctx.balancer_hash_cookie.key, "Foo")
+        assert.are.same(ngx.ctx.balancer_hash_cookie.path, "/")
+      end)
+    end)
     it("multi-header", function()
       local value = { "some header value", "another value" }
       headers.HeaderName = value
@@ -562,6 +583,29 @@ describe("Balancer", function()
         })
         assert.are.same(crc32(table.concat(value)), hash)
       end)
+      describe("cookie", function()
+        it("uses the cookie when present in the request", function()
+          local value = "some cookie value"
+          ngx.var.cookie_Foo = value
+          local hash = balancer._create_hash({
+            hash_on = "consumer",
+            hash_fallback = "cookie",
+            hash_on_cookie = "Foo",
+          })
+          assert.are.same(crc32(value), hash)
+          assert.is_nil(ngx.ctx.balancer_hash_cookie)
+        end)
+        it("creates the cookie when not present in the request", function()
+          balancer._create_hash({
+            hash_on = "consumer",
+            hash_fallback = "cookie",
+            hash_on_cookie = "Foo",
+            hash_on_cookie_path = "/",
+          })
+          assert.are.same(ngx.ctx.balancer_hash_cookie.key, "Foo")
+          assert.are.same(ngx.ctx.balancer_hash_cookie.path, "/")
+        end)
+      end)
     end)
   end)