Rate-Limiting-Advanced-OK Plugin for Kong

01_rate-limiting-advanced-ok.md

Reference: https://docs.konghq.com/hub/kong-inc/rate-limiting-advanced/

Overview

The rate-limiting-advanced-ok plugin is designed to provide enhanced rate limiting capabilities to services managed by Kong API Gateway. This plugin extends the basic rate limiting features by offering more granular control over rate limits based on a variety of criteria including IP, consumer, credential, and more. It supports both sliding and fixed window types for flexibility in how rate limits are enforced.

NOTE: The code was rewritten based on the original "rate-limiting-advanced" on May 27, 2024.

Features

• Multiple Identifiers: Limit traffic based on IP, consumer, credentials, and other request attributes.
• Window Types: Supports both sliding and fixed windows to manage rate limiting intervals.
• Scalable Configuration: Allows synchronization with a central data store for distributed rate limiting across multiple nodes.
• Customizable Limits: Define limits and window sizes directly in the plugin configuration.
• Namespace Isolation: Configurations can be isolated using namespaces, ensuring that settings do not conflict across different instances or deployments.

Installation

Prerequisites

• Kong API Gateway installed and running.
• Docker environment (if deploying using Docker).
• Access to the Kong configuration files or environment variables.

Step-by-Step Installation

1. Prepare Plugin Files Clone or copy the rate-limiting-advanced-ok plugin files into a directory on your system.

2. Update Docker-Compose Configuration Modify your docker-compose.yml to include the plugin in the Kong container:

services:
  kong:
    environment:
      KONG_LUA_PACKAGE_PATH: "/usr/local/share/lua/5.1/kong/plugins/rate-limiting-advanced-ok/?.lua;;"
      KONG_PLUGINS: "bundled,rate-limiting-advanced-ok"
    volumes:
      - ./rate-limiting-advanced-ok:/usr/local/share/lua/5.1/kong/plugins/rate-limiting-advanced-ok

3. Configure the Plugin In your kong.yml declarative configuration file, add the plugin configuration:

plugins:
  - name: rate-limiting-advanced-ok
    config:
      identifier: "consumer"
      window_size:
        - 60    # Window size in seconds
      limit:
        - 100   # Limit for the specified window
      window_type: "sliding"
      strategy: "local"
      namespace: "unique_namespace"
      dictionary_name: "kong_rate_limiting_counters"

4. Deploy Changes Apply the changes by restarting Kong. If using Docker, you can do this by running:

docker-compose down
docker-compose up -d

5. Verify Installation After deployment, check that the plugin is active by accessing the Kong Admin API:

curl http://localhost:8001/plugins/enabled

This endpoint should list rate-limiting-advanced-ok as an enabled plugin.

Usage

After installation, the rate-limiting-advanced-ok plugin will automatically start applying rate limits as configured to the traffic passing through Kong. You can adjust the plugin settings in the kong.yml file to change limits, window sizes, or other parameters based on your requirements.

02_handler.lua

-- filename: rate-limiting-advanced-ok/handler.lua
local base_rate_limiting = require("kong.tools.public.rate-limiting").new("rate-limiting-advanced")
local schema = require("kong.plugins.rate-limiting-advanced.schema")
local event_hooks = require("kong.enterprise_edition.event_hooks")
local consumer_groups_helpers = require("kong.enterprise_edition.consumer_groups_helpers")
local kong_meta = require("kong.meta")
local uuid = require("kong.tools.uuid")
local ngx_var = ngx
local null = ngx_var.null
local kong = kong
local math_ceil = math.ceil
local math_floor = math.floor
local math_max = math.max
local math_min = math.min
local math_random = math.random
local ngx_time = ngx_var.time
local pcall_safe = pcall
local pairs_iterator = pairs
local ipairs_iterator = ipairs
local tonumber_safe = tonumber

local plugin = {
    PRIORITY = 910,
    VERSION = kong_meta.core_version,
}

local headers = {
    limit = "X-RateLimit-Limit",
    remaining = "X-RateLimit-Remaining",
    legacy_limit = "RateLimit-Limit",
    legacy_remaining = "RateLimit-Remaining",
    reset = "RateLimit-Reset",
    retry_after = "Retry-After"
}

local time_units = {
    [60] = "minute",
    [3600] = "hour",
    [31536000] = "year",
    [2592000] = "month",
    [86400] = "day",
    "second"
}

local identifiers = {
    ip = function()
        return kong.client.get_forwarded_ip()
    end,
    credential = function()
        local credential = kong.client.get_credential()
        return credential and credential.id or nil
    end,
    consumer = function()
        local consumer = kong.client.get_consumer()
        return consumer and consumer.id or nil
    end,
    service = function()
        local service = kong.router.get_service()
        return service and service.id or nil
    end,
    header = function(args)
        return kong.request.get_header(args.header_name)
    end,
    path = function(args)
        local path = kong.request.get_path()
        return path == args.path and args.path or false
    end,
    ["consumer-group"] = function(args)
        local consumer_group_id = args.consumer_group_id
        if not consumer_group_id then
            return nil
        end
        for _, group in ipairs_iterator(kong.client.get_consumer_groups()) do
            if group.id == consumer_group_id then
                return group.id
            end
        end
        return nil
    end,
}

local function setup_timer(args)
    local sync_rate = args.sync_rate
    local namespace = args.namespace
    local timer_id = uuid()
    local current_time = ngx_time()
    local initial_sync = sync_rate - current_time - math_floor(current_time / sync_rate) * sync_rate
    kong.log.debug("creating timer for namespace ", namespace, ", timer_id: ", timer_id, ", initial sync in ", initial_sync, " seconds")
    ngx_var.timer.at(initial_sync, base_rate_limiting.sync, base_rate_limiting, namespace, timer_id)
    base_rate_limiting.config[namespace].timer_id = timer_id
    base_rate_limiting.fetch(nil, base_rate_limiting, namespace, current_time, math_min(sync_rate - 0.001, 2), true)
end

local function add_namespace(args, previous_timer_id)
    if not args then
        kong.log.warn("[rate-limiting-advanced] no config was specified.", " Skipping the namespace creation.")
        return false
    end
    kong.log.debug("attempting to add namespace ", args.namespace)
    local success, error_message = pcall_safe(function()
        local strategy = args.strategy or "redis"
        local strategy_options = nil
        if strategy == "redis" then
            strategy_options = args.redis
        elseif strategy == "local" then
            args.sync_rate = -1
        end
        local dictionary_name = args.dictionary_name or schema.fields.dictionary_name.default
        if not dictionary_name then
            kong.log.warn("[rate-limiting-advanced] no shared dictionary was specified.", " Trying the default value '", dictionary_name, "'...")
        end
        if ngx_var.shared[dictionary_name] == nil then
            kong.log.notice("[rate-limiting-advanced] specified shared dictionary '", dictionary_name, "' doesn't exist. Falling back to the 'kong' shared dictionary")
            dictionary_name = "kong"
        end
        kong.log.notice("[rate-limiting-advanced] using shared dictionary '" .. dictionary_name .. "'")
        base_rate_limiting:new_namespace({
            namespace = args.namespace,
            sync_rate = args.sync_rate,
            strategy = strategy,
            strategy_opts = strategy_options,
            dict = dictionary_name,
            window_sizes = args.window_size,
            db = kong.db,
            timer_id = previous_timer_id,
        })
    end)
    if not success then
        kong.log.err("Error in creating new ratelimit namespace: ", error_message)
        return false
    end
    return true
end

function plugin.init_worker()
    event_hooks.publish("rate-limiting-advanced", "rate-limit-exceeded", {
        description = "Run an event when a rate limit has been exceeded",
        fields = {
            "consumer",
            "ip",
            "service",
            "rate",
            "limit",
            "window"
        },
        unique = {
            "consumer",
            "ip",
            "service"
        },
    })
end

function plugin.configure(args, previous_config)
    local new_config = {}
    if previous_config then
        for _, config in ipairs_iterator(previous_config) do
            local namespace = config.namespace
            local sync_rate = config.sync_rate or null
            new_config[namespace] = true
            kong.log.debug("clear and reset ", namespace)
            if not base_rate_limiting.config[namespace] then
                add_namespace(config)
            else
                local current_timer_id = nil
                if sync_rate > 0 then
                    current_timer_id = base_rate_limiting.config[namespace].timer_id
                end
                base_rate_limiting:clear_config(namespace)
                add_namespace(config, current_timer_id)
                if sync_rate > 0 and sync_rate < 1 then
                    kong.log.warn("Config option 'sync_rate' ", sync_rate, " is between 0 and 1; a config update is recommended")
                end
            end
        end
    end
    for namespace in pairs_iterator(base_rate_limiting.config) do
        if not new_config[namespace] then
            kong.log.debug("clearing old namespace ", namespace)
            base_rate_limiting.config[namespace].kill = true
            base_rate_limiting.config[namespace].timer_id = nil
        end
    end
end

function plugin.access(args)
    local namespace = args.namespace
    local current_time = ngx_time()
    local identifier = identifiers[args.identifier](args)
    if not identifier then
        identifier = identifiers.ip()
    end
    if not base_rate_limiting.config[namespace] then
        add_namespace(args)
    end
    if args.sync_rate > 0 and not base_rate_limiting.config[namespace].timer_id then
        setup_timer(args)
    end
    local was_limited = nil
    if args.enforce_consumer_groups and kong.client.get_consumer() and args.consumer_groups then
        while true do
            local consumer = kong.client.get_consumer()
            for _, group_name in ipairs_iterator(args.consumer_groups) do
                local group = consumer_groups_helpers.get_consumer_group(group_name)
                if group and consumer_groups_helpers.is_consumer_in_group(consumer.id, group.id) then
                    local group_config = consumer_groups_helpers.get_consumer_group_config(group.id, "rate-limiting-advanced")
                    if group_config then
                        was_limited = group_config.config
                        break
                    else
                        kong.log.warn("Consumer group ", group.name, " enforced but no consumer group configurations provided. Original plugin configurations will apply.")
                        break
                    end
                end
            end
        end
    end
    if not was_limited then
        was_limited = args
    end
    local headers_to_set = {}
    local minimum_remaining = nil
    local minimum_window = nil
    local minimum_limit = nil
    for _, window_size in ipairs_iterator(was_limited.window_size) do
        local limit = tonumber_safe(was_limited.limit[window_size])
        local current_count = base_rate_limiting:increment(identifier, window_size, 1, namespace, was_limited.window_type == "fixed" and 0 or nil)
        local window_start_timestamp = math_floor(current_time / window_size) * window_size
        local window_start_key = "timestamp:" .. window_size .. ":window_start"
        if limit < current_count and was_limited.window_type == "sliding" then
            ngx_var.shared[namespace]:add(window_start_key, window_start_timestamp)
            window_start_timestamp = ngx_var.shared[namespace]:get(window_start_key)
        else
            ngx_var.shared[namespace]:delete(window_start_key)
        end
        local time_unit = time_units[window_size] or window_size
        local remaining_requests = math_max(limit - current_count, 0)
        if not args.hide_client_headers then
            headers_to_set[headers.limit .. "-" .. time_unit] = limit
            headers_to_set[headers.remaining .. "-" .. time_unit] = remaining_requests
            if not minimum_remaining or remaining_requests < minimum_remaining or remaining_requests == minimum_remaining and minimum_window < window_size then
                minimum_remaining = remaining_requests
                minimum_window = window_size
                minimum_limit = limit
                local reset_time = math_max(1, minimum_window - current_time - window_start_timestamp)
                if was_limited.window_type == "sliding" then
                    reset_time = math_ceil(reset_time + math_max(0, (current_count - minimum_limit) / minimum_limit * minimum_window))
                end
            end
        end
        if limit < current_count then
            was_limited = _
            local emit_event = event_hooks.emit
            local event_name = "rate-limiting-advanced"
            local event_type = "rate-limit-exceeded"
            local event_data = {
                consumer = kong.client.get_consumer() or {},
                ip = kong.client.get_forwarded_ip(),
                service = kong.router.get_service() or {},
                rate = current_count,
                limit = limit,
                window = time_unit
            }
            emit_event(event_name, event_type, event_data)
        end
    end
    headers_to_set[headers.legacy_limit] = minimum_limit
    headers_to_set[headers.legacy_remaining] = minimum_remaining
    headers_to_set[headers.reset] = minimum_window
    if was_limited then
        local reset_time = minimum_window
        local max_jitter = was_limited.retry_after_jitter_max
        if reset_time and max_jitter > 0 then
            reset_time = reset_time + math_random(max_jitter)
        end
        headers_to_set[headers.retry_after] = reset_time
        if args.disable_penalty and was_limited.window_type == "sliding" then
            for limit_index = 1, was_limited, 1 do
                base_rate_limiting:increment(identifier, tonumber_safe(was_limited.window_size[limit_index]), -1, namespace, 0)
            end
        end
        return kong.response.exit(args.error_code, {
            message = args.error_message,
        }, headers_to_set)
    else
        kong.response.set_headers(headers_to_set)
    end
end

return plugin

03_schema.lua

-- filename: rate-limiting-advanced-ok/schema.lua
local redis = require("kong.enterprise_edition.redis")
local typedefs = require("kong.db.schema.typedefs")
local ngx_shared = ngx

-- Function to validate if a shared dictionary exists
local function validate_shared_dict(dict_name)
    if not ngx_shared.shared[dict_name] then
        return false, "missing shared dict '" .. dict_name .. "'"
    end
    return true
end

-- Schema definition for the plugin
local schema = {
    name = "rate-limiting-advanced-ok",
    fields = {{
        config = {
            type = "record",
            fields = {
                { identifier = {
                    type = "string",
                    required = true,
                    default = "consumer",
                    one_of = { "ip", "credential", "consumer", "service", "header", "path", "consumer-group" },
                    description = "The type of identifier used to generate the rate limit key.",
                }},
                { window_size = {
                    type = "array",
                    required = true,
                    elements = { type = "number" },
                    description = "One or more window sizes to apply a limit to (defined in seconds).",
                }},
                { window_type = {
                    type = "string",
                    default = "sliding",
                    one_of = { "fixed", "sliding" },
                    description = "Sets the time window type.",
                }},
                { limit = {
                    type = "array",
                    required = true,
                    elements = { type = "number" },
                    description = "One or more requests-per-window limits to apply.",
                }},
                { sync_rate = {
                    type = "number",
                    description = "How often to sync counter data to the central data store.",
                }},
                { namespace = {
                    type = "string",
                    required = true,
                    auto = true,
                    description = "The rate limiting library namespace to use for this plugin instance.",
                }},
                { strategy = {
                    type = "string",
                    required = true,
                    default = "local",
                    one_of = { "cluster", "redis", "local" },
                    description = "The rate-limiting strategy to use.",
                }},
                { dictionary_name = {
                    type = "string",
                    required = true,
                    default = "kong_rate_limiting_counters",
                    description = "The shared dictionary where counters are stored.",
                }},
                { hide_client_headers = {
                    type = "boolean",
                    default = false,
                    description = "Optionally hide informative response headers.",
                }},
                { retry_after_jitter_max = {
                    type = "number",
                    default = 0,
                    description = "The upper bound of a jitter in seconds added to the `Retry-After` header.",
                }},
                { header_name = typedefs.header_name },
                { path = typedefs.path },
                { redis = redis.config_schema },
                { enforce_consumer_groups = {
                    type = "boolean",
                    default = false,
                    description = "Determines if consumer groups are allowed to override the settings.",
                }},
                { consumer_groups = {
                    type = "array",
                    description = "List of consumer groups allowed to override the settings.",
                    elements = { type = "string" },
                }},
                { disable_penalty = {
                    type = "boolean",
                    default = false,
                    description = "If set to `true`, does not count denied requests.",
                }},
                { error_code = {
                    type = "number",
                    gt = 0,
                    default = 429,
                    description = "Set a custom error code to return when the rate limit is exceeded.",
                }},
                { error_message = {
                    type = "string",
                    default = "API rate limit exceeded",
                    description = "Set a custom error message to return when the rate limit is exceeded.",
                }},
            },
        },
    }},
    entity_checks = {{
        custom_entity_check = {
            field_sources = { "config" },
            fn = function(args)
                local validation_status, error_message = validate_shared_dict(args.config.dictionary_name)
                if not validation_status then
                    return false, error_message
                end
                return true
            end,
        },
    }},
}

return schema