115 lines
3.1 KiB
Lua
115 lines
3.1 KiB
Lua
-- Based on https://github.com/Klaessen/openresty-loadbalancers/blob/main/sticky-balancer.lua
|
|
|
|
local _M = {}
|
|
|
|
local balancer = require "ngx.balancer"
|
|
local cookie_name = "INVIDIOUS_SERVER_ID"
|
|
local servers
|
|
local weighted_servers
|
|
local domain
|
|
|
|
-- Generate a weighted server list based on weights
|
|
local function generate_weighted_server_list(servers)
|
|
local weighted_servers = {}
|
|
for _, server in ipairs(servers) do
|
|
for i = 1, server.weight do
|
|
table.insert(weighted_servers, server)
|
|
end
|
|
end
|
|
return weighted_servers
|
|
end
|
|
|
|
-- Hash function to select server
|
|
local function hash(key, num_buckets)
|
|
local hash = ngx.crc32_long(key)
|
|
return (hash % num_buckets) + 1
|
|
end
|
|
|
|
-- Select server based on cookie or assign a new one
|
|
local function select_server()
|
|
local cookie = ngx.var["cookie_" .. cookie_name]
|
|
local server_index
|
|
|
|
|
|
math.randomseed(os.time())
|
|
|
|
if cookie then
|
|
server_index = tonumber(cookie)
|
|
ngx.header["X-Server-Id"] = server_index
|
|
else
|
|
server_index = math.random(#servers)
|
|
-- https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Set-Cookie#partitioned
|
|
ngx.header["Set-Cookie"] = cookie_name .. "=" .. server_index .. "; domain=" .. domain .. "; Path=/; HttpOnly; SameSite=None; Secure; Partitioned"
|
|
ngx.header["X-Server-Id"] = server_index
|
|
end
|
|
|
|
local server = weighted_servers[server_index]
|
|
return server
|
|
end
|
|
|
|
function _M.run(upstreams)
|
|
domain = ".nadeko.net"
|
|
local host = ngx.req.get_headers()["Host"]
|
|
|
|
-- TOR Support
|
|
if string.match(host, ".onion") then
|
|
domain = host
|
|
end
|
|
|
|
servers = upstreams
|
|
weighted_servers = generate_weighted_server_list(servers)
|
|
local ok, err
|
|
local args = ngx.req.get_uri_args()
|
|
|
|
if args then
|
|
for key, server_index in pairs(args) do
|
|
if key == "backend" then
|
|
server_index = tonumber(server_index)
|
|
|
|
-- To redirect to another backend if user inputs a backend that doesn't exists
|
|
-- Ex: ?backend=5 will give you X-Server-Id=1 (Backend 1)
|
|
val = val % #servers
|
|
if val == 0 then
|
|
val = #servers
|
|
end
|
|
|
|
ok, err = balancer.set_current_peer(servers[val][1], servers[val][2])
|
|
|
|
if not ok then
|
|
-- ngx.say("No peer available")
|
|
return ngx.exit(502)
|
|
end
|
|
|
|
ngx.header["Set-Cookie"] = cookie_name .. "=" .. server_index .. "; domain=" .. domain .. "; Path=/; HttpOnly; SameSite=None; Secure; Partitioned"
|
|
ngx.header["X-Server-Id"] = server_index
|
|
end
|
|
end
|
|
end
|
|
|
|
local server = select_server()
|
|
if not server then
|
|
-- ngx.say("No peer available")
|
|
return ngx.exit(502)
|
|
end
|
|
|
|
if string.match(server[1], 'unix:') then
|
|
ok, err = balancer.set_current_peer(server[1])
|
|
else
|
|
ok, err = balancer.set_current_peer(server[1], server[2])
|
|
end
|
|
|
|
if not ok then
|
|
ngx.log(ngx.ERR, "Failed to set the current peer: ", err)
|
|
-- ngx.say("Failed to set the current peer")
|
|
return ngx.exit(500)
|
|
end
|
|
|
|
-- https://github.com/openresty/lua-resty-core/blob/master/lib/ngx/balancer.md#enable_keepalive
|
|
ok, err = balancer.enable_keepalive(60, 1000)
|
|
if not ok then
|
|
ngx.log(ngx.ERR, "failed to set keepalive: ", err)
|
|
return
|
|
end
|
|
end
|
|
|
|
return _M
|