implement token bucket rate limiting

app/controllers/application_controller.rb

@@ -19,7 +19,7 @@ class ApplicationController < ActionController::Base
   rescue_from SessionLoader::AuthenticationFailure, :with => :authentication_failed
   rescue_from Danbooru::Paginator::PaginationError, :with => :render_pagination_limit
 
-protected
+  protected
   def show_moderation_notice?
     CurrentUser.can_approve_posts? && (cookies[:moderated].blank? || Time.at(cookies[:moderated].to_i) < 1.day.ago)
   end
@@ -40,19 +40,37 @@ protected
   end
   
   def api_check
-    if request.format.to_s =~ /\/json|\/xml/ || params[:controller] == "iqdb"
+    if !CurrentUser.is_anonymous? && !request.get? && !request.head?
-      if ApiLimiter.throttled?(CurrentUser.id || request.remote_ip, request.request_method)
+      if CurrentUser.user.token_bucket.nil?
-        render :text => "429 Too Many Requests\n", :layout => false, :status => 429
+        TokenBucket.create_default(CurrentUser.user)
+        CurrentUser.user.reload
+      end
+
+      throttled = CurrentUser.user.token_bucket.throttled?
+      headers["X-Api-Limit"] = CurrentUser.user.token_bucket.token_count
+
+      if throttled
+        respond_to do |format|
+          format.json do
+            render json: {success: false, reason: "too many requests"}.to_json, status: 429
+          end
+
+          format.xml do
+            render xml: {success: false, reason: "too many requests"}.to_xml(:root => "response"), status: 429
+          end
+
+          format.html do
+            render :template => "static/too_many_requests", :status => 429
+          end
+        end
+
         return false
       end
-    # elsif request.format.to_s =~ /\/html/ && !ApiLimiter.idempotent?(request.request_method)
-    #   if ApiLimiter.throttled?(CurrentUser.id || request.remote_ip, request.request_method)
-    #     render :template => "static/too_many_requests", :status => 429
-    #   end
     end
-    
+
     return true
   end
+
   def rescue_exception(exception)
     @exception = exception
 

app/controllers/favorites_controller.rb

@@ -1,6 +1,7 @@
 class FavoritesController < ApplicationController
   before_filter :member_only
   respond_to :html, :xml, :json
+  skip_before_filter :api_check
 
   def index
     if params[:tags]

app/controllers/post_votes_controller.rb

@@ -1,5 +1,6 @@
 class PostVotesController < ApplicationController
   before_filter :voter_only
+  skip_before_filter :api_check
 
   def create
     @post = Post.find(params[:post_id])

app/controllers/users_controller.rb

@@ -2,6 +2,7 @@ class UsersController < ApplicationController
   respond_to :html, :xml, :json
   before_filter :member_only, :only => [:edit, :update, :upgrade]
   rescue_from User::PrivilegeError, :with => :access_denied
+  skip_before_filter :api_check
 
   def new
     @user = User.new

app/logical/anonymous_user.rb

@@ -180,13 +180,13 @@ class AnonymousUser
   def enable_sequential_post_navigation
     true
   end
+
+  def api_regen_multiplier
+    1
+  end
   
-  def api_hourly_limit(idempotent = false)
+  def api_burst_limit
-    if idempotent
+    5
-      500
-    else
-      5
-    end
   end
   
   def statement_timeout

app/logical/api_limiter.rb

@@ -1,26 +0,0 @@
-module ApiLimiter
-  def self.idempotent?(method)
-    case method
-    when "POST", "PUT", "DELETE", "PATCH"
-      false
-
-    else
-      true
-    end
-  end
-
-  def throttled?(user_key, http_method = "GET")
-    idempotent = ApiLimiter.idempotent?(http_method)
-    key = "api/#{user_key}/#{Time.now.hour}/#{idempotent}"
-    MEMCACHE.fetch(key, 1.hour, :raw => true) {0}
-    MEMCACHE.incr(key).to_i > CurrentUser.user.api_hourly_limit(idempotent)
-  end
-
-  def remaining_hourly_limit(user_key, idempotent = true)
-    key = "api/#{user_key}/#{Time.now.hour}/#{idempotent}"
-    requests = MEMCACHE.fetch(key, 1.hour, :raw => true) {0}.to_i
-    CurrentUser.user.api_hourly_limit(idempotent) - requests
-  end
-  
-  module_function :throttled?, :remaining_hourly_limit
-end

app/logical/daily_maintenance.rb

@@ -17,5 +17,6 @@ class DailyMaintenance
     Tag.clean_up_negative_post_counts!
     PostApproval.prune!
     SuperVoter.init!
+    TokenBucket.prune!
   end
 end

app/models/comment.rb

@@ -11,12 +11,12 @@ class Comment < ActiveRecord::Base
   before_validation :initialize_creator, :on => :create
   before_validation :initialize_updater
   after_create :update_last_commented_at_on_create
-  after_update(:if => lambda {|rec| CurrentUser.id != rec.creator_id}) do
+  after_update(:if => lambda {|rec| CurrentUser.id != rec.creator_id}) do |rec|
-    ModAction.log("comment ##{id} updated by #{CurrentUser.name}")
+    ModAction.log("comment ##{rec.id} updated by #{CurrentUser.name}")
   end
   after_destroy :update_last_commented_at_on_destroy
-  after_destroy(:if => lambda {|rec| CurrentUser.id != rec.creator_id}) do
+  after_destroy(:if => lambda {|rec| CurrentUser.id != rec.creator_id}) do |rec|
-    ModAction.log("comment ##{id} deleted by #{CurrentUser.name}")
+    ModAction.log("comment ##{rec.id} deleted by #{CurrentUser.name}")
   end
   attr_accessible :body, :post_id, :do_not_bump_post, :is_deleted, :as => [:member, :gold, :platinum, :builder, :janitor, :moderator, :admin]
   attr_accessible :is_sticky, :as => [:moderator, :admin]

app/models/forum_post.rb

@@ -19,11 +19,11 @@ class ForumPost < ActiveRecord::Base
   validate :topic_is_not_restricted, :on => :create
   before_destroy :validate_topic_is_unlocked
   after_save :delete_topic_if_original_post
-  after_update(:if => lambda {|rec| rec.updater_id != rec.creator_id}) do
+  after_update(:if => lambda {|rec| rec.updater_id != rec.creator_id}) do |rec|
-    ModAction.log("#{CurrentUser.name} updated forum post ##{id}")
+    ModAction.log("#{CurrentUser.name} updated forum post ##{rec.id}")
   end
-  after_destroy(:if => lambda {|rec| rec.updater_id != rec.creator_id}) do
+  after_destroy(:if => lambda {|rec| rec.updater_id != rec.creator_id}) do |rec|
-    ModAction.log("#{CurrentUser.name} deleted forum post ##{id}")
+    ModAction.log("#{CurrentUser.name} deleted forum post ##{rec.id}")
   end
   mentionable(
     :message_field => :body, 

app/models/ip_ban.rb

@@ -6,11 +6,11 @@ class IpBan < ActiveRecord::Base
   validates_format_of :ip_addr, :with => IP_ADDR_REGEX
   validates_uniqueness_of :ip_addr, :if => lambda {|rec| rec.ip_addr =~ IP_ADDR_REGEX}
   attr_accessible :ip_addr, :reason
-  after_create do
+  after_create do |rec|
-    ModAction.log("#{CurrentUser.name} created ip ban for #{ip_addr}")
+    ModAction.log("#{CurrentUser.name} created ip ban for #{rec.ip_addr}")
   end
-  after_destroy do
+  after_destroy do |rec|
-    ModAction.log("#{CurrentUser.name} deleted ip ban for #{ip_addr}")
+    ModAction.log("#{CurrentUser.name} deleted ip ban for ##{rec.ip_addr}")
   end
 
   def self.is_banned?(ip_addr)

app/models/token_bucket.rb

@@ -0,0 +1,40 @@
+class TokenBucket < ActiveRecord::Base
+  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!
+    TokenBucket.where(user_id: user_id).update_all(["token_count = least(token_count + (? * extract(epoch from now() - last_touched_at)), ?), last_touched_at = now()", user.api_regen_multiplier, user.api_burst_limit])
+
+    # estimate the token count to avoid reloading
+    self.token_count += (Time.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

app/models/user.rb

@@ -83,6 +83,7 @@ class User < ActiveRecord::Base
   has_one :api_key
   has_one :dmail_filter
   has_one :super_voter
+  has_one :token_bucket
   has_many :subscriptions, lambda {order("tag_subscriptions.name")}, :class_name => "TagSubscription", :foreign_key => "creator_id"
   has_many :note_versions, :foreign_key => "updater_id"
   has_many :dmails, lambda {order("dmails.id desc")}, :foreign_key => "owner_id"
@@ -598,32 +599,31 @@ class User < ActiveRecord::Base
       end
     end
 
-    def api_hourly_limit(idempotent = true)
+    def api_regen_multiplier
-      base = if is_platinum? && api_key.present?
+      # regen this amount per second
-        5000
+      if is_platinum? && api_key.present?
+        4
       elsif is_gold? && api_key.present?
-        1000
+        2
       else
-        300
+        1
-      end
-
-      if idempotent
-        base * 10
-      else
-        base
       end
     end
 
-    def remaining_api_hourly_limit
+    def api_burst_limit
-      ApiLimiter.remaining_hourly_limit(CurrentUser.ip_addr, true)
+      # can make this many api calls at once before being bound by
+      # api_regen_multiplier refilling your pool
+      if is_platinum? && api_key.present?
+        60
+      elsif is_gold? && api_key.present?
+        30
+      else
+        10
+      end
     end
 
-    def remaining_api_hourly_limit_read
+    def remaining_api_limit
-      ApiLimiter.remaining_hourly_limit(CurrentUser.ip_addr, true)
+      token_bucket.try(:token_count) || api_burst_limit
-    end
-
-    def remaining_api_hourly_limit_write
-      ApiLimiter.remaining_hourly_limit(CurrentUser.ip_addr, false)
     end
 
     def statement_timeout
@@ -645,7 +645,7 @@ class User < ActiveRecord::Base
     def method_attributes
       list = super + [:is_banned, :can_approve_posts, :can_upload_free, :is_super_voter, :level_string]
       if id == CurrentUser.user.id
-        list += [:remaining_api_hourly_limit, :remaining_api_hourly_limit_read, :remaining_api_hourly_limit_write]
+        list += [:remaining_api_limit, :api_burst_limit]
       end
       list
     end

app/models/user_feedback.rb

@@ -10,11 +10,11 @@ class UserFeedback < ActiveRecord::Base
   validate :creator_is_gold
   validate :user_is_not_creator
   after_create :create_dmail
-  after_update(:if => lambda {|rec| CurrentUser.id != rec.creator_id}) do
+  after_update(:if => lambda {|rec| rec.updater_id != rec.creator_id}) do |rec|
-    ModAction.log(%{#{CurrentUser.name} updated user feedback for "#{user_name}":/users/#{user_id}})
+    ModAction.log("#{CurrentUser.name} updated user feedback for #{rec.user_name}")
   end
-  after_destroy(:if => lambda {|rec| CurrentUser.id != rec.creator_id}) do
+  after_destroy(:if => lambda {|rec| rec.updater_id != rec.creator_id}) do |rec|
-    ModAction.log(%{#{CurrentUser.name} deleted user feedback for "#{user_name}":/users/#{user_id}})
+    ModAction.log("#{CurrentUser.name} deleted user feedback for #{rec.user_name}")
   end
 
   module SearchMethods

app/views/static/too_many_requests.html.erb

@@ -1,2 +1,2 @@
 <h1>Too Many Requests</h1>
-<p>You can only make <%= CurrentUser.api_hourly_limit(false) %> updates and <%= CurrentUser.api_hourly_limit(true) %> reads per hour.</p>
+<p>Please rate limit yourself</p>

app/views/users/_statistics.html.erb

@@ -153,6 +153,14 @@
             (<%= link_to "help", wiki_pages_path(title: "help:api") %>)
           </td>
         </tr>
+
+        <tr>
+          <th>API Limits</th>
+          <td>
+            <%= CurrentUser.user.remaining_api_limit %>
+            / <%= CurrentUser.user.api_burst_limit %>,
+          </td>
+        </tr>
       <% end %>
   </table>
 </div>

db/migrate/20170106012138_create_token_buckets.rb

@@ -0,0 +1,10 @@
+class CreateTokenBuckets < ActiveRecord::Migration
+  def up
+  	execute "create unlogged table token_buckets (user_id integer, last_touched_at timestamp not null, token_count real not null)"
+  	add_index :token_buckets, :user_id, :unique => true
+  end
+
+  def down
+  	raise NotImplementedError
+  end
+end

db/structure.sql

@@ -3054,6 +3054,17 @@ CREATE SEQUENCE tags_id_seq
 ALTER SEQUENCE tags_id_seq OWNED BY tags.id;
 
 
+--
+-- Name: token_buckets; Type: TABLE; Schema: public; Owner: -
+--
+
+CREATE UNLOGGED TABLE 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: -
 --
@@ -6963,6 +6974,13 @@ CREATE UNIQUE INDEX index_tags_on_name ON tags USING btree (name);
 CREATE INDEX index_tags_on_name_pattern ON tags USING btree (name text_pattern_ops);
 
 
+--
+-- Name: index_token_buckets_on_user_id; Type: INDEX; Schema: public; Owner: -
+--
+
+CREATE UNIQUE INDEX index_token_buckets_on_user_id ON token_buckets USING btree (user_id);
+
+
 --
 -- Name: index_uploads_on_uploader_id; Type: INDEX; Schema: public; Owner: -
 --
@@ -7389,3 +7407,5 @@ INSERT INTO schema_migrations (version) VALUES ('20161227003428');
 
 INSERT INTO schema_migrations (version) VALUES ('20161229001201');
 
+INSERT INTO schema_migrations (version) VALUES ('20170106012138');
+

test/functional/posts_controller_test.rb

@@ -19,16 +19,18 @@ class PostsControllerTest < ActionController::TestCase
     context "for api calls" do
       context "passing the api limit" do
         setup do
-          User.any_instance.stubs(:api_hourly_limit).returns(5)
+          @post = FactoryGirl.create(:post)
+          @bucket = TokenBucket.create(user_id: @user.id, token_count: 5, last_touched_at: Time.now)
+          User.any_instance.stubs(:api_burst_limit).returns(5)
         end
         
         should "work" do
-          CurrentUser.user.api_hourly_limit.times do
+          @user.api_burst_limit.times do
-            get :index, {:format => "json", :login => @user.name, :api_key => @user.api_key.key}
+            post :update, {:format => "json", :id => @post.id, :post => {:rating => "q"}, :login => @user.name, :api_key => @user.api_key.key}
             assert_response :success
           end
 
-          get :index, {:format => "json", :login => @user.name, :api_key => @user.api_key.key}
+            post :update, {:format => "json", :id => @post.id, :post => {:rating => "q"}, :login => @user.name, :api_key => @user.api_key.key}
           assert_response 429
         end
       end

test/unit/api_limiter_test.rb

@@ -1,43 +0,0 @@
-require 'test_helper'
-
-class ApiLimiterTest < ActiveSupport::TestCase
-  context "for reads" do
-    context "for an anonymous user" do
-      setup do
-        @count = 5
-        @user = AnonymousUser.new
-        CurrentUser.user = @user
-      end
-
-      should "respect api limits" do
-        @user.expects(:api_hourly_limit).with(true).times(@count + 1).returns(@count)
-
-        @count.times do
-          assert_equal(false, ApiLimiter.throttled?(CurrentUser.id || "127.0.0.1", "GET"))
-        end
-
-        assert_equal(true, ApiLimiter.throttled?(CurrentUser.id || "127.0.0.1", "GET"))
-      end
-    end
-  end
-
-  context "for writes" do
-    context "for an anonymous user" do
-      setup do
-        @count = 5
-        @user = AnonymousUser.new
-        CurrentUser.user = @user
-      end
-
-      should "respect api limits" do
-        @user.expects(:api_hourly_limit).with(false).times(@count + 1).returns(@count)
-
-        @count.times do
-          assert_equal(false, ApiLimiter.throttled?(CurrentUser.id || "127.0.0.1", "POST"))
-        end
-
-        assert_equal(true, ApiLimiter.throttled?(CurrentUser.id || "127.0.0.1", "POST"))
-      end
-    end
-  end
-end

test/unit/token_bucket_test.rb

@@ -0,0 +1,45 @@
+require 'test_helper'
+
+class TokenBucketTest < ActiveSupport::TestCase
+  context "#add!" do
+    setup do
+      @user = FactoryGirl.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 = FactoryGirl.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 = FactoryGirl.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