118 lines
3.1 KiB
Lua
118 lines
3.1 KiB
Lua
local balancer_resty = require("balancer.resty")
|
|
local resty_chash = require("resty.chash")
|
|
local util = require("util")
|
|
local ck = require("resty.cookie")
|
|
|
|
local _M = balancer_resty:new({ factory = resty_chash, name = "sticky" })
|
|
|
|
function _M.new(self, backend)
|
|
local nodes = util.get_nodes(backend.endpoints)
|
|
local digest_func = util.md5_digest
|
|
if backend["sessionAffinityConfig"]["cookieSessionAffinity"]["hash"] == "sha1" then
|
|
digest_func = util.sha1_digest
|
|
end
|
|
|
|
local o = {
|
|
instance = self.factory:new(nodes),
|
|
cookie_name = backend["sessionAffinityConfig"]["cookieSessionAffinity"]["name"] or "route",
|
|
cookie_expires = backend["sessionAffinityConfig"]["cookieSessionAffinity"]["expires"],
|
|
cookie_max_age = backend["sessionAffinityConfig"]["cookieSessionAffinity"]["maxage"],
|
|
digest_func = digest_func,
|
|
}
|
|
setmetatable(o, self)
|
|
self.__index = self
|
|
return o
|
|
end
|
|
|
|
local function encrypted_endpoint_string(self, endpoint_string)
|
|
local encrypted, err = self.digest_func(endpoint_string)
|
|
if err ~= nil then
|
|
ngx.log(ngx.ERR, err)
|
|
end
|
|
|
|
return encrypted
|
|
end
|
|
|
|
local function parse_cookie_expires(expires)
|
|
local time = tonumber(string.sub(expires, 0, string.len(expires) - 1))
|
|
if time == nil then
|
|
return nil, string.format("the time of expires (%s) is wrong", expires)
|
|
end
|
|
|
|
local unit = string.sub(expires, -1)
|
|
if unit == "y" then
|
|
return time * 60 * 60 * 24 * 365
|
|
elseif unit == "M" then
|
|
return time * 60 * 60 * 24 * 30
|
|
elseif unit == "w" then
|
|
return time * 60 * 60 * 24 * 7
|
|
elseif unit == "d" then
|
|
return time * 60 * 60 * 24
|
|
elseif unit == "h" then
|
|
return time * 60 * 60
|
|
elseif unit == "m" then
|
|
return time * 60
|
|
elseif unit == "s" then
|
|
return time
|
|
else
|
|
return nil, string.format("the unit of expires (%s) is wrong, validated unit includes: y, M, w, d, h, m and s", expires)
|
|
end
|
|
end
|
|
|
|
local function set_cookie(self, value)
|
|
local cookie, err = ck:new()
|
|
if not cookie then
|
|
ngx.log(ngx.ERR, err)
|
|
end
|
|
|
|
local cookie_data = {
|
|
key = self.cookie_name,
|
|
value = value,
|
|
path = ngx.var.location_path,
|
|
domain = ngx.var.host,
|
|
httponly = true,
|
|
}
|
|
|
|
if self.cookie_expires then
|
|
local expires, err = parse_cookie_expires(self.cookie_expires)
|
|
if err then
|
|
ngx.log(ngx.WARN, string.format("error when parsing cookie expires: %s, ignoring it", tostring(err)))
|
|
else
|
|
cookie_data.expires = ngx.cookie_time(expires)
|
|
end
|
|
end
|
|
|
|
if self.cookie_max_age then
|
|
cookie_data.max_age = tonumber(self.cookie_max_age)
|
|
end
|
|
|
|
local ok
|
|
ok, err = cookie:set(cookie_data)
|
|
if not ok then
|
|
ngx.log(ngx.ERR, err)
|
|
end
|
|
end
|
|
|
|
local function pick_random(instance)
|
|
local index = math.random(instance.npoints)
|
|
return instance:next(index)
|
|
end
|
|
|
|
function _M.balance(self)
|
|
local cookie, err = ck:new()
|
|
if not cookie then
|
|
ngx.log(ngx.ERR, err)
|
|
return pick_random(self.instance)
|
|
end
|
|
|
|
local key = cookie:get(self.cookie_name)
|
|
if not key then
|
|
local tmp_endpoint_string = pick_random(self.instance)
|
|
key = encrypted_endpoint_string(self, tmp_endpoint_string)
|
|
set_cookie(self, key)
|
|
end
|
|
|
|
return self.instance:find(key)
|
|
end
|
|
|
|
return _M
|