Skip to content

Commit

Permalink
Fixing some use-cases of the response rate limiting plugin
Browse files Browse the repository at this point in the history
  • Loading branch information
subnetmarco authored and thibaultcha committed May 27, 2016
1 parent a458c15 commit b6c0b85
Show file tree
Hide file tree
Showing 5 changed files with 72 additions and 5 deletions.
3 changes: 3 additions & 0 deletions kong/plugins/rate-limiting/handler.lua
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,9 @@ local responses = require "kong.tools.responses"
local singletons = require "kong.singletons"
local BasePlugin = require "kong.plugins.base_plugin"

local pairs = pairs
local tostring = tostring

local RATELIMIT_LIMIT = "X-RateLimit-Limit"
local RATELIMIT_REMAINING = "X-RateLimit-Remaining"

Expand Down
22 changes: 22 additions & 0 deletions kong/plugins/response-ratelimiting/access.lua
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,13 @@ local timestamp = require "kong.tools.timestamp"
local responses = require "kong.tools.responses"
local utils = require "kong.tools.utils"

local pairs = pairs
local tostring = tostring

local _M = {}

local RATELIMIT_REMAINING = "X-RateLimit-Remaining"

local function get_identifier()
local identifier

Expand Down Expand Up @@ -64,6 +69,23 @@ function _M.execute(conf)
return responses.send_HTTP_INTERNAL_SERVER_ERROR(err)
end
end

-- Append usage headers to the upstream request. Also checks "block_on_first_violation".
for k, v in pairs(conf.limits) do
local remaining
for lk, lv in pairs(usage[k]) do
if conf.block_on_first_violation and lv.remaining == 0 then
return responses.send(429, "API rate limit exceeded for '"..k.."'")
end

if not remaining or lv.remaining < remaining then
remaining = lv.remaining
end
end

ngx.req.set_header(RATELIMIT_REMAINING.."-"..k, remaining)
end

ngx.ctx.usage = usage -- For later use
end

Expand Down
7 changes: 6 additions & 1 deletion kong/plugins/response-ratelimiting/header_filter.lua
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,11 @@ local utils = require "kong.tools.utils"
local stringy = require "stringy"
local responses = require "kong.tools.responses"

local pairs = pairs
local ipairs = ipairs
local tonumber = tonumber
local math_max = math.max

local RATELIMIT_LIMIT = "X-RateLimit-Limit"
local RATELIMIT_REMAINING = "X-RateLimit-Remaining"

Expand Down Expand Up @@ -42,7 +47,7 @@ function _M.execute(conf)
for limit_name, v in pairs(usage) do
for period_name, lv in pairs(usage[limit_name]) do
ngx.header[RATELIMIT_LIMIT.."-"..limit_name.."-"..period_name] = lv.limit
ngx.header[RATELIMIT_REMAINING.."-"..limit_name.."-"..period_name] = math.max(0, lv.remaining - (increments[limit_name] and increments[limit_name] or 0)) -- increment_value for this current request
ngx.header[RATELIMIT_REMAINING.."-"..limit_name.."-"..period_name] = math_max(0, lv.remaining - (increments[limit_name] and increments[limit_name] or 0)) -- increment_value for this current request

if increments[limit_name] and increments[limit_name] > 0 and lv.remaining <= 0 then
stop = true -- No more
Expand Down
1 change: 1 addition & 0 deletions kong/plugins/response-ratelimiting/schema.lua
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ return {
fields = {
header_name = { type = "string", default = "x-kong-limit" },
continue_on_error = { type = "boolean", default = false },
block_on_first_violation = { type = "boolean", default = false},
limits = { type = "table",
schema = {
flexible = true,
Expand Down
44 changes: 40 additions & 4 deletions spec/plugins/response-ratelimiting/access_spec.lua
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,10 @@ describe("RateLimiting Plugin", function()
{ name = "tests-response-ratelimiting2", request_host = "test2.com", upstream_url = "http://httpbin.org/" },
{ name = "tests-response-ratelimiting3", request_host = "test3.com", upstream_url = "http://httpbin.org/" },
{ name = "tests-response-ratelimiting4", request_host = "test4.com", upstream_url = "http://httpbin.org/" },
{ name = "tests-response-ratelimiting5", request_host = "test5.com", upstream_url = "http://httpbin.org/" }
{ name = "tests-response-ratelimiting5", request_host = "test5.com", upstream_url = "http://httpbin.org/" },
{ name = "tests-response-ratelimiting6", request_host = "test6.com", upstream_url = "http://httpbin.org/" },
{ name = "tests-response-ratelimiting7", request_host = "test7.com", upstream_url = "http://httpbin.org/" },
{ name = "tests-response-ratelimiting8", request_host = "test8.com", upstream_url = "http://httpbin.org/" }
},
consumer = {
{ custom_id = "consumer_123" },
Expand All @@ -43,7 +46,10 @@ describe("RateLimiting Plugin", function()
{ name = "response-ratelimiting", config = { limits = { video = { minute = 6 } } }, __api = 3 },
{ name = "response-ratelimiting", config = { limits = { video = { minute = 2 } } }, __api = 3, __consumer = 1 },
{ name = "response-ratelimiting", config = { continue_on_error = false, limits = { video = { minute = 6 } } }, __api = 4 },
{ name = "response-ratelimiting", config = { continue_on_error = true, limits = { video = { minute = 6 } } }, __api = 5 }
{ name = "response-ratelimiting", config = { continue_on_error = true, limits = { video = { minute = 6 } } }, __api = 5 },
{ name = "response-ratelimiting", config = { continue_on_error = true, limits = { video = { minute = 2 } } }, __api = 6 },
{ name = "response-ratelimiting", config = { continue_on_error = false, block_on_first_violation = true, limits = { video = { minute = 6, hour = 10 }, image = { minute = 4 } } }, __api = 7 },
{ name = "response-ratelimiting", config = { limits = { video = { minute = 6, hour = 10 }, image = { minute = 4 } } }, __api = 8 }
},
keyauth_credential = {
{ key = "apikey123", __consumer = 1 },
Expand Down Expand Up @@ -80,7 +86,6 @@ describe("RateLimiting Plugin", function()
-- Additonal request, while limit is 6/minute
local _, status = http_client.get(PROXY_URL.."/response-headers", {["x-kong-limit"] = "video=1"}, {host = "test1.com"})
assert.are.equal(429, status)

end)

it("should handle multiple limits", function()
Expand Down Expand Up @@ -167,6 +172,38 @@ describe("RateLimiting Plugin", function()
end)
end)

describe("Upstream usage headers", function()
it("should append the headers with multiple limits", function()
local response, status = http_client.get(PROXY_URL.."/get", {}, {host = "test8.com"})
assert.are.equal(200, status)
local body = cjson.decode(response)
assert.are.equal("4", body.headers["X-Ratelimit-Remaining-Image"])
assert.are.equal("6", body.headers["X-Ratelimit-Remaining-Video"])

-- Actually consume the limits
local _, status = http_client.get(PROXY_URL.."/response-headers", {["x-kong-limit"] = "video=2, image=1"}, {host = "test8.com"})
assert.are.equal(200, status)

os.execute("sleep "..SLEEP_VALUE) -- The increment happens in log_by_lua, give it some time

local response, status = http_client.get(PROXY_URL.."/get", {}, {host = "test8.com"})
assert.are.equal(200, status)
local body = cjson.decode(response)

assert.are.equal("3", body.headers["X-Ratelimit-Remaining-Image"])
assert.are.equal("4", body.headers["X-Ratelimit-Remaining-Video"])
end)
end)

it("should block on first violation", function()
local _, status = http_client.get(PROXY_URL.."/response-headers", {["x-kong-limit"] = "image=4, video=2"}, {host = "test7.com"})
assert.are.equal(200, status)

local response, status = http_client.get(PROXY_URL.."/response-headers", {["x-kong-limit"] = "video=2"}, {host = "test7.com"})
assert.are.equal(429, status)
assert.are.equal("API rate limit exceeded for 'image'", cjson.decode(response).message)
end)

describe("Continue on error", function()
after_each(function()
dao_factory:drop_schema()
Expand Down Expand Up @@ -205,7 +242,6 @@ describe("RateLimiting Plugin", function()
assert.falsy(headers["x-ratelimit-limit-video-minute"])
assert.falsy(headers["x-ratelimit-remaining-video-minute"])
end)

end)

end)

0 comments on commit b6c0b85

Please sign in to comment.