Ограничение количества запросов — Raterlimiter

Если вы опасаетесь, что ваш веб-сервис могут заDOSить нерадивые пользователи, или у вас просто слабенький сервер, то вы уже задумывались над ограничением количества запросов от каждого пользователя. По-хорошему — это только один из необходимых эшелонов обороны. Конечно, от серьёзной атаки такое ограничение не убережёт, но с точки зрения цена/качество вполне подходящее

Недавно я начал активно заниматься Эрлангом. Ну и, как обычно, для закрепления материала реализовал несложный веб-сервис на Mochiweb. Mochiweb — вполне достойный фреймворк для создания веб-приложений, но возможности лимитировать количество запросов от одного клиента я не нашёл. Вот и сделал это самостоятельно.

Т.к. функционал лимитирования скорости запросов вполне себе изолированный и не привязан к какой-то конкретной задаче, я выделил сделанный модуль в независимое приложение и решил выложить его исходный код.

Задача

Итак, имеем Erlang/OTP, Mochiweb, rebar. Хочется считать количество запросов от конкретного пользователя и отдавать ему 413 код ошибки, если запросы идут слишком часто. Клиент идентифицируется своим IP адресом. Тем самым, который отдает mochiweb_request:get(peer).

Задача не такая сложная, но, возможно, готовое решение сэкономит кому-то время.

Подход

Я решил использовать алгоритм Token bucket.
Если описать его кратко — считаем определенные интервалы времени для каждого клиента корзинами, в которые будем складывать запросы пользователя за этот интервал времени. Корзины имеют лимитированный объём. Как только корзина забивается запросами до лимита, перестаём облуживать клиента. Т.к. время идет, корзины для клиента постоянно меняются, и клиент имеет шанс быть обслуженным в новый интервал времени.


На картинке выше красным обозначено время, когда каждый из клиентов не обслуживался.
Как видно из картинки, Клиент А какое-то время назад получил ограничение запросов и больше их не слал, Клиент Б спокойно работает и не превышает лимита, а клиент В выбрал на текущий момент лимит запросов и получает отказы.

Корзины будем создавать по мере появления запросов от клиента и с учетом времени. Каждый клиент может иметь свой масштаб времени, читай время жизни корзины. Объем корзины, то есть лимит запросов, можно задавать для каждого клиента отдельно.

Хранить информацию о корзинах клиентов будем в словаре:
IP + № корзины в качестве ключа
counter — собственно, счетчик запросов

API для клиента состоит из одной функции check_rate(IP, TimeScale, Limit).
Где
IP — это ip адрес клиента (впрочем может быть любым идентификатором)
TimeScale — Время жизни корзины в мс.
Limit — максимальное количество запросов в корзине

Вызов этой функции инкрементирует счетчик и возвращает клиенту надо ли обслуживать запрос.

Реализация

Для хранения данных я использую ETS таблицу типа ordered_set. Она отсортирована по ключам и не допускает дублирования ключей. В таблице лежит запись

[BucketID, Counter, CreatedAt, ModifiedAt]

Где
BucketID — {IP, BucketNum}, также является ключем для ordered_set
Counter — количество запросов от клиента
CreatedAt — Время создания корзины (в мс)
ModifiedAt — Время изменения корзины (в мс)
Время создания служит больше для отладки, от него можно отказаться

Основная функция:

-spec count_hit(Id::binary(), Scale::integer(), Limit::integer()) -> {ok, continue} | {fail, Count::integer()}.
%% Counts request by ID and blocks it if rate limiter fires
%% ID of the client
%% Scale of time (1000 = new bucket every second, 60000 = bucket every minute)
%% Limit - max size of bucket
count_hit(Id, Scale, Limit) ->
  Stamp = timestamp(),                    %% milliseconds since 00:00 GMT, January 1, 1970
  BucketNumber = trunc(Stamp / Scale),    %% with scale=1 bucket changes every millisecond
  Key = {BucketNumber, Id},
  case   ets:member(?RATERLIMITER_TABLE, Key) of
    false ->
      ets:insert(?RATERLIMITER_TABLE, {Key, 1, Stamp, Stamp }),
      {ok, continue};
    true ->
      % increment counter by 1, created_time by 0, and changed_time by current one
      Counters = ets:update_counter(?RATERLIMITER_TABLE, Key, [{2,1},{3,0},{4,1,0, Stamp}]),
      % Counter, created at, changed at
      [BucketSize, _, _] = Counters,
      if
        (BucketSize > Limit) ->
                {fail, Limit};
        true ->
                {ok, continue}
      end
    end.

Сначала проверяем есть ли нужная нам корзина {IP, BucketNumber}.
В случае новой корзины все довольно банально — создаем новую корзину с начальными значениями и говорим ОК.

Добавление счетчика к существующей корзине осуществляется чуть интереснее. Модуль ets позволяет модифицировать счетчики по указанным правилам. Мне понравилось, как изящно можно совместить инкремент счетчика с изменением даты последнего обращения.

Итак код:

      Counters = ets:update_counter(?RATERLIMITER_TABLE, Key, [{2,1},{3,0},{4,1,0, Stamp}])

инкрементирует счетчик №2 на 1 (это как раз счетчик обращений), добавляет 0 ко времени создания, добавляет 1 к счетчику №4 с проверкой максимального значения. Т.к. максимальное значение указано 0, четвертый счетчик всегда устанавливается равным Stamp (текущему времени). На выходе получаем список из всех изменяемых счетчиков по порядку, т.е.

[Counter, CreatedTime, ChangedTime]

В итоге за два обращения к таблицам мы обновили либо создали корзину и получили на выходе количество обращений. В текущей реализации raterlimiter атомарность изменения таблиц не критична, он просто работает синхронно (вызов gen_server:call/2 синхронен).

Удаление старых корзин

Со временем корзины клиентов всё копятся и копятся. Надо их удалить.

Удаление происходит периодически. Удаляются просто все корзины старее указанного лимита.

Надо сказать, на Эрланге всё выглядит красиво и кратко —

remove_old_limiters(Timeout) ->
  NowStamp = timestamp(),
  Matcher = ets:fun2ms(fun ({_,_,_,Accessed}) when Accessed < (NowStamp - Timeout) -> true end),
  ets:select_delete(?RATERLIMITER_TABLE, Matcher).

ets:select_delete удаляет все записи, удовлетворяющие функции соответствия (MatchSpec). Писать эти match_spec вручную — сущий ад. Гораздо проще использовать функцию ets:fun2ms, которая преобразует обычную функцию Эрланга в функцию соответствия. Например, указанная выше функция

fun ({_,_,_,Accessed}) when Accessed < (NowStamp - Timeout) -> true end

в конечном виде в формате MatchSpec, для NowStamp=1000, Timeout=0 выглядит как

[{{'_','_','_','$1'},[{'<','$1',{'-',1000,0}}],[true]}]

Главное, при использовании ets:fun2ms не забыть включить в исходник дополнительный файл —

-include_lib("stdlib/include/ms_transform.hrl").

Использование

Использование сводится к одному вызову check_rate, не считая, конечно, кода и настроек для запуска самого приложения raterlimiter. Вот пример из моего приложения на базе Mochiweb:

        {Status, _Reason} = raterlimiter:check_rate(Req:get(peer), 30000, 500),  % 500 requests per 30 seconds allowed
        case Status of
          ok ->
            serve_request(Req, DocRoot, Path);
          _ -> % client wants too much information
            Req:respond({413,[{"Content-Type", "text/html"}, {"Retry-After", "900 (Too many requests)"}], "С вашего IP слишком много запросов, подождите 15 минут."})
        end

С полным кодом можно ознакомиться на Github. Там же можно сообщить об ошибках, или сделать pull request.