/
sticky.lua
146 lines (116 loc) · 3.63 KB
/
sticky.lua
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
local balancer_resty = require("balancer.resty")
local ck = require("resty.cookie")
local ngx_balancer = require("ngx.balancer")
local split = require("util.split")
local _M = balancer_resty:new()
local DEFAULT_COOKIE_NAME = "route"
function _M.cookie_name(self)
return self.cookie_session_affinity.name or DEFAULT_COOKIE_NAME
end
function _M.new(self)
local o = {
alternative_backends = nil,
cookie_session_affinity = nil,
traffic_shaping_policy = nil
}
setmetatable(o, self)
self.__index = self
return o
end
function _M.get_cookie(self)
local cookie, err = ck:new()
if not cookie then
ngx.log(ngx.ERR, err)
end
return cookie:get(self:cookie_name())
end
function _M.set_cookie(self, value)
local cookie, err = ck:new()
if not cookie then
ngx.log(ngx.ERR, err)
end
local cookie_path = self.cookie_session_affinity.path
if not cookie_path then
cookie_path = ngx.var.location_path
end
local cookie_data = {
key = self:cookie_name(),
value = value,
path = cookie_path,
httponly = true,
secure = ngx.var.https == "on",
}
if self.cookie_session_affinity.expires and self.cookie_session_affinity.expires ~= "" then
cookie_data.expires = ngx.cookie_time(ngx.time() + tonumber(self.cookie_session_affinity.expires))
end
if self.cookie_session_affinity.maxage and self.cookie_session_affinity.maxage ~= "" then
cookie_data.max_age = tonumber(self.cookie_session_affinity.maxage)
end
local ok
ok, err = cookie:set(cookie_data)
if not ok then
ngx.log(ngx.ERR, err)
end
end
function _M.get_last_failure()
return ngx_balancer.get_last_failure()
end
local function get_failed_upstreams()
local indexed_upstream_addrs = {}
local upstream_addrs = split.split_upstream_var(ngx.var.upstream_addr) or {}
for _, addr in ipairs(upstream_addrs) do
indexed_upstream_addrs[addr] = true
end
return indexed_upstream_addrs
end
local function should_set_cookie(self)
if self.cookie_session_affinity.locations and ngx.var.host then
local locs = self.cookie_session_affinity.locations[ngx.var.host]
if locs == nil then
-- Based off of wildcard hostname in ../certificate.lua
local wildcard_host, _, err = ngx.re.sub(ngx.var.host, "^[^\\.]+\\.", "*.", "jo")
if err then
ngx.log(ngx.ERR, "error: ", err);
elseif wildcard_host then
locs = self.cookie_session_affinity.locations[wildcard_host]
end
end
if locs ~= nil then
for _, path in pairs(locs) do
if ngx.var.location_path == path then
return true
end
end
end
end
return false
end
function _M.balance(self)
local upstream_from_cookie
local key = self:get_cookie()
if key then
upstream_from_cookie = self.instance:find(key)
end
local last_failure = self.get_last_failure()
local should_pick_new_upstream = last_failure ~= nil and self.cookie_session_affinity.change_on_failure or
upstream_from_cookie == nil
if not should_pick_new_upstream then
return upstream_from_cookie
end
local new_upstream
new_upstream, key = self:pick_new_upstream(get_failed_upstreams())
if not new_upstream then
ngx.log(ngx.WARN, string.format("failed to get new upstream; using upstream %s", new_upstream))
elseif should_set_cookie(self) then
self:set_cookie(key)
end
return new_upstream
end
function _M.sync(self, backend)
-- reload balancer nodes
balancer_resty.sync(self, backend)
self.traffic_shaping_policy = backend.trafficShapingPolicy
self.alternative_backends = backend.alternativeBackends
self.cookie_session_affinity = backend.sessionAffinityConfig.cookieSessionAffinity
end
return _M