8 files changed, 36 insertions, 5 deletions
diff --git a/app/controllers/account_follow_controller.rb b/app/controllers/account_follow_controller.rb
index 185a355f8..33394074d 100644
--- a/app/controllers/account_follow_controller.rb
+++ b/app/controllers/account_follow_controller.rb
@@ -6,7 +6,7 @@ class AccountFollowController < ApplicationController
   before_action :authenticate_user!
 
   def create
-    FollowService.new.call(current_user.account, @account.acct)
+    FollowService.new.call(current_user.account, @account, with_rate_limit: true)
     redirect_to account_path(@account)
   end
 end
diff --git a/app/controllers/api/base_controller.rb b/app/controllers/api/base_controller.rb
index 68bf425f4..153ade253 100644
--- a/app/controllers/api/base_controller.rb
+++ b/app/controllers/api/base_controller.rb
@@ -44,6 +44,10 @@ class Api::BaseController < ApplicationController
     render json: { error: 'There was a temporary problem serving your request, please try again' }, status: 503
   end
 
+  rescue_from Mastodon::RateLimitExceededError do
+    render json: { error: I18n.t('errors.429') }, status: 429
+  end
+
   rescue_from ActionController::ParameterMissing do |e|
     render json: { error: e.to_s }, status: 400
   end
diff --git a/app/controllers/api/v1/accounts_controller.rb b/app/controllers/api/v1/accounts_controller.rb
index f5b06862b..0080faf33 100644
--- a/app/controllers/api/v1/accounts_controller.rb
+++ b/app/controllers/api/v1/accounts_controller.rb
@@ -14,6 +14,8 @@ class Api::V1::AccountsController < Api::BaseController
 
   skip_before_action :require_authenticated_user!, only: :create
 
+  override_rate_limit_headers :follow, family: :follows
+
   def show
     render json: @account, serializer: REST::AccountSerializer
   end
@@ -29,7 +31,7 @@ class Api::V1::AccountsController < Api::BaseController
   end
 
   def follow
-    FollowService.new.call(current_user.account, @account, reblogs: truthy_param?(:reblogs))
+    FollowService.new.call(current_user.account, @account, reblogs: truthy_param?(:reblogs), with_rate_limit: true)
 
     options = @account.locked? || current_user.account.silenced? ? {} : { following_map: { @account.id => { reblogs: truthy_param?(:reblogs) } }, requested_map: { @account.id => false } }
 
diff --git a/app/controllers/api/v1/statuses/reblogs_controller.rb b/app/controllers/api/v1/statuses/reblogs_controller.rb
index 9abeb0759..7fa774a4d 100644
--- a/app/controllers/api/v1/statuses/reblogs_controller.rb
+++ b/app/controllers/api/v1/statuses/reblogs_controller.rb
@@ -7,8 +7,11 @@ class Api::V1::Statuses::ReblogsController < Api::BaseController
   before_action :require_user!
   before_action :set_reblog
 
+  override_rate_limit_headers :create, family: :statuses
+
   def create
     @status = ReblogService.new.call(current_account, @reblog, reblog_params)
+
     render json: @status, serializer: REST::StatusSerializer
   end
 
diff --git a/app/controllers/api/v1/statuses_controller.rb b/app/controllers/api/v1/statuses_controller.rb
index 85521fadf..2f55e95fd 100644
--- a/app/controllers/api/v1/statuses_controller.rb
+++ b/app/controllers/api/v1/statuses_controller.rb
@@ -8,6 +8,8 @@ class Api::V1::StatusesController < Api::BaseController
   before_action :require_user!, except:  [:show, :context]
   before_action :set_status, only:       [:show, :context]
 
+  override_rate_limit_headers :create, family: :statuses
+
   # This API was originally unlimited, pagination cannot be introduced without
   # breaking backwards-compatibility. Arbitrarily high number to cover most
   # conversations as quasi-unlimited, it would be too much work to render more
@@ -42,7 +44,8 @@ class Api::V1::StatusesController < Api::BaseController
                                          scheduled_at: status_params[:scheduled_at],
                                          application: doorkeeper_token.application,
                                          poll: status_params[:poll],
-                                         idempotency: request.headers['Idempotency-Key'])
+                                         idempotency: request.headers['Idempotency-Key'],
+                                         with_rate_limit: true)
 
     render json: @status, serializer: @status.is_a?(ScheduledStatus) ? REST::ScheduledStatusSerializer : REST::StatusSerializer
   end
diff --git a/app/controllers/application_controller.rb b/app/controllers/application_controller.rb
index 0cfa2b386..973db6aca 100644
--- a/app/controllers/application_controller.rb
+++ b/app/controllers/application_controller.rb
@@ -29,6 +29,7 @@ class ApplicationController < ActionController::Base
   rescue_from Mastodon::NotPermittedError, with: :forbidden
   rescue_from HTTP::Error, OpenSSL::SSL::SSLError, with: :internal_server_error
   rescue_from Mastodon::RaceConditionError, with: :service_unavailable
+  rescue_from Mastodon::RateLimitExceededError, with: :too_many_requests
 
   before_action :store_current_location, except: :raise_not_found, unless: :devise_controller?
   before_action :require_functional!, if: :user_signed_in?
@@ -111,6 +112,10 @@ class ApplicationController < ActionController::Base
     respond_with_error(503)
   end
 
+  def too_many_requests
+    respond_with_error(429)
+  end
+
   def single_user_mode?
     @single_user_mode ||= Rails.configuration.x.single_user_mode && Account.where('id > 0').exists?
   end
diff --git a/app/controllers/authorize_interactions_controller.rb b/app/controllers/authorize_interactions_controller.rb
index e27366ea3..29c0288d0 100644
--- a/app/controllers/authorize_interactions_controller.rb
+++ b/app/controllers/authorize_interactions_controller.rb
@@ -20,7 +20,7 @@ class AuthorizeInteractionsController < ApplicationController
   end
 
   def create
-    if @resource.is_a?(Account) && FollowService.new.call(current_account, @resource)
+    if @resource.is_a?(Account) && FollowService.new.call(current_account, @resource, with_rate_limit: true)
       render :success
     else
       render :error
diff --git a/app/controllers/concerns/rate_limit_headers.rb b/app/controllers/concerns/rate_limit_headers.rb
index b79c558d8..86fe58a71 100644
--- a/app/controllers/concerns/rate_limit_headers.rb
+++ b/app/controllers/concerns/rate_limit_headers.rb
@@ -3,6 +3,20 @@
 module RateLimitHeaders
   extend ActiveSupport::Concern
 
+  class_methods do
+    def override_rate_limit_headers(method_name, options = {})
+      around_action(only: method_name, if: :current_account) do |_controller, block|
+        begin
+          block.call
+        ensure
+          rate_limiter = RateLimiter.new(current_account, options)
+          rate_limit_headers = rate_limiter.to_headers
+          response.headers.merge!(rate_limit_headers) unless response.headers['X-RateLimit-Remaining'].present? && rate_limit_headers['X-RateLimit-Remaining'].to_i > response.headers['X-RateLimit-Remaining'].to_i
+        end
+      end
+    end
+  end
+
   included do
     before_action :set_rate_limit_headers, if: :rate_limited_request?
   end
@@ -44,7 +58,7 @@ module RateLimitHeaders
   end
 
   def api_throttle_data
-    most_limited_type, = request.env['rack.attack.throttle_data'].min_by { |_, v| v[:limit] }
+    most_limited_type, = request.env['rack.attack.throttle_data'].min_by { |_, v| v[:limit] - v[:count] }
     request.env['rack.attack.throttle_data'][most_limited_type]
   end