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 ApiLimiter.throttled?(CurrentUser.id || request.remote_ip, request.request_method)
-        render :text => "429 Too Many Requests\n", :layout => false, :status => 429
+    if !CurrentUser.is_anonymous? && !request.get? && !request.head?
+      if CurrentUser.user.token_bucket.nil?
+        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)
-    if idempotent
-      500
-    else
-      5
-    end
+  def api_burst_limit
+    5
   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
-    ModAction.log("comment ##{id} updated by #{CurrentUser.name}")
+  after_update(:if => lambda {|rec| CurrentUser.id != rec.creator_id}) do |rec|
+    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
-    ModAction.log("comment ##{id} deleted by #{CurrentUser.name}")
+  after_destroy(:if => lambda {|rec| CurrentUser.id != rec.creator_id}) do |rec|
+    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
-    ModAction.log("#{CurrentUser.name} updated forum post ##{id}")
+  after_update(:if => lambda {|rec| rec.updater_id != rec.creator_id}) do |rec|
+    ModAction.log("#{CurrentUser.name} updated forum post ##{rec.id}")
   end
-  after_destroy(:if => lambda {|rec| rec.updater_id != rec.creator_id}) do
-    ModAction.log("#{CurrentUser.name} deleted forum post ##{id}")
+  after_destroy(:if => lambda {|rec| rec.updater_id != rec.creator_id}) do |rec|
+    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
-    ModAction.log("#{CurrentUser.name} created ip ban for #{ip_addr}")
+  after_create do |rec|
+    ModAction.log("#{CurrentUser.name} created ip ban for #{rec.ip_addr}")
   end
-  after_destroy do
-    ModAction.log("#{CurrentUser.name} deleted ip ban for #{ip_addr}")
+  after_destroy do |rec|
+    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)
-      base = if is_platinum? && api_key.present?
-        5000
+    def api_regen_multiplier
+      # regen this amount per second
+      if is_platinum? && api_key.present?
+        4
       elsif is_gold? && api_key.present?
-        1000
+        2
       else
-        300
-      end
-
-      if idempotent
-        base * 10
-      else
-        base
+        1
       end
     end
 
-    def remaining_api_hourly_limit
-      ApiLimiter.remaining_hourly_limit(CurrentUser.ip_addr, true)
+    def api_burst_limit
+      # 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
-      ApiLimiter.remaining_hourly_limit(CurrentUser.ip_addr, true)
-    end
-
-    def remaining_api_hourly_limit_write
-      ApiLimiter.remaining_hourly_limit(CurrentUser.ip_addr, false)
+    def remaining_api_limit
+      token_bucket.try(:token_count) || api_burst_limit
     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
-    ModAction.log(%{#{CurrentUser.name} updated user feedback for "#{user_name}":/users/#{user_id}})
+  after_update(:if => lambda {|rec| rec.updater_id != rec.creator_id}) do |rec|
+    ModAction.log("#{CurrentUser.name} updated user feedback for #{rec.user_name}")
   end
-  after_destroy(:if => lambda {|rec| CurrentUser.id != rec.creator_id}) do
-    ModAction.log(%{#{CurrentUser.name} deleted user feedback for "#{user_name}":/users/#{user_id}})
+  after_destroy(:if => lambda {|rec| rec.updater_id != rec.creator_id}) do |rec|
+    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>