plugins/mod_limits.lua

   1 -- Because we deal with pre-authed sessions and streams we can't be host-specific
   2 module:set_global();
   3
   4 local filters = require "util.filters";
   5 local throttle = require "util.throttle";
   6 local timer = require "util.timer";
   7 local ceil = math.ceil;
   8
   9 local limits_cfg = module:get_option("limits", {});
  10 local limits_resolution = module:get_option_number("limits_resolution", 1);
  11
  12 local default_bytes_per_second = 3000;
  13 local default_burst = 2;
  14
  15 local rate_units = { b = 1, k = 3, m = 6, g = 9, t = 12 } -- Plan for the future.
  16 local function parse_rate(rate, sess_type)
  17         local quantity, unit, exp;
  18         if rate then
  19                 quantity, unit = rate:match("^(%d+) ?([^/]+)/s$");
  20                 exp = quantity and rate_units[unit:sub(1,1):lower()];
  21         end
  22         if not exp then
  23                 module:log("error", "Error parsing rate for %s: %q, using default rate (%d bytes/s)", sess_type, rate, default_bytes_per_second);
  24                 return default_bytes_per_second;
  25         end
  26         return quantity*(10^exp);
  27 end
  28
  29 local function parse_burst(burst, sess_type)
  30         if type(burst) == "string" then
  31                 burst = burst:match("^(%d+) ?s$");
  32         end
  33         local n_burst = tonumber(burst);
  34         if not n_burst then
  35                 module:log("error", "Unable to parse burst for %s: %q, using default burst interval (%ds)", sess_type, burst, default_burst);
  36         end
  37         return n_burst or default_burst;
  38 end
  39
  40 -- Process config option into limits table:
  41 -- limits = { c2s = { bytes_per_second = X, burst_seconds = Y } }
  42 local limits = {};
  43
  44 for sess_type, sess_limits in pairs(limits_cfg) do
  45         limits[sess_type] = {
  46                 bytes_per_second = parse_rate(sess_limits.rate, sess_type);
  47                 burst_seconds = parse_burst(sess_limits.burst, sess_type);
  48         };
  49 end
  50
  51 local default_filter_set = {};
  52
  53 function default_filter_set.bytes_in(bytes, session)
  54         local sess_throttle = session.throttle;
  55         if sess_throttle then
  56                 local ok, balance, outstanding = sess_throttle:poll(#bytes, true);
  57                 if not ok then
  58                         session.log("debug", "Session over rate limit (%d) with %d (by %d), pausing", sess_throttle.max, #bytes, outstanding);
  59                         outstanding = ceil(outstanding);
  60                         session.conn:pause(); -- Read no more data from the connection until there is no outstanding data
  61                         local outstanding_data = bytes:sub(-outstanding);
  62                         bytes = bytes:sub(1, #bytes-outstanding);
  63                         timer.add_task(limits_resolution, function ()
  64                                 if not session.conn then return; end
  65                                 if sess_throttle:peek(#outstanding_data) then
  66                                         session.log("debug", "Resuming paused session");
  67                                         session.conn:resume();
  68                                 end
  69                                 -- Handle what we can of the outstanding data
  70                                 session.data(outstanding_data);
  71                         end);
  72                 end
  73         end
  74         return bytes;
  75 end
  76
  77 local type_filters = {
  78         c2s = default_filter_set;
  79         s2sin = default_filter_set;
  80         s2sout = default_filter_set;
  81 };
  82
  83 local function filter_hook(session)
  84         local session_type = session.type:match("^[^_]+");
  85         local filter_set, opts = type_filters[session_type], limits[session_type];
  86         if opts then
  87                 if session.conn and session.conn.setlimit then
  88                         session.conn:setlimit(opts.bytes_per_second);
  89                         -- Currently no burst support
  90                 else
  91                         session.throttle = throttle.create(opts.bytes_per_second * opts.burst_seconds, opts.burst_seconds);
  92                         filters.add_filter(session, "bytes/in", filter_set.bytes_in, 1000);
  93                 end
  94         end
  95 end
  96
  97 function module.load()
  98         filters.add_filter_hook(filter_hook);
  99 end
 100
 101 function module.unload()
 102         filters.remove_filter_hook(filter_hook);
 103 end
 104
 105 function module.add_host(module)
 106         local unlimited_jids = module:get_option_inherited_set("unlimited_jids", {});
 107
 108         if not unlimited_jids:empty() then
 109                 module:hook("authentication-success", function (event)
 110                         local session = event.session;
 111                         local session_type = session.type:match("^[^_]+");
 112                         local jid = session.username .. "@" .. session.host;
 113                         if unlimited_jids:contains(jid) then
 114                                 if session.conn and session.conn.setlimit then
 115                                         session.conn:setlimit(0);
 116                                         -- Currently no burst support
 117                                 else
 118                                         local filter_set = type_filters[session_type];
 119                                         filters.remove_filter(session, "bytes/in", filter_set.bytes_in);
 120                                         session.throttle = nil;
 121                                 end
 122                         end
 123                 end);
 124         end
 125 end