refactor(rate-limiting) remove kong.tools.* dependency from rate-limi…

…ting plugin

kong/plugins/rate-limiting/daos.lua

@@ -1,3 +1,5 @@
 return {
-  tables = {"ratelimiting_metrics"}
+  tables = {
+    "ratelimiting_metrics"
+  },
 }

kong/plugins/rate-limiting/handler.lua

@@ -1,49 +1,64 @@
 -- Copyright (C) Kong Inc.
-
 local policies = require "kong.plugins.rate-limiting.policies"
-local timestamp = require "kong.tools.timestamp"
-local responses = require "kong.tools.responses"
 local BasePlugin = require "kong.plugins.base_plugin"
 
-local ngx_log = ngx.log
+
+local kong = kong
+local ngx = ngx
+local max = math.max
+local time = ngx.time
 local pairs = pairs
 local tostring = tostring
-local ngx_timer_at = ngx.timer.at
+local timer_at = ngx.timer.at
+
 
 local RATELIMIT_LIMIT = "X-RateLimit-Limit"
 local RATELIMIT_REMAINING = "X-RateLimit-Remaining"
 
+
 local RateLimitingHandler = BasePlugin:extend()
 
+
 RateLimitingHandler.PRIORITY = 901
-RateLimitingHandler.VERSION = "0.1.0"
+RateLimitingHandler.VERSION = "0.2.0"
+
 
 local function get_identifier(conf)
   local identifier
 
-  -- Consumer is identified by ip address or authenticated_credential id
-  if conf.limit_by == "consumer" then
-    identifier = ngx.ctx.authenticated_consumer and ngx.ctx.authenticated_consumer.id
-    if not identifier and ngx.ctx.authenticated_credential then -- Fallback on credential
-      identifier = ngx.ctx.authenticated_credential.id
+  if conf.limit_by == "consumer" or conf.limit_by == "credential" then
+    local shared_ctx = kong.ctx.shared
+    local ngx_ctx = ngx.ctx
+
+    local consumer = shared_ctx.authenticated_consumer or
+                     ngx_ctx.authenticated_consumer
+
+    if conf.limit_by == "consumer" then
+      identifier = consumer and consumer.id
+    end
+
+    if not identifier then
+      local credential = shared_ctx.authenticated_credential or
+                         ngx_ctx.authenticated_credential
+
+      identifier = credential and credential.id
     end
-  elseif conf.limit_by == "credential" then
-    identifier = ngx.ctx.authenticated_credential and ngx.ctx.authenticated_credential.id
   end
 
   if not identifier then
-    identifier = ngx.var.remote_addr
+    identifier = kong.client.get_forwarded_ip()
   end
 
   return identifier
 end
 
+
 local function get_usage(conf, identifier, current_timestamp, limits)
   local usage = {}
   local stop
 
-  for name, limit in pairs(limits) do
-    local current_usage, err = policies[conf.policy].usage(conf, identifier, current_timestamp, name)
+  for period, limit in pairs(limits) do
+    local current_usage, err = policies[conf.policy].usage(conf, identifier, period, current_timestamp)
     if err then
       return nil, nil, err
     end
@@ -52,30 +67,41 @@ local function get_usage(conf, identifier, current_timestamp, limits)
     local remaining = limit - current_usage
 
     -- Recording usage
-    usage[name] = {
+    usage[period] = {
       limit = limit,
-      remaining = remaining
+      remaining = remaining,
     }
 
     if remaining <= 0 then
-      stop = name
+      stop = period
     end
   end
 
   return usage, stop
 end
 
+
+local function increment(premature, conf, ...)
+  if premature then
+    return
+  end
+
+  policies[conf.policy].increment(conf, ...)
+end
+
+
 function RateLimitingHandler:new()
   RateLimitingHandler.super.new(self, "rate-limiting")
 end
 
+
 function RateLimitingHandler:access(conf)
   RateLimitingHandler.super.access(self)
-  local current_timestamp = timestamp.get_utc()
+
+  local current_timestamp = time() * 1000
 
   -- Consumer is identified by ip address or authenticated_credential id
   local identifier = get_identifier(conf)
-  local policy = conf.policy
   local fault_tolerant = conf.fault_tolerant
 
   -- Load current metric for configured period
@@ -85,45 +111,65 @@ function RateLimitingHandler:access(conf)
     hour = conf.hour,
     day = conf.day,
     month = conf.month,
-    year = conf.year
+    year = conf.year,
   }
 
   local usage, stop, err = get_usage(conf, identifier, current_timestamp, limits)
   if err then
     if fault_tolerant then
-      ngx_log(ngx.ERR, "failed to get usage: ", tostring(err))
+      kong.log.err("failed to get usage: ", tostring(err))
     else
-      return responses.send_HTTP_INTERNAL_SERVER_ERROR(err)
+      kong.log.err(err)
+      return kong.response.exit(500, { message = "An unexpected error occurred" })
     end
   end
 
   if usage then
     -- Adding headers
     if not conf.hide_client_headers then
+      local headers = {}
       for k, v in pairs(usage) do
-        ngx.header[RATELIMIT_LIMIT .. "-" .. k] = v.limit
-        ngx.header[RATELIMIT_REMAINING .. "-" .. k] = math.max(0, (stop == nil or stop == k) and v.remaining - 1 or v.remaining) -- -increment_value for this current request
+        if stop == nil or stop == k then
+          v.remaining = v.remaining - 1
+        end
+
+        headers[RATELIMIT_LIMIT .. "-" .. k] = v.limit
+        headers[RATELIMIT_REMAINING .. "-" .. k] = max(0, v.remaining)
       end
+
+      kong.ctx.plugin.headers = headers
     end
 
     -- If limit is exceeded, terminate the request
     if stop then
-      return responses.send(429, "API rate limit exceeded")
+      return kong.response.exit(429, { message = "API rate limit exceeded" })
     end
   end
 
-  local incr = function(premature, conf, limits, identifier, current_timestamp, value)
-    if premature then
-      return
+  kong.ctx.plugin.timer = function()
+    local ok, err = timer_at(0, increment, conf, limits, identifier, current_timestamp, 1)
+    if not ok then
+      kong.log.err("failed to create timer: ", err)
     end
-    policies[policy].increment(conf, limits, identifier, current_timestamp, value)
   end
+end
 
-  -- Increment metrics for configured periods if the request goes through
-  local ok, err = ngx_timer_at(0, incr, conf, limits, identifier, current_timestamp, 1)
-  if not ok then
-    ngx_log(ngx.ERR, "failed to create timer: ", err)
+
+function RateLimitingHandler:header_filter(_)
+  RateLimitingHandler.super.header_filter(self)
+
+  local headers = kong.ctx.plugin.headers
+  if headers then
+    kong.response.set_headers(headers)
   end
 end
 
+
+function RateLimitingHandler:log(_)
+  if kong.ctx.plugin.timer then
+    kong.ctx.plugin.timer()
+  end
+end
+
+
 return RateLimitingHandler

kong/plugins/rate-limiting/migrations/001_14_to_15.lua

@@ -4,6 +4,15 @@ return {
       ALTER TABLE IF EXISTS ONLY "ratelimiting_metrics"
         ALTER "period_date" TYPE TIMESTAMP WITH TIME ZONE USING "period_date" AT TIME ZONE 'UTC';
     ]],
+    teardown = function(connector)
+      assert(connector:connect_migrations())
+      assert(connector:query([[
+        DROP FUNCTION IF EXISTS "increment_rate_limits_api" (UUID, TEXT, TEXT, TIMESTAMP WITH TIME ZONE, INTEGER) CASCADE;
+        DROP FUNCTION IF EXISTS "increment_rate_limits" (UUID, TEXT, TEXT, TIMESTAMP WITHOUT TIME ZONE, INTEGER) CASCADE;
+        DROP FUNCTION IF EXISTS "increment_rate_limits" (UUID, TEXT, TEXT, TIMESTAMP WITH TIME ZONE, INTEGER) CASCADE;
+        DROP FUNCTION IF EXISTS "increment_rate_limits" (UUID, UUID, TEXT, TEXT, TIMESTAMP WITH TIME ZONE, INTEGER) CASCADE;
+      ]]))
+    end,
   },
 
   cassandra = {