Sindbad~EG File Manager

--[[
    The file will contains the code that counts the requests / ip and requests / domain.

    Worry Points : 
    - Call to the get_keys() is blocking. We'll need to see if it's a problem later and if it scale with the nb of ipairs
    - Does it scale with the number of IPs in the memory_zone ? Maybe we should delete IP from here when TmpBanned
    - With a lot of workers, does the Nginx workers blocked themself because of the lock on the check() call ? 
      Maybe we should switch to a local lua counter in a table (L3) + regular sync to the L2 dict using a timer
      It may be a problem on an ipxtender server... 
    - TTL management of entries. I'm not 100% sure about the current implementation.  
--]]

local _M = {}

local timer = require("resty.timer")
local o2debug = require "lib/o2switch_debug"
local o2config = require "lib/o2switch_config"
local ipUtils = require "lib/o2switch_ip_utils"
local ngx = ngx
local domainCounter = ngx.shared.domain_req_counter_dict
local ipCounter = ngx.shared.ip_req_counter_dict
local miscDict = ngx.shared.misc_dict
-- local sha1_bin = ngx.sha1_bin 
local type = type
local ipairs = ipairs
local pairs = pairs
local config = o2config.autoMitigationConfig
local tonumber = tonumber
local math_abs = math.abs
local math_floor = math.floor 
local sub = string.sub


-- Increment the counter type for the given key
-- @counterType either 'd' for domain or 'i' for ip address
-- @key either the binary ip address from nginx, or the domain name (full, untruncated)
-- @increment the increment value, default to 1
-- @return {newValue|nil, errMsg|nil}
function _M.increment(counterType, key, increment)
    local timeframe = config[counterType][1] + 5
    if counterType == 'i' then
        local v, r = ipCounter:incr(
            key, 
            (increment or 1), 
            0, 
            timeframe
        )
        ipCounter:expire(key, timeframe)
        return v, r
    else
        local v, r = domainCounter:incr(
            sub(key, 1, 30), -- sha1_bin(key), -- sha1 is better but not usable if we want to export the data from memory
            (increment or 1), 
            0, 
            timeframe
        )
        domainCounter:expire(key, timeframe)
        return v, r
    end
    return nil, 'Unexpected error'
end


-- Check if we need a mitigation for the type&key
-- @counterType 'i' for ip and 'd' for domain
-- @key The raw key, so not encoded.So binary IP address or full untruncated domain name
-- @return nil or the mitigation to use, number or nil
function _M.check(counterType, key)
    -- --o2debug.debug('check called')
    local c = config[counterType]
    local value, err = _M.increment(counterType, key)
    local mitigation = nil
    for i = 1, #c[4] do
        if c[4][i][1] > value then
            return mitigation, value
        else
            mitigation = c[4][i][2]
        end
    end

    return mitigation, value
end

-- Decrement the counter type/key. Also delete the key from the dict if number < 0 and set a max value to number to avoid overflow
-- @counterType 'i' for ip or 'd' for domain
-- @key The already encoded key. So the binary IP address and for domain, the sub(domain, 1, 30)
-- @decrement The decrement value, default to 1
-- @allowNegative Set to true to allow setting a negative value on the counter. Useful for whitelisting for instance.
-- @return {newValue|nil, errMsg|nil}
function _M.decrement(counterType, key, decrement, allowNegative)
    local decrement = (decrement > 0) and decrement * -1 or decrement
    local dict = counterType == 'i' and ipCounter or domainCounter
    local leakagePeriod = config[counterType][2] + 5
    local allowNegative = allowNegative or false

    local value, err = dict:incr(key, (decrement or -1), 0, leakagePeriod)

    if(type(value) ~= 'number') then
        dict:delete(key)
        return nil, err or 'Not a number'
    end

    -- For IP (i), the negative value can be legitimate, for instance, when a challenge is passed
    if(not allowNegative and value < 0) then
        dict:delete(key)
        return nil, nil
    end

    -- Not sure about this one. Only keep negative (whitelist) value, to keep the "whiltelist" until it re-become > 0
    if(value < 0) then
        dict:expire(key, leakagePeriod * 2)
    end

    return value, err
end

-- Return all the keys in the shared dict used to store all the counting stuff
-- @counterType either 'i' or 'd'
-- @return table with the list for keys found
function _M.getKeys(counterType) 
    local dict = counterType == 'i' and ipCounter or domainCounter
    return dict:get_keys(0)
end

-- Leak the 'counterType' counter with the right amount (based on the configuration). Also flush expired counter.
-- @counterType either 'i' or 'd'
-- @return bool, nb of keys in counter|nil, errMsg|nil
function _M.leak(counterType) 
    --o2debug.debug('Leak called for : ' .. counterType)
    --ngx.update_time()
    --local begin = ngx.now()

    local dict = counterType == 'i' and ipCounter or domainCounter
    local c = config[counterType]
    if (type(c) ~= 'table') then
        return false, nil, "Cant find the configuration"
    end
    
    local leakageCount = miscDict:get(counterType) or 0
    local leakageAmount = c[3]
    local leakageInterval = c[2]
    local reset = false
    if(type(leakageAmount) ~= 'number') then
        return false, nil, 'No leakage amount'
    end

    -- For every cycle (every timeframe), we reset the value of the count to a max value, 
    -- to avoid having a domain / ip blocked too long, if a attack has stopped.
    if((leakageCount or 0) > math_floor(c[1] / leakageInterval)) then
        reset = true
        --o2debug.debug('Reset loop this time')
        if(leakageCount > 900000) then
            miscDict:set(counterType, 0)
        end
    end

    local getResetValue = function(currentVal)
        local prev = 0
        local val = 0
        for i = 1, #c[4] do
            if currentVal > c[4][i][1] and currentVal > prev then
                val = c[4][i][1]
            else
                return (val - leakageAmount) > 0 and (val - leakageAmount) or val 
            end
            prev = c[4][i][1]
        end
        return (val - leakageAmount) > 0 and (val - leakageAmount) or val  
    end

    dict:flush_expired(0)

    local count = 1
    for _, key in ipairs(dict:get_keys(0)) do
        -- Peak value first. Not great (performance wise)
        local val, err = dict:get(key)
        if err then
            dict:delete(key)
            goto zcontinue
        end

        -- Dont touch negative value (because it came from a captcha/js challenge probably)
        if val < 0 then
            dict:expire(key, 86400) -- We'll need to see if it's a problem having a huge TTL and the IP piling up in the memory
            goto zcontinue
        end

        -- Reset : each timeframe circle, we reset the value
        if reset then
            local newValue = getResetValue(val)
            if newValue <= 0 then
                dict:delete(key)
            else
                dict:set(key, newValue)
            end
            goto zcontinue
        end

        -- Normal cycle
        _M.decrement(counterType, key, leakageAmount)

        ::zcontinue::
        count = count + 1
    end

    --ngx.update_time()
    ----o2debug.debug(tostring(leakageCount) .. ' : leakage elapsed seconds for "' .. counterType .. '" ' .. tostring(count) .. ' keys : ' .. tostring(ngx.now() - begin))
    miscDict:incr(counterType, 1, 0)

    return true, count, nil
end

-- This will registrer all the timers for the counter management. The timers will be launched only in one worker (node-wide).
-- @return bool
function _M.registerTimers() 
    --o2debug.debug('o2switch_counter.registerTimers() called')
    for i, v in pairs(o2config.autoMitigationConfig) do
        -- https://kong.github.io/lua-resty-timer/topics/README.md.html
        timer({
            interval = v[2],  
            recurring = true, 
            immediate = false,
            detached = true, 
            jitter = 0.1,
            shm_name = "timer_shm", -- shm to use for node-wide timers
            key_name = 'counter-' .. i .. '-leakage', -- key-name to use for node-wide timers
            expire = function(counterType) 
                local o2counter = require "lib/o2switch_counter"
                o2counter.leak(counterType)
            end
        }, i)        
    end
    return true
end

-- Export data from a counter
--  @return list|false, errorMsg|nil
function _M.exportDataFromCounter(counterType, filterFunction)
    if filterFunction == nil then
        -- Default filter function (Not really meant to be used, it's more of an example/documentation)
        filterFunction = function(count) 
            return count > 1000
        end
    end

    if type(filterFunction) ~= 'function' then
        return false, 'The filterFunction argument must be LUA Function'
    end

    local dict = counterType == 'i' and ipCounter or domainCounter
    local keys = _M.getKeys(counterType)
    local result = {}

    for _, k in ipairs(keys) do
        local value = dict:get(k)
        if counterType == 'i' and type(value) ~= 'number' then
            dict:delete(k)
            value = -1
        end
        if filterFunction(value) then
            local kk = counterType == 'i' and ipUtils.binaryToStringIp(k) or k
            result[kk] = value
        end
    end

    return result, nil
end

return _M