Fixing some use-cases of the response rate limiting plugin

kong/plugins/rate-limiting/handler.lua

@@ -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"
 

kong/plugins/response-ratelimiting/access.lua

@@ -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
 
@@ -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
 

kong/plugins/response-ratelimiting/header_filter.lua

@@ -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"
 
@@ -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

kong/plugins/response-ratelimiting/schema.lua

@@ -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,

spec/plugins/response-ratelimiting/access_spec.lua

@@ -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" },
@@ -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 },
@@ -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()
@@ -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()
@@ -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)