diff --git a/app/controllers/application_controller.rb b/app/controllers/application_controller.rb index 5b36a3175..8c8237bdd 100644 --- a/app/controllers/application_controller.rb +++ b/app/controllers/application_controller.rb @@ -2,8 +2,6 @@ class ApplicationController < ActionController::Base include Pundit helper_method :search_params - class ApiLimitError < StandardError; end - self.responder = ApplicationResponder skip_forgery_protection if: -> { SessionLoader.new(request).has_api_authentication? } @@ -74,17 +72,16 @@ class ApplicationController < ActionController::Base def api_check return if CurrentUser.is_anonymous? || request.get? || request.head? - if CurrentUser.user.token_bucket.nil? - TokenBucket.create_default(CurrentUser.user) - CurrentUser.user.reload - end + rate_limiter = RateLimiter.new( + "write", + [CurrentUser.user.cache_key], + cost: 1, + rate: CurrentUser.user.api_regen_multiplier, + burst: CurrentUser.user.api_burst_limit + ) - throttled = CurrentUser.user.token_bucket.throttled? - headers["X-Api-Limit"] = CurrentUser.user.token_bucket.token_count.to_s - - if throttled - raise ApiLimitError, "too many requests" - end + headers["X-Rate-Limit"] = rate_limiter.to_json + rate_limiter.limit! end def rescue_exception(exception) @@ -113,7 +110,7 @@ class ApplicationController < ActionController::Base render_error_page(410, exception, template: "static/pagination_error", message: "You cannot go beyond page #{CurrentUser.user.page_limit}.") when Post::SearchError render_error_page(422, exception, template: "static/tag_limit_error", message: "You cannot search for more than #{CurrentUser.tag_query_limit} tags at a time.") - when ApiLimitError + when RateLimiter::RateLimitError render_error_page(429, exception) when NotImplementedError render_error_page(501, exception, message: "This feature isn't available: #{exception.message}") diff --git a/app/logical/danbooru_maintenance.rb b/app/logical/danbooru_maintenance.rb index c990dcfae..48acb7c0e 100644 --- a/app/logical/danbooru_maintenance.rb +++ b/app/logical/danbooru_maintenance.rb @@ -5,13 +5,13 @@ module DanbooruMaintenance safely { Upload.prune! } safely { PostPruner.prune! } safely { PostAppealForumUpdater.update_forum! } + safely { RateLimit.prune! } safely { regenerate_post_counts! } end def daily safely { Delayed::Job.where('created_at < ?', 45.days.ago).delete_all } safely { PostDisapproval.prune! } - safely { TokenBucket.prune! } safely { BulkUpdateRequestPruner.warn_old } safely { BulkUpdateRequestPruner.reject_expired } safely { Ban.prune! } diff --git a/app/logical/rate_limiter.rb b/app/logical/rate_limiter.rb new file mode 100644 index 000000000..e8fc52430 --- /dev/null +++ b/app/logical/rate_limiter.rb @@ -0,0 +1,30 @@ +class RateLimiter + class RateLimitError < StandardError; end + + attr_reader :action, :keys, :cost, :rate, :burst + + def initialize(action, keys = ["*"], cost: 1, rate: 1, burst: 1) + @action = action + @keys = keys + @cost = cost + @rate = rate + @burst = burst + end + + def limit! + raise RateLimitError if limited? + end + + def limited? + rate_limits.any?(&:limited?) + end + + def as_json(options = {}) + hash = rate_limits.map { |limit| [limit.key, limit.points] }.to_h + super(options).except("keys", "rate_limits").merge(limits: hash) + end + + def rate_limits + @rate_limits ||= RateLimit.create_or_update!(action: action, keys: keys, cost: cost, rate: rate, burst: burst) + end +end diff --git a/app/models/rate_limit.rb b/app/models/rate_limit.rb new file mode 100644 index 000000000..3236f1ac4 --- /dev/null +++ b/app/models/rate_limit.rb @@ -0,0 +1,53 @@ +class RateLimit < ApplicationRecord + scope :expired, -> { where("updated_at < ?", 1.hour.ago) } + + def self.prune! + expired.delete_all + end + + # `action` is the action being limited. Usually a controller endpoint. + # `keys` is who is being limited. Usually a [user, ip] pair, meaning the action is limited both by the user's ID and their IP. + # `cost` is the number of points the action costs. + # `rate` is the number of points per second that are refilled. + # `burst` is the maximum number of points that can be saved up. + def self.create_or_update!(action:, keys:, cost:, rate:, burst:) + # { key0: keys[0], ..., keyN: keys[N] } + key_params = keys.map.with_index { |key, i| [:"key#{i}", key] }.to_h + + # (created_at, updated_at, action, keyN, points) + values = keys.map.with_index { |key, i| "(:now, :now, :action, :key#{i}, :points)" } + + # Do an upsert, creating a new rate limit object for each key that doesn't + # already exist, and updating the limit for each limit that already exists. + # + # https://www.postgresql.org/docs/current/sql-insert.html#SQL-ON-CONFLICT + sql = <<~SQL + INSERT INTO rate_limits (created_at, updated_at, action, key, points) + VALUES #{values.join(", ")} + ON CONFLICT (action, key) DO UPDATE SET + updated_at = :now, + limited = rate_limits.points + :rate * EXTRACT(epoch FROM (:now - rate_limits.updated_at)) < 0, + points = + CASE + WHEN rate_limits.points + :rate * EXTRACT(epoch FROM (:now - rate_limits.updated_at)) < 0 THEN + LEAST(:burst, rate_limits.points + :rate * EXTRACT(epoch FROM (:now - rate_limits.updated_at))) + ELSE + LEAST(:burst, rate_limits.points + :rate * EXTRACT(epoch FROM (:now - rate_limits.updated_at))) - :cost + END + RETURNING * + SQL + + sql_params = { + now: Time.zone.now, + action: action, + rate: rate, + burst: burst, + cost: cost, + points: burst - cost, + **key_params + } + + rate_limits = RateLimit.find_by_sql([sql, sql_params]) + rate_limits + end +end diff --git a/app/models/token_bucket.rb b/app/models/token_bucket.rb deleted file mode 100644 index 96ddf74ab..000000000 --- a/app/models/token_bucket.rb +++ /dev/null @@ -1,41 +0,0 @@ -class TokenBucket < ApplicationRecord - self.primary_key = "user_id" - belongs_to :user - - def self.prune! - where("last_touched_at < ?", 1.day.ago).delete_all - end - - def self.create_default(user) - TokenBucket.create(user_id: user.id, token_count: user.api_burst_limit, last_touched_at: Time.now) - end - - def accept? - token_count >= 1 - end - - def add! - now = Time.now - TokenBucket.where(user_id: user_id).update_all(["token_count = least(token_count + (? * extract(epoch from ? - last_touched_at)), ?), last_touched_at = ?", user.api_regen_multiplier, now, user.api_burst_limit, now]) - - # estimate the token count to avoid reloading - self.token_count += user.api_regen_multiplier * (now - last_touched_at) - self.token_count = user.api_burst_limit if token_count > user.api_burst_limit - end - - def consume! - TokenBucket.where(user_id: user_id).update_all("token_count = greatest(0, token_count - 1)") - self.token_count -= 1 - end - - def throttled? - add! - - if accept? - consume! - return false - else - return true - end - end -end diff --git a/app/models/user.rb b/app/models/user.rb index d8ad0b507..22dfa82b1 100644 --- a/app/models/user.rb +++ b/app/models/user.rb @@ -134,7 +134,6 @@ class User < ApplicationRecord has_many :user_events, dependent: :destroy has_one :recent_ban, -> {order("bans.id desc")}, :class_name => "Ban" - has_one :token_bucket has_one :email_address, dependent: :destroy has_many :api_keys, dependent: :destroy has_many :note_versions, :foreign_key => "updater_id" @@ -539,10 +538,6 @@ class User < ApplicationRecord User.api_burst_limit(level) end - def remaining_api_limit - token_bucket.try(:token_count) || api_burst_limit - end - def statement_timeout User.statement_timeout(level) end diff --git a/app/policies/user_policy.rb b/app/policies/user_policy.rb index faf36bec7..b5bef3ba9 100644 --- a/app/policies/user_policy.rb +++ b/app/policies/user_policy.rb @@ -62,7 +62,7 @@ class UserPolicy < ApplicationPolicy comment_threshold default_image_size favorite_tags blacklisted_tags time_zone per_page custom_style favorite_count api_regen_multiplier - api_burst_limit remaining_api_limit statement_timeout + api_burst_limit statement_timeout favorite_group_limit favorite_limit tag_query_limit max_saved_searches theme ] diff --git a/app/views/users/_statistics.html.erb b/app/views/users/_statistics.html.erb index f72744ad8..cc7625f80 100644 --- a/app/views/users/_statistics.html.erb +++ b/app/views/users/_statistics.html.erb @@ -265,14 +265,6 @@ (<%= link_to_wiki "help", "help:api" %>) - - - API Limits - - <%= CurrentUser.user.remaining_api_limit %> - / <%= CurrentUser.user.api_burst_limit %> (may not be up to date) - - <% end %> diff --git a/db/migrate/20170106012138_create_token_buckets.rb b/db/migrate/20170106012138_create_token_buckets.rb index a41a7067a..0b442ae6f 100644 --- a/db/migrate/20170106012138_create_token_buckets.rb +++ b/db/migrate/20170106012138_create_token_buckets.rb @@ -5,6 +5,6 @@ class CreateTokenBuckets < ActiveRecord::Migration[4.2] end def down - raise NotImplementedError + drop_table :token_buckets end end diff --git a/db/migrate/20210303195217_replace_token_buckets_with_rate_limits.rb b/db/migrate/20210303195217_replace_token_buckets_with_rate_limits.rb new file mode 100644 index 000000000..7002d5a75 --- /dev/null +++ b/db/migrate/20210303195217_replace_token_buckets_with_rate_limits.rb @@ -0,0 +1,24 @@ +require_relative "20170106012138_create_token_buckets" + +class ReplaceTokenBucketsWithRateLimits < ActiveRecord::Migration[6.1] + def change + revert CreateTokenBuckets + + create_table :rate_limits do |t| + t.timestamps null: false + t.boolean :limited, null: false, default: false + t.float :points, null: false + t.string :action, null: false + t.string :key, null: false + + t.index [:key, :action], unique: true + end + + reversible do |dir| + dir.up do + execute "ALTER TABLE rate_limits SET UNLOGGED" + execute "ALTER TABLE rate_limits SET (fillfactor = 50)" + end + end + end +end diff --git a/db/structure.sql b/db/structure.sql index e7549407b..13c02f9c8 100644 --- a/db/structure.sql +++ b/db/structure.sql @@ -2929,6 +2929,41 @@ CREATE SEQUENCE public.posts_id_seq ALTER SEQUENCE public.posts_id_seq OWNED BY public.posts.id; +-- +-- Name: rate_limits; Type: TABLE; Schema: public; Owner: - +-- + +CREATE UNLOGGED TABLE public.rate_limits ( + id bigint NOT NULL, + created_at timestamp(6) without time zone NOT NULL, + updated_at timestamp(6) without time zone NOT NULL, + limited boolean DEFAULT false NOT NULL, + points double precision NOT NULL, + action character varying NOT NULL, + key character varying NOT NULL +) +WITH (fillfactor='50'); + + +-- +-- Name: rate_limits_id_seq; Type: SEQUENCE; Schema: public; Owner: - +-- + +CREATE SEQUENCE public.rate_limits_id_seq + START WITH 1 + INCREMENT BY 1 + NO MINVALUE + NO MAXVALUE + CACHE 1; + + +-- +-- Name: rate_limits_id_seq; Type: SEQUENCE OWNED BY; Schema: public; Owner: - +-- + +ALTER SEQUENCE public.rate_limits_id_seq OWNED BY public.rate_limits.id; + + -- -- Name: saved_searches; Type: TABLE; Schema: public; Owner: - -- @@ -3085,17 +3120,6 @@ CREATE SEQUENCE public.tags_id_seq ALTER SEQUENCE public.tags_id_seq OWNED BY public.tags.id; --- --- Name: token_buckets; Type: TABLE; Schema: public; Owner: - --- - -CREATE UNLOGGED TABLE public.token_buckets ( - user_id integer, - last_touched_at timestamp without time zone NOT NULL, - token_count real NOT NULL -); - - -- -- Name: uploads; Type: TABLE; Schema: public; Owner: - -- @@ -4352,6 +4376,13 @@ ALTER TABLE ONLY public.post_votes ALTER COLUMN id SET DEFAULT nextval('public.p ALTER TABLE ONLY public.posts ALTER COLUMN id SET DEFAULT nextval('public.posts_id_seq'::regclass); +-- +-- Name: rate_limits id; Type: DEFAULT; Schema: public; Owner: - +-- + +ALTER TABLE ONLY public.rate_limits ALTER COLUMN id SET DEFAULT nextval('public.rate_limits_id_seq'::regclass); + + -- -- Name: saved_searches id; Type: DEFAULT; Schema: public; Owner: - -- @@ -4739,6 +4770,14 @@ ALTER TABLE ONLY public.posts ADD CONSTRAINT posts_pkey PRIMARY KEY (id); +-- +-- Name: rate_limits rate_limits_pkey; Type: CONSTRAINT; Schema: public; Owner: - +-- + +ALTER TABLE ONLY public.rate_limits + ADD CONSTRAINT rate_limits_pkey PRIMARY KEY (id); + + -- -- Name: saved_searches saved_searches_pkey; Type: CONSTRAINT; Schema: public; Owner: - -- @@ -7280,6 +7319,13 @@ CREATE INDEX index_posts_on_uploader_id ON public.posts USING btree (uploader_id CREATE INDEX index_posts_on_uploader_ip_addr ON public.posts USING btree (uploader_ip_addr); +-- +-- Name: index_rate_limits_on_key_and_action; Type: INDEX; Schema: public; Owner: - +-- + +CREATE UNIQUE INDEX index_rate_limits_on_key_and_action ON public.rate_limits USING btree (key, action); + + -- -- Name: index_saved_searches_on_labels; Type: INDEX; Schema: public; Owner: - -- @@ -7385,13 +7431,6 @@ CREATE INDEX index_tags_on_name_trgm ON public.tags USING gin (name public.gin_t CREATE INDEX index_tags_on_post_count ON public.tags USING btree (post_count); --- --- Name: index_token_buckets_on_user_id; Type: INDEX; Schema: public; Owner: - --- - -CREATE UNIQUE INDEX index_token_buckets_on_user_id ON public.token_buckets USING btree (user_id); - - -- -- Name: index_uploads_on_referer_url; Type: INDEX; Schema: public; Owner: - -- @@ -7964,6 +8003,7 @@ INSERT INTO "schema_migrations" (version) VALUES ('20210127000201'), ('20210127012303'), ('20210214095121'), -('20210214101614'); +('20210214101614'), +('20210303195217'); diff --git a/test/factories/rate_limit.rb b/test/factories/rate_limit.rb new file mode 100644 index 000000000..cfb37ff3f --- /dev/null +++ b/test/factories/rate_limit.rb @@ -0,0 +1,8 @@ +FactoryBot.define do + factory(:rate_limit) do + limited { false } + points { 0 } + action { "test" } + key { "1234" } + end +end diff --git a/test/functional/application_controller_test.rb b/test/functional/application_controller_test.rb index bd96c4a24..0783431f6 100644 --- a/test/functional/application_controller_test.rb +++ b/test/functional/application_controller_test.rb @@ -227,7 +227,7 @@ class ApplicationControllerTest < ActionDispatch::IntegrationTest should "fail with a 429 error" do user = create(:user) post = create(:post, rating: "s") - TokenBucket.any_instance.stubs(:throttled?).returns(true) + RateLimit.any_instance.stubs(:limited?).returns(true) put_auth post_path(post), user, params: { post: { rating: "e" } } diff --git a/test/unit/rate_limit_test.rb b/test/unit/rate_limit_test.rb new file mode 100644 index 000000000..d120a494a --- /dev/null +++ b/test/unit/rate_limit_test.rb @@ -0,0 +1,72 @@ +require 'test_helper' + +class RateLimitTest < ActiveSupport::TestCase + context "RateLimit: " do + context "#limit! method" do + should "create a new rate limit object if none exists, or update it if it already exists" do + assert_difference("RateLimit.count", 1) do + RateLimiter.new("write", ["users/1"]).limited? + end + + assert_difference("RateLimit.count", 0) do + RateLimiter.new("write", ["users/1"]).limited? + end + + assert_difference("RateLimit.count", 1) do + RateLimiter.new("write", ["users/1", "ip/1.2.3.4"]).limited? + end + + assert_difference("RateLimit.count", 0) do + RateLimiter.new("write", ["users/1", "ip/1.2.3.4"]).limited? + end + end + + should "include the cost of the first action when initializing the limit" do + limiter = RateLimiter.new("write", ["users/1"], burst: 10, cost: 1) + assert_equal(9, limiter.rate_limits.first.points) + end + + should "be limited if the point count is negative" do + freeze_time + create(:rate_limit, action: "write", key: "users/1", points: -1) + limiter = RateLimiter.new("write", ["users/1"], cost: 1) + + assert_equal(true, limiter.limited?) + assert_equal(-1, limiter.rate_limits.first.points) + end + + should "not be limited if the point count was positive before the action" do + freeze_time + create(:rate_limit, action: "write", key: "users/1", points: 0.01) + limiter = RateLimiter.new("write", ["users/1"], cost: 1) + + assert_equal(false, limiter.limited?) + assert_equal(-0.99, limiter.rate_limits.first.points) + end + + should "refill the points at the correct rate" do + freeze_time + create(:rate_limit, action: "write", key: "users/1", points: -2) + + limiter = RateLimiter.new("write", ["users/1"], cost: 1, rate: 1, burst: 10) + assert_equal(true, limiter.limited?) + assert_equal(-2, limiter.rate_limits.first.points) + + travel 1.second + limiter = RateLimiter.new("write", ["users/1"], cost: 1, rate: 1, burst: 10) + assert_equal(true, limiter.limited?) + assert_equal(-1, limiter.rate_limits.first.points) + + travel 5.second + limiter = RateLimiter.new("write", ["users/1"], cost: 1, rate: 1, burst: 10) + assert_equal(false, limiter.limited?) + assert_equal(3, limiter.rate_limits.first.points) + + travel 60.second + limiter = RateLimiter.new("write", ["users/1"], cost: 1, rate: 1, burst: 10) + assert_equal(false, limiter.limited?) + assert_equal(9, limiter.rate_limits.first.points) + end + end + end +end diff --git a/test/unit/token_bucket_test.rb b/test/unit/token_bucket_test.rb deleted file mode 100644 index 3c96b371a..000000000 --- a/test/unit/token_bucket_test.rb +++ /dev/null @@ -1,45 +0,0 @@ -require 'test_helper' - -class TokenBucketTest < ActiveSupport::TestCase - context "#add!" do - setup do - @user = FactoryBot.create(:user) - TokenBucket.create(user_id: @user.id, last_touched_at: 1.minute.ago, token_count: 0) - end - - should "work" do - @user.token_bucket.add! - assert_operator(@user.token_bucket.token_count, :>, 0) - @user.reload - assert_operator(@user.token_bucket.token_count, :>, 0) - end - end - - context "#consume!" do - setup do - @user = FactoryBot.create(:user) - TokenBucket.create(user_id: @user.id, last_touched_at: 1.minute.ago, token_count: 1) - end - - should "work" do - @user.token_bucket.consume! - assert_operator(@user.token_bucket.token_count, :<, 1) - @user.reload - assert_operator(@user.token_bucket.token_count, :<, 1) - end - end - - context "#throttled?" do - setup do - @user = FactoryBot.create(:user) - TokenBucket.create(user_id: @user.id, last_touched_at: 1.minute.ago, token_count: 0) - end - - should "work" do - assert(!@user.token_bucket.throttled?) - assert_operator(@user.token_bucket.token_count, :<, 60) - @user.reload - assert_operator(@user.token_bucket.token_count, :<, 60) - end - end -end