Add Global Rate Limiting support

rootfs/etc/nginx/lua/global_throttle.lua

@@ -0,0 +1,131 @@
+local resty_global_throttle = require("resty.global_throttle")
+local resty_ipmatcher = require("resty.ipmatcher")
+local util = require("util")
+
+local ngx = ngx
+local ngx_exit = ngx.exit
+local ngx_log = ngx.log
+local ngx_ERR = ngx.ERR
+local ngx_INFO = ngx.INFO
+
+local _M = {}
+
+local DECISION_CACHE = ngx.shared.global_throttle_cache
+
+-- it does not make sense to cache decision for too little time
+-- the benefit of caching likely is negated if we cache for too little time
+-- Lua Shared Dict's time resolution for expiry is 0.001.
+local CACHE_THRESHOLD = 0.001
+
+local DEFAULT_RAW_KEY = "remote_addr"
+
+local function should_ignore_request(ignored_cidrs)
+  if not ignored_cidrs or #ignored_cidrs == 0 then
+    return false
+  end
+
+  local ignored_cidrs_matcher, err = resty_ipmatcher.new(ignored_cidrs)
+  if not ignored_cidrs_matcher then
+    ngx_log(ngx_ERR, "failed to initialize resty-ipmatcher: ", err)
+    return false
+  end
+
+  local is_ignored
+  is_ignored, err = ignored_cidrs_matcher:match(ngx.var.remote_addr)
+  if err then
+    ngx_log(ngx_ERR, "failed to match ip: '",
+      ngx.var.remote_addr, "': ", err)
+    return false
+  end
+
+  return is_ignored
+end
+
+local function is_enabled(config, location_config)
+  if config.memcached.host == "" or config.memcached.port == 0 then
+    return false
+  end
+  if location_config.limit == 0 or
+    location_config.window_size == 0 then
+    return false
+  end
+
+  if should_ignore_request(location_config.ignored_cidrs) then
+    return false
+  end
+
+  return true
+end
+
+local function get_namespaced_key_value(namespace, key_value)
+  return namespace .. key_value
+end
+
+function _M.throttle(config, location_config)
+  if not is_enabled(config, location_config) then
+    return
+  end
+
+  local key_value = util.generate_var_value(location_config.key)
+  if not key_value or key_value == "" then
+    key_value = ngx.var[DEFAULT_RAW_KEY]
+  end
+
+  local namespaced_key_value =
+    get_namespaced_key_value(location_config.namespace, key_value)
+
+  local is_limit_exceeding = DECISION_CACHE:get(namespaced_key_value)
+  if is_limit_exceeding then
+    ngx.var.global_rate_limit_exceeding = "c"
+    return ngx_exit(config.status_code)
+  end
+
+  local my_throttle, err = resty_global_throttle.new(
+    location_config.namespace,
+    location_config.limit,
+    location_config.window_size,
+    {
+      provider = "memcached",
+      host = config.memcached.host,
+      port = config.memcached.port,
+      connect_timeout = config.memcached.connect_timeout,
+      max_idle_timeout = config.memcached.max_idle_timeout,
+      pool_size = config.memcached.pool_size,
+    }
+  )
+  if err then
+    ngx.log(ngx.ERR, "faled to initialize resty_global_throttle: ", err)
+    -- fail open
+    return
+  end
+
+  local desired_delay, estimated_final_count
+  estimated_final_count, desired_delay, err = my_throttle:process(key_value)
+  if err then
+    ngx.log(ngx.ERR, "error while processing key: ", err)
+    -- fail open
+    return
+  end
+
+  if desired_delay then
+    if desired_delay > CACHE_THRESHOLD then
+      local ok
+      ok, err =
+        DECISION_CACHE:safe_add(namespaced_key_value, true, desired_delay)
+      if not ok then
+        if err ~= "exists" then
+          ngx_log(ngx_ERR, "failed to cache decision: ", err)
+        end
+      end
+    end
+
+    ngx.var.global_rate_limit_exceeding = "y"
+    ngx_log(ngx_INFO, "limit is exceeding for ",
+      location_config.namespace, "/", key_value,
+      " with estimated_final_count: ", estimated_final_count)
+
+    return ngx_exit(config.status_code)
+  end
+end
+
+return _M