Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(bk-cache/jwt-key.lua): ha refactor to make the dp robust #89

Merged
merged 7 commits into from
Nov 22, 2024
45 changes: 39 additions & 6 deletions src/apisix/editions/ee/plugins/bk-cache/access-token.lua
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ local core = require("apisix.core")
local access_token_define = require("apisix.plugins.bk-define.access-token")
-- local bkauth_component = require("apisix.plugins.bk-components.bkauth")
local ssm_component = require("apisix.plugins.bk-components.ssm")
local lru_new = require("resty.lrucache").new

local ACCESS_TOKEN_CACHE_TTL = 600
local ACCESS_TOKEN_CACHE_COUNT = 2000
Expand All @@ -31,11 +32,14 @@ local access_token_lrucache = core.lrucache.new(
}
)

local ACCESS_TOKEN_FALLBACK_CACHE_TTL = 60 * 60 * 24
local ACCESS_TOKEN_FALLBACK_CACHE_COUNT = 2000
local access_token_fallback_lrucache = lru_new(ACCESS_TOKEN_FALLBACK_CACHE_COUNT)

local _M = {}

local function get_access_token(access_token)
local err
err = "authentication based on access_token is not supported"

-- local bkauth_token, err = bkauth_component.verify_access_token(access_token)
-- if bkauth_token ~= nil then
Expand All @@ -50,26 +54,55 @@ local function get_access_token(access_token)
if ssm_token ~= nil then
return {
token = access_token_define.new(ssm_token.bk_app_code, ssm_token.username, ssm_token.expires_in),
}
}, nil
end

return nil, err
end

return {
err = err,
}
err = "authentication based on access_token is not supported"
return nil, err
end

function _M.get_access_token(access_token)
local key = access_token
local result, err = access_token_lrucache(key, nil, get_access_token, access_token)
if result == nil then
-- if the service is down(100% down), we can use the fallback cache, make the dp robust
if err == "connection refused" then
-- try to use the fallback cache
result = access_token_fallback_lrucache:get(key)
if result ~= nil then
core.log.error("the ssm down, error: ", err, " use the fallback cache. ",
"key=", key, " result=", core.json.delay_encode(result))
return result.token, nil
end

err = "get_access_token failed, error: " .. err
end

return nil, err
end
return result.token, result.err

-- if the service is ok, update the fallback cache, keep it the newest
-- currently, the access_token(ee) is 24 hours, so the expires_in < 24 hours, and maybe the expires_in < 0
local expires_in = result.token:get_expires_in()
if expires_in > 0 and expires_in <= ACCESS_TOKEN_FALLBACK_CACHE_TTL then
-- if the access_token will expire in 24 hours, set the ttl shorter,
-- otherwise, when ssm down, some access_token will valid even it's already expired
access_token_fallback_lrucache:set(key, result, expires_in)
else
-- if expires_in < 0, also set 24 hours ttl
access_token_fallback_lrucache:set(key, result, ACCESS_TOKEN_FALLBACK_CACHE_TTL)
end


return result.token, nil
end

if _TEST then -- luacheck: ignore
_M._get_access_token = get_access_token
_M._access_token_fallback_lrucache = access_token_fallback_lrucache
end

return _M
27 changes: 27 additions & 0 deletions src/apisix/editions/ee/plugins/bk-cache/bk-token.lua
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
--
local core = require("apisix.core")
local bklogin_component = require("apisix.plugins.bk-components.bklogin")
local lru_new = require("resty.lrucache").new

local BK_TOKEN_CACHE_TTL = 300
local BK_TOKEN_CACHE_COUNT = 2000
Expand All @@ -29,15 +30,41 @@ local bk_token_lrucache = core.lrucache.new(
}
)

local BK_TOKEN_FALLBACK_CACHE_TTL = 60 * 60 * 24
local BK_TOKEN_FALLBACK_CACHE_COUNT = 2000
local bk_token_fallback_lrucache = lru_new(BK_TOKEN_FALLBACK_CACHE_COUNT)

local _M = {}

function _M.get_username_by_bk_token(bk_token)
local key = bk_token
local result, err = bk_token_lrucache(key, nil, bklogin_component.get_username_by_bk_token, bk_token)
if result == nil then
-- if the service is down(100% down), we can use the fallback cache, make the dp robust
-- if the bklogin down, no new bk_token will be generated; but the old bk_token maybe expired
if err == "connection refused" then
-- try to use the fallback cache
result = bk_token_fallback_lrucache:get(key)
if result ~= nil then
core.log.error("the bklogin down, error: ", err, " use the fallback cache. ",
"key=", key, " result=", core.json.delay_encode(result))
return result.username, result.error_message
end

err = "get_username_by_bk_token failed, error: " .. err
end

return nil, err
end

-- if the service is ok, update the fallback cache, keep it the newest
bk_token_fallback_lrucache:set(key, result, BK_TOKEN_FALLBACK_CACHE_TTL)

return result.username, result.error_message
end

if _TEST then -- luacheck: ignore
_M._bk_token_fallback_lrucache = bk_token_fallback_lrucache
end

return _M
24 changes: 24 additions & 0 deletions src/apisix/editions/ee/plugins/bk-components/bklogin.lua
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,30 @@ function _M.get_username_by_bk_token(bk_token)
}
)

-- retry if some connection error, while the bklogin in bk-user 2.x
if err == "closed" or err == "connection reset by peer" then
res, err = http_client:request_uri(
url, {
method = "GET",
query = core.string.encode_args(
{
bk_token = bk_token,
}
),
ssl_verify = false,
headers = {
["Content-Type"] = "application/x-www-form-urlencoded",
},
}
)
end

-- if the ssm is down, return the raw error
if err == "connection refused" then
core.log.error("failed to request url: %s, err: %s, response: nil", url, err)
return nil, err
end

local result, _err = bk_components_utils.parse_response(res, err, true)
if result == nil then
core.log.error(
Expand Down
6 changes: 6 additions & 0 deletions src/apisix/editions/ee/plugins/bk-components/ssm.lua
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,12 @@ function _M.verify_access_token(access_token)
}
)

-- if the ssm is down, return the raw error
if err == "connection refused" then
core.log.error("failed to request url: %s, err: %s, response: nil", url, err)
return nil, err
end

local result, _err = bk_components_utils.parse_response(res, err, true)
if result == nil then
core.log.error(
Expand Down
61 changes: 50 additions & 11 deletions src/apisix/editions/ee/tests/bk-cache/test-access-token.lua
Original file line number Diff line number Diff line change
Expand Up @@ -83,9 +83,9 @@ describe(
ssm_verify_access_token_err = "ssm error"
ssm_is_configured = true

local result = access_token_cache._get_access_token("fake-access-token")
assert.is_nil(result.token)
assert.is_equal(result.err, "ssm error")
local result, err = access_token_cache._get_access_token("fake-access-token")
assert.is_nil(result)
assert.is_equal(err, "ssm error")
end
)

Expand All @@ -95,9 +95,9 @@ describe(
ssm_verify_access_token_err = "ssm error"
ssm_is_configured = nil

local result = access_token_cache._get_access_token("fake-access-token")
assert.is_nil(result.token)
assert.is_equal(result.err, "authentication based on access_token is not supported")
local result, err = access_token_cache._get_access_token("fake-access-token")
assert.is_nil(result)
assert.is_equal(err, "authentication based on access_token is not supported")
end
)

Expand All @@ -111,9 +111,9 @@ describe(
ssm_verify_access_token_err = nil
ssm_is_configured = nil

local result = access_token_cache._get_access_token("fake-access-token")
assert.is_nil(result.token)
assert.is_equal(result.err, "authentication based on access_token is not supported")
local result, err = access_token_cache._get_access_token("fake-access-token")
assert.is_nil(result)
assert.is_equal(err, "authentication based on access_token is not supported")
end
)
end
Expand Down Expand Up @@ -164,16 +164,55 @@ describe(
local result, err = access_token_cache.get_access_token(access_token)
assert.is_nil(result)
assert.is_not_nil(err)
assert.stub(ssm_component.verify_access_token).was_called(1)

-- get from cache
access_token_cache.get_access_token(access_token)
assert.stub(ssm_component.verify_access_token).was_called(1)
assert.stub(ssm_component.verify_access_token).was_called(2)

-- get from func
access_token_cache.get_access_token(uuid.generate_v4())
assert.stub(ssm_component.verify_access_token).was_called(2)
assert.stub(ssm_component.verify_access_token).was_called(3)
end
)

it(
"connection refused, miss in fallback cache", function()
ssm_verify_access_token_result = nil
ssm_verify_access_token_err = "connection refused"
ssm_is_configured = true

local access_token = uuid.generate_v4()
local result, err = access_token_cache.get_access_token(access_token)
assert.is_nil(result)
assert.is_equal(err, "get_access_token failed, error: connection refused")
assert.stub(ssm_component.verify_access_token).was_called_with(access_token)
end
)

it(
"connection refused, hit in fallback cache", function()
local cached_access_token_result = {
token = {
app_code = "my-app",
user_id = "admin",
expires_in = 100,
}
}
ssm_verify_access_token_result = nil
ssm_verify_access_token_err = "connection refused"
ssm_is_configured = true

local access_token = uuid.generate_v4()
access_token_cache._access_token_fallback_lrucache:set(access_token, cached_access_token_result, 60 * 60 * 24)

local result, err = access_token_cache.get_access_token(access_token)
assert.is_same(result, cached_access_token_result.token)
assert.is_nil(err)
assert.stub(ssm_component.verify_access_token).was_called_with(access_token)
end
)

end
)
end
Expand Down
32 changes: 32 additions & 0 deletions src/apisix/editions/ee/tests/bk-cache/test-bk-token.lua
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,38 @@ describe(
assert.stub(bklogin_component.get_username_by_bk_token).was_called(3)
end
)

it(
"connection refused, miss in fallback cache", function()
get_username_by_bk_token_result = nil
get_username_by_bk_token_err = "connection refused"

local bk_token = uuid.generate_v4()
local result, err = bk_token_cache.get_username_by_bk_token(bk_token)
assert.is_nil(result)
assert.is_equal(err, "get_username_by_bk_token failed, error: connection refused")
assert.stub(bklogin_component.get_username_by_bk_token).was_called_with(bk_token)
end
)

it(
"connection refused, hit in fallback cache", function()
local cached_get_username_by_bk_token_result = {
username = "admin",
}
get_username_by_bk_token_result = nil
get_username_by_bk_token_err = "connection refused"

local bk_token = uuid.generate_v4()
bk_token_cache._bk_token_fallback_lrucache:set(bk_token, cached_get_username_by_bk_token_result, 60 * 60 * 24)

local result, err = bk_token_cache.get_username_by_bk_token(bk_token)
assert.is_equal(result, "admin")
assert.is_nil(err)
assert.stub(bklogin_component.get_username_by_bk_token).was_called_with(bk_token)
end
)

end
)
end
Expand Down
11 changes: 11 additions & 0 deletions src/apisix/editions/ee/tests/bk-components/test-bklogin.lua
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,17 @@ describe(
end
)

it(
"connection refused", function()
response = nil
response_err = "connection refused"

local result, err = bklogin.get_username_by_bk_token("fake-bk-token")
assert.is_nil(result)
assert.equals(err, "connection refused")
end
)

it(
"bk_error_code is not 0", function()
response = {
Expand Down
11 changes: 11 additions & 0 deletions src/apisix/editions/ee/tests/bk-components/test-ssm.lua
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,17 @@ describe(
end
)

it(
"connection refused", function()
response = nil
response_err = "connection refused"

local result, err = ssm.verify_access_token("fake-access-token")
assert.is_nil(result)
assert.equals(err, "connection refused")
end
)

it(
"code is not 0", function()
response = {
Expand Down
Loading
Loading