Merge pull request Kong#953 from Mashape/plugins/ratelimiting-continu…

…eonerror

Rate-Limiting and Response Rate-Limiting can continue on DB error

kong/plugins/rate-limiting/handler.lua

@@ -26,8 +26,9 @@ local function increment(api_id, identifier, current_timestamp, value)
   -- Increment metrics for all periods if the request goes through
   local _, stmt_err = dao.ratelimiting_metrics:increment(api_id, identifier, current_timestamp, value)
   if stmt_err then
-    return responses.send_HTTP_INTERNAL_SERVER_ERROR(stmt_err)
+    return false, stmt_err
   end
+  return true
 end
 
 local function increment_async(premature, api_id, identifier, current_timestamp, value)
@@ -46,7 +47,7 @@ local function get_usage(api_id, identifier, current_timestamp, limits)
   for name, limit in pairs(limits) do
     local current_metric, err = dao.ratelimiting_metrics:find_one(api_id, identifier, current_timestamp, name)
     if err then
-      return responses.send_HTTP_INTERNAL_SERVER_ERROR(err)
+      return nil, nil, err
     end
 
     -- What is the current usage for the configured limit name?
@@ -79,30 +80,48 @@ function RateLimitingHandler:access(conf)
   local identifier = get_identifier()
   local api_id = ngx.ctx.api.id
   local is_async = conf.async
+  local is_continue_on_error = conf.continue_on_error
 
   -- Load current metric for configured period
   conf.async = nil
-  local usage, stop = get_usage(api_id, identifier, current_timestamp, conf)
-
-  -- Adding headers
-  for k, v in pairs(usage) do
-    ngx.header[constants.HEADERS.RATELIMIT_LIMIT.."-"..k] = v.limit
-    ngx.header[constants.HEADERS.RATELIMIT_REMAINING.."-"..k] = math.max(0, (stop == nil or stop == k) and v.remaining - 1 or v.remaining) -- -increment_value for this current request
+  conf.continue_on_error = nil
+  local usage, stop, err = get_usage(api_id, identifier, current_timestamp, conf)
+  if err then
+    if is_continue_on_error then
+      ngx.log(ngx.ERR, "failed to get usage: ", tostring(err))
+    else
+      return responses.send_HTTP_INTERNAL_SERVER_ERROR(err)
+    end
   end
 
-  -- If limit is exceeded, terminate the request
-  if stop then
-    return responses.send(429, "API rate limit exceeded")
-  end
+  if usage then
+    -- Adding headers
+    for k, v in pairs(usage) do
+      ngx.header[constants.HEADERS.RATELIMIT_LIMIT.."-"..k] = v.limit
+      ngx.header[constants.HEADERS.RATELIMIT_REMAINING.."-"..k] = math.max(0, (stop == nil or stop == k) and v.remaining - 1 or v.remaining) -- -increment_value for this current request
+    end
 
+    -- If limit is exceeded, terminate the request
+    if stop then
+      return responses.send(429, "API rate limit exceeded")
+    end
+  end
+
   -- Increment metrics for all periods if the request goes through
   if is_async then
     local ok, err = ngx.timer.at(0, increment_async, api_id, identifier, current_timestamp, 1)
     if not ok then
       ngx.log(ngx.ERR, "failed to create timer: ", err)
     end
   else
-    increment(api_id, identifier, current_timestamp, 1)
+    local _, err = increment(api_id, identifier, current_timestamp, 1)
+    if err then
+      if is_continue_on_error then
+        ngx.log(ngx.ERR, "failed to increment: ", tostring(err))
+      else
+        return responses.send_HTTP_INTERNAL_SERVER_ERROR(err)
+      end
+    end
   end
 end
 

kong/plugins/rate-limiting/schema.lua

@@ -9,7 +9,8 @@ return {
     day = { type = "number" },
     month = { type = "number" },
     year = { type = "number" },
-    async = { type = "boolean", default = false }
+    async = { type = "boolean", default = false },
+    continue_on_error = { type = "boolean", default = false }
   },
   self_check = function(schema, plugin_t, dao, is_update)
     local ordered_periods = { "second", "minute", "hour", "day", "month", "year"}

kong/plugins/response-ratelimiting/access.lua

@@ -24,7 +24,7 @@ local function get_current_usage(api_id, identifier, current_timestamp, limits)
     for lk, lv in pairs(v) do -- Iterare over periods
       local current_metric, err = dao.response_ratelimiting_metrics:find_one(api_id, identifier, current_timestamp, lk, k)
       if err then
-        return responses.send_HTTP_INTERNAL_SERVER_ERROR(err)
+        return false, err
       end
 
       local current_usage = current_metric and current_metric.value or 0
@@ -54,7 +54,15 @@ function _M.execute(conf)
   ngx.ctx.identifier = identifier -- For later use
 
   -- Load current metric for configured period
-  local usage = get_current_usage(api_id, identifier, current_timestamp, conf.limits)
+  local usage, err = get_current_usage(api_id, identifier, current_timestamp, conf.limits)
+  if err then
+    if conf.continue_on_error then
+      ngx.log(ngx.ERR, "failed to get usage: ", tostring(err))
+      return
+    else
+      return responses.send_HTTP_INTERNAL_SERVER_ERROR(err)
+    end
+  end
   ngx.ctx.usage = usage -- For later use
 end
 

kong/plugins/response-ratelimiting/handler.lua

@@ -22,7 +22,7 @@ function ResponseRateLimitingHandler:header_filter(conf)
 end
 
 function ResponseRateLimitingHandler:log(conf)
-  if not ngx.ctx.stop_log then
+  if not ngx.ctx.stop_log and ngx.ctx.usage then
     ResponseRateLimitingHandler.super.log(self)
     log.execute(ngx.ctx.api.id, ngx.ctx.identifier, ngx.ctx.current_timestamp, ngx.ctx.increments, ngx.ctx.usage)
   end

kong/plugins/response-ratelimiting/header_filter.lua

@@ -32,7 +32,8 @@ function _M.execute(conf)
   ngx.ctx.increments = increments
 
   local usage = ngx.ctx.usage -- Load current usage
-
+  if not usage then return end
+
   local stop
   for limit_name, v in pairs(usage) do
     for period_name, lv in pairs(usage[limit_name]) do

kong/plugins/response-ratelimiting/schema.lua

@@ -37,6 +37,7 @@ end
 return {
   fields = {
     header_name = { type = "string", default = "x-kong-limit" },
+    continue_on_error = { type = "boolean", default = false },
     limits = { type = "table",
       schema = {
         flexible = true,

spec/plugins/rate-limiting/access_spec.lua

@@ -17,15 +17,17 @@ end
 
 describe("RateLimiting Plugin", function()
 
-  setup(function()
+  local function prepare_db()
     spec_helper.prepare_db()
     spec_helper.insert_fixtures {
       api = {
         { name = "tests-rate-limiting1", request_host = "test3.com", upstream_url = "http://mockbin.com" },
         { name = "tests-rate-limiting2", request_host = "test4.com", upstream_url = "http://mockbin.com" },
         { name = "tests-rate-limiting3", request_host = "test5.com", upstream_url = "http://mockbin.com" },
         { name = "tests-rate-limiting4", request_host = "test6.com", upstream_url = "http://mockbin.com" },
-        { name = "tests-rate-limiting5", request_host = "test7.com", upstream_url = "http://mockbin.com" }
+        { name = "tests-rate-limiting5", request_host = "test7.com", upstream_url = "http://mockbin.com" },
+        { name = "tests-rate-limiting6", request_host = "test8.com", upstream_url = "http://mockbin.com" },
+        { name = "tests-rate-limiting7", request_host = "test9.com", upstream_url = "http://mockbin.com" }
       },
       consumer = {
         { custom_id = "provider_123" },
@@ -38,16 +40,20 @@ describe("RateLimiting Plugin", function()
         { name = "rate-limiting", config = { minute = 6 }, __api = 2 },
         { name = "rate-limiting", config = { minute = 3, hour = 5 }, __api = 3 },
         { name = "rate-limiting", config = { minute = 33 }, __api = 4 },
-        { name = "rate-limiting", config = { minute = 6, async = true }, __api = 5 }
+        { name = "rate-limiting", config = { minute = 6, async = true }, __api = 5 },
+        { name = "rate-limiting", config = { minute = 6, continue_on_error = false }, __api = 6 },
+        { name = "rate-limiting", config = { minute = 6, continue_on_error = true }, __api = 7 }
       },
       keyauth_credential = {
         { key = "apikey122", __consumer = 1 },
         { key = "apikey123", __consumer = 2 }
       }
     }
+  end
 
+  setup(function()
+    prepare_db()
     spec_helper.start_kong()
-
     wait()
   end)
 
@@ -148,7 +154,6 @@ describe("RateLimiting Plugin", function()
   end)
 
   describe("Async increment", function()
-
     it("should increment asynchronously", function()
       -- Default rate-limiting plugin for this API says 6/minute
         local limit = 6
@@ -167,6 +172,63 @@ describe("RateLimiting Plugin", function()
         assert.are.equal(429, status)
         assert.are.equal("API rate limit exceeded", body.message)
     end)
+  end)
+
+  describe("Continue on error", function()
+
+    local session, err, configuration
+
+    setup(function()
+      local cassandra = require "cassandra"
+      local TEST_CONF = spec_helper.get_env().conf_file
+      local env = spec_helper.get_env(TEST_CONF)
+      configuration = env.configuration
+      session, err = cassandra.spawn_session {
+        shm = "ratelimiting_specs",
+        keyspace = configuration.dao_config.keyspace,
+        contact_points = configuration.dao_config.contact_points
+      }
+      assert.falsy(err)
+    end)
 
+    after_each(function()
+      session:execute("DROP KEYSPACE "..configuration.dao_config.keyspace)
+      prepare_db()
+    end)
+
+    teardown(function()
+      session:shutdown()
+    end)
+
+    it("should not continue if an error occurs", function()
+      local _, status, headers = http_client.get(STUB_GET_URL, {}, {host = "test8.com"})
+      assert.are.equal(200, status)
+      assert.are.same(tostring(6), headers["x-ratelimit-limit-minute"])
+      assert.are.same(tostring(5), headers["x-ratelimit-remaining-minute"])
+
+      -- Simulate an error on the database
+      session:execute("DROP TABLE ratelimiting_metrics")
+
+      -- Make another request
+      local res, status, _ = http_client.get(STUB_GET_URL, {}, {host = "test8.com"})
+      assert.equal("An unexpected error occurred", cjson.decode(res).message)
+      assert.are.equal(500, status)
+    end)
+
+    it("should continue if an error occurs", function()
+      local _, status, headers = http_client.get(STUB_GET_URL, {}, {host = "test9.com"})
+      assert.are.equal(200, status)
+      assert.falsy(headers["x-ratelimit-limit-minute"])
+      assert.falsy(headers["x-ratelimit-remaining-minute"])
+
+      -- Simulate an error on the database
+      session:execute("DROP TABLE ratelimiting_metrics")
+
+      -- Make another request
+      local _, status, headers = http_client.get(STUB_GET_URL, {}, {host = "test9.com"})
+      assert.are.equal(200, status)
+      assert.falsy(headers["x-ratelimit-limit-minute"])
+      assert.falsy(headers["x-ratelimit-remaining-minute"])
+    end)
   end)
 end)