增加sql的waf功能

resty/limit/conn.lua

@@ -0,0 +1,125 @@
+-- Copyright (C) Yichun Zhang (agentzh)
+--
+-- This library is an enhanced Lua port of the standard ngx_limit_conn
+-- module.
+
+
+local math = require "math"
+
+
+local setmetatable = setmetatable
+local floor = math.floor
+local ngx_shared = ngx.shared
+local assert = assert
+
+
+local _M = {
+    _VERSION = '0.07'
+}
+
+
+local mt = {
+    __index = _M
+}
+
+
+function _M.new(dict_name, max, burst, default_conn_delay)
+    local dict = ngx_shared[dict_name]
+    if not dict then
+        return nil, "shared dict not found"
+    end
+
+    assert(max > 0 and burst >= 0 and default_conn_delay > 0)
+
+    local self = {
+        dict = dict,
+        max = max + 0,    -- just to ensure the param is good
+        burst = burst,
+        unit_delay = default_conn_delay,
+    }
+
+    return setmetatable(self, mt)
+end
+
+
+function _M.incoming(self, key, commit)
+    local dict = self.dict
+    local max = self.max
+
+    self.committed = false
+
+    local conn, err
+    if commit then
+        conn, err = dict:incr(key, 1, 0)
+        if not conn then
+            return nil, err
+        end
+
+        if conn > max + self.burst then
+            conn, err = dict:incr(key, -1)
+            if not conn then
+                return nil, err
+            end
+            return nil, "rejected"
+        end
+        self.committed = true
+
+    else
+        conn = (dict:get(key) or 0) + 1
+        if conn > max + self.burst then
+            return nil, "rejected"
+        end
+    end
+
+    if conn > max then
+        -- make the exessive connections wait
+        return self.unit_delay * floor((conn - 1) / max), conn
+    end
+
+    -- we return a 0 delay by default
+    return 0, conn
+end
+
+
+function _M.is_committed(self)
+    return self.committed
+end
+
+
+function _M.leaving(self, key, req_latency)
+    assert(key)
+    local dict = self.dict
+
+    local conn, err = dict:incr(key, -1)
+    if not conn then
+        return nil, err
+    end
+
+    if req_latency then
+        local unit_delay = self.unit_delay
+        self.unit_delay = (req_latency + unit_delay) / 2
+    end
+
+    return conn
+end
+
+
+function _M.uncommit(self, key)
+    assert(key)
+    local dict = self.dict
+
+    return dict:incr(key, -1)
+end
+
+
+function _M.set_conn(self, conn)
+    self.max = conn
+end
+
+
+function _M.set_burst(self, burst)
+    self.burst = burst
+end
+
+
+return _M

resty/limit/count.lua

@@ -0,0 +1,103 @@
+-- implement GitHub request rate limiting:
+--    https://developer.github.com/v3/#rate-limiting
+
+local ngx_shared = ngx.shared
+local setmetatable = setmetatable
+local assert = assert
+
+
+local _M = {
+   _VERSION = '0.07'
+}
+
+
+local mt = {
+    __index = _M
+}
+
+
+-- the "limit" argument controls number of request allowed in a time window.
+-- time "window" argument controls the time window in seconds.
+function _M.new(dict_name, limit, window)
+    local dict = ngx_shared[dict_name]
+    if not dict then
+        return nil, "shared dict not found"
+    end
+
+    assert(limit > 0 and window > 0)
+
+    local self = {
+        dict = dict,
+        limit = limit,
+        window = window,
+    }
+
+    return setmetatable(self, mt)
+end
+
+
+function _M.incoming(self, key, commit)
+    local dict = self.dict
+    local limit = self.limit
+    local window = self.window
+
+    local remaining, ok, err
+
+    if commit then
+        remaining, err = dict:incr(key, -1, limit)
+        if not remaining then
+            return nil, err
+        end
+
+        if remaining == limit - 1 then
+            ok, err = dict:expire(key, window)
+            if not ok then
+                if err == "not found" then
+                    remaining, err = dict:incr(key, -1, limit)
+                    if not remaining then
+                        return nil, err
+                    end
+
+                    ok, err = dict:expire(key, window)
+                    if not ok then
+                        return nil, err
+                    end
+
+                else
+                    return nil, err
+                end
+            end
+        end
+
+    else
+        remaining = (dict:get(key) or limit) - 1
+    end
+
+    if remaining < 0 then
+        return nil, "rejected"
+    end
+
+    return 0, remaining
+end
+
+
+-- uncommit remaining and return remaining value
+function _M.uncommit(self, key)
+    assert(key)
+    local dict = self.dict
+    local limit = self.limit
+
+    local remaining, err = dict:incr(key, 1)
+    if not remaining then
+        if err == "not found" then
+            remaining = limit
+        else
+            return nil, err
+        end
+    end
+
+    return remaining
+end
+
+
+return _M

resty/limit/req.lua

@@ -0,0 +1,153 @@
+-- Copyright (C) Yichun Zhang (agentzh)
+--
+-- This library is an approximate Lua port of the standard ngx_limit_req
+-- module.
+
+
+local ffi = require "ffi"
+local math = require "math"
+
+
+local ngx_shared = ngx.shared
+local ngx_now = ngx.now
+local setmetatable = setmetatable
+local ffi_cast = ffi.cast
+local ffi_str = ffi.string
+local abs = math.abs
+local tonumber = tonumber
+local type = type
+local assert = assert
+local max = math.max
+
+
+-- TODO: we could avoid the tricky FFI cdata when lua_shared_dict supports
+-- hash-typed values as in redis.
+ffi.cdef[[
+    struct lua_resty_limit_req_rec {
+        unsigned long        excess;
+        uint64_t             last;  /* time in milliseconds */
+        /* integer value, 1 corresponds to 0.001 r/s */
+    };
+]]
+local const_rec_ptr_type = ffi.typeof("const struct lua_resty_limit_req_rec*")
+local rec_size = ffi.sizeof("struct lua_resty_limit_req_rec")
+
+-- we can share the cdata here since we only need it temporarily for
+-- serialization inside the shared dict:
+local rec_cdata = ffi.new("struct lua_resty_limit_req_rec")
+
+
+local _M = {
+    _VERSION = '0.07'
+}
+
+
+local mt = {
+    __index = _M
+}
+
+
+function _M.new(dict_name, rate, burst)
+    local dict = ngx_shared[dict_name]
+    if not dict then
+        return nil, "shared dict not found"
+    end
+
+    assert(rate > 0 and burst >= 0)
+
+    local self = {
+        dict = dict,
+        rate = rate * 1000,
+        burst = burst * 1000,
+    }
+
+    return setmetatable(self, mt)
+end
+
+
+-- sees an new incoming event
+-- the "commit" argument controls whether should we record the event in shm.
+-- FIXME we have a (small) race-condition window between dict:get() and
+-- dict:set() across multiple nginx worker processes. The size of the
+-- window is proportional to the number of workers.
+function _M.incoming(self, key, commit)
+    local dict = self.dict
+    local rate = self.rate
+    local now = ngx_now() * 1000
+
+    local excess
+
+    -- it's important to anchor the string value for the read-only pointer
+    -- cdata:
+    local v = dict:get(key)
+    if v then
+        if type(v) ~= "string" or #v ~= rec_size then
+            return nil, "shdict abused by other users"
+        end
+        local rec = ffi_cast(const_rec_ptr_type, v)
+        local elapsed = now - tonumber(rec.last)
+
+        -- print("elapsed: ", elapsed, "ms")
+
+        -- we do not handle changing rate values specifically. the excess value
+        -- can get automatically adjusted by the following formula with new rate
+        -- values rather quickly anyway.
+        excess = max(tonumber(rec.excess) - rate * abs(elapsed) / 1000 + 1000,
+                     0)
+
+        -- print("excess: ", excess)
+
+        if excess > self.burst then
+            return nil, "rejected"
+        end
+
+    else
+        excess = 0
+    end
+
+    if commit then
+        rec_cdata.excess = excess
+        rec_cdata.last = now
+        dict:set(key, ffi_str(rec_cdata, rec_size))
+    end
+
+    -- return the delay in seconds, as well as excess
+    return excess / rate, excess / 1000
+end
+
+
+function _M.uncommit(self, key)
+    assert(key)
+    local dict = self.dict
+
+    local v = dict:get(key)
+    if not v then
+        return nil, "not found"
+    end
+
+    if type(v) ~= "string" or #v ~= rec_size then
+        return nil, "shdict abused by other users"
+    end
+
+    local rec = ffi_cast(const_rec_ptr_type, v)
+
+    local excess = max(tonumber(rec.excess) - 1000, 0)
+
+    rec_cdata.excess = excess
+    rec_cdata.last = rec.last
+    dict:set(key, ffi_str(rec_cdata, rec_size))
+    return true
+end
+
+
+function _M.set_rate(self, rate)
+    self.rate = rate * 1000
+end
+
+
+function _M.set_burst(self, burst)
+    self.burst = burst * 1000
+end
+
+
+return _M

resty/limit/traffic.lua

@@ -0,0 +1,58 @@
+-- Copyright (C) Yichun Zhang (agentzh)
+--
+-- This is an aggregator for various concrete traffic limiter instances
+-- (like instances of the resty.limit.req, resty.limit.count and
+-- resty.limit.conn classes).
+
+
+local max = math.max
+
+
+local _M = {
+    _VERSION = '0.07'
+}
+
+
+-- the states table is user supplied. each element stores the 2nd return value
+-- of each limiter if there is no error returned. for resty.limit.req, the state
+-- is the "excess" value (i.e., the number of excessive requests each second),
+-- and for resty.limit.conn, the state is the current concurrency level
+-- (including the current new connection).
+function _M.combine(limiters, keys, states)
+    local n = #limiters
+    local max_delay = 0
+    for i = 1, n do
+        local lim = limiters[i]
+        local delay, err = lim:incoming(keys[i], i == n)
+        if not delay then
+            return nil, err
+        end
+        if i == n then
+            if states then
+                states[i] = err
+            end
+            max_delay = delay
+        end
+    end
+    for i = 1, n - 1 do
+        local lim = limiters[i]
+        local delay, err = lim:incoming(keys[i], true)
+        if not delay then
+            for j = 1, i - 1 do
+                -- we intentionally ignore any errors returned below.
+                limiters[j]:uncommit(keys[j])
+            end
+            limiters[n]:uncommit(keys[n])
+            return nil, err
+        end
+        if states then
+            states[i] = err
+        end
+
+        max_delay = max(max_delay, delay)
+    end
+    return max_delay
+end
+
+
+return _M