Add specific rate limits for posting and following (#13172)

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

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

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 } }
 

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
 

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

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

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

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
 

app/lib/exceptions.rb

@@ -8,6 +8,7 @@ module Mastodon
   class LengthValidationError < ValidationError; end
   class DimensionsValidationError < ValidationError; end
   class RaceConditionError < Error; end
+  class RateLimitExceededError < Error; end
 
   class UnexpectedResponseError < Error
     def initialize(response = nil)

app/lib/rate_limiter.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+class RateLimiter
+  include Redisable
+
+  FAMILIES = {
+    follows: {
+      limit: 400,
+      period: 24.hours.freeze,
+    }.freeze,
+
+    statuses: {
+      limit: 300,
+      period: 3.hours.freeze,
+    }.freeze,
+
+    media: {
+      limit: 30,
+      period: 30.minutes.freeze,
+    }.freeze,
+  }.freeze
+
+  def initialize(by, options = {})
+    @by     = by
+    @family = options[:family]
+    @limit  = FAMILIES[@family][:limit]
+    @period = FAMILIES[@family][:period].to_i
+  end
+
+  def record!
+    count = redis.get(key)
+
+    if count.nil?
+      redis.set(key, 0)
+      redis.expire(key, (@period - (last_epoch_time % @period) + 1).to_i)
+    end
+
+    raise Mastodon::RateLimitExceededError if count.present? && count.to_i >= @limit
+
+    redis.incr(key)
+  end
+
+  def rollback!
+    redis.decr(key)
+  end
+
+  def to_headers(now = Time.now.utc)
+    {
+      'X-RateLimit-Limit' => @limit.to_s,
+      'X-RateLimit-Remaining' => (@limit - (redis.get(key) || 0).to_i).to_s,
+      'X-RateLimit-Reset' => (now + (@period - now.to_i % @period)).iso8601(6),
+    }
+  end
+
+  private
+
+  def key
+    @key ||= "rate_limit:#{@by.id}:#{@family}:#{(last_epoch_time / @period).to_i}"
+  end
+
+  def last_epoch_time
+    @last_epoch_time ||= Time.now.to_i
+  end
+end

app/models/concerns/account_interactions.rb

@@ -87,10 +87,22 @@ module AccountInteractions
     has_many :announcement_mutes, dependent: :destroy
   end
 
-  def follow!(other_account, reblogs: nil, uri: nil)
+  def follow!(other_account, reblogs: nil, uri: nil, rate_limit: false)
     reblogs = true if reblogs.nil?
 
-    rel = active_relationships.create_with(show_reblogs: reblogs, uri: uri)
+    rel = active_relationships.create_with(show_reblogs: reblogs, uri: uri, rate_limit: rate_limit)
+                              .find_or_create_by!(target_account: other_account)
+
+    rel.update!(show_reblogs: reblogs)
+    remove_potential_friendship(other_account)
+
+    rel
+  end
+
+  def request_follow!(other_account, reblogs: nil, uri: nil, rate_limit: false)
+    reblogs = true if reblogs.nil?
+
+    rel = follow_requests.create_with(show_reblogs: reblogs, uri: uri, rate_limit: rate_limit)
                          .find_or_create_by!(target_account: other_account)
 
     rel.update!(show_reblogs: reblogs)

app/models/concerns/rate_limitable.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module RateLimitable
+  extend ActiveSupport::Concern
+
+  def rate_limit=(value)
+    @rate_limit = value
+  end
+
+  def rate_limit?
+    @rate_limit
+  end
+
+  def rate_limiter(by, options = {})
+    return @rate_limiter if defined?(@rate_limiter)
+
+    @rate_limiter = RateLimiter.new(by, options)
+  end
+
+  class_methods do
+    def rate_limit(options = {})
+      after_create do
+        by = public_send(options[:by])
+
+        if rate_limit? && by&.local?
+          rate_limiter(by, options).record!
+          @rate_limit_recorded = true
+        end
+      end
+
+      after_rollback do
+        rate_limiter(public_send(options[:by]), options).rollback! if @rate_limit_recorded
+      end
+    end
+  end
+end

app/models/follow.rb

@@ -15,6 +15,9 @@
 class Follow < ApplicationRecord
   include Paginable
   include RelationshipCacheable
+  include RateLimitable
+
+  rate_limit by: :account, family: :follows
 
   belongs_to :account
   belongs_to :target_account, class_name: 'Account'

app/models/follow_request.rb

@@ -15,6 +15,9 @@
 class FollowRequest < ApplicationRecord
   include Paginable
   include RelationshipCacheable
+  include RateLimitable
+
+  rate_limit by: :account, family: :follows
 
   belongs_to :account
   belongs_to :target_account, class_name: 'Account'

app/models/status.rb

@@ -32,6 +32,9 @@ class Status < ApplicationRecord
   include Paginable
   include Cacheable
   include StatusThreadingConcern
+  include RateLimitable
+
+  rate_limit by: :account, family: :statuses
 
   self.discard_column = :deleted_at
 

app/services/follow_service.rb

@@ -7,54 +7,68 @@ class FollowService < BaseService
   # Follow a remote user, notify remote user about the follow
   # @param [Account] source_account From which to follow
   # @param [String, Account] uri User URI to follow in the form of username@domain (or account record)
-  # @param [true, false, nil] reblogs Whether or not to show reblogs, defaults to true
+  # @param [Hash] options
-  def call(source_account, target_account, reblogs: nil, bypass_locked: false)
+  # @option [Boolean] :reblogs Whether or not to show reblogs, defaults to true
-    reblogs = true if reblogs.nil?
+  # @option [Boolean] :bypass_locked
-    target_account = ResolveAccountService.new.call(target_account, skip_webfinger: true)
+  # @option [Boolean] :with_rate_limit
-
+  def call(source_account, target_account, options = {})
-    raise ActiveRecord::RecordNotFound if target_account.nil? || target_account.id == source_account.id || target_account.suspended?
+    @source_account = source_account
-    raise Mastodon::NotPermittedError  if target_account.blocking?(source_account) || source_account.blocking?(target_account) || target_account.moved? || (!target_account.local? && target_account.ostatus?) || source_account.domain_blocking?(target_account.domain)
+    @target_account = ResolveAccountService.new.call(target_account, skip_webfinger: true)
-
+    @options        = { reblogs: true, bypass_locked: false, with_rate_limit: false }.merge(options)
-    if source_account.following?(target_account)
+
-      # We're already following this account, but we'll call follow! again to
+    raise ActiveRecord::RecordNotFound if following_not_possible?
-      # make sure the reblogs status is set correctly.
+    raise Mastodon::NotPermittedError  if following_not_allowed?
-      return source_account.follow!(target_account, reblogs: reblogs)
+
-    elsif source_account.requested?(target_account)
+    if @source_account.following?(@target_account)
-      # This isn't managed by a method in AccountInteractions, so we modify it
+      return change_follow_options!
-      # ourselves if necessary.
+    elsif @source_account.requested?(@target_account)
-      req = source_account.follow_requests.find_by(target_account: target_account)
+      return change_follow_request_options!
-      req.update!(show_reblogs: reblogs)
-      return req
     end
 
     ActivityTracker.increment('activity:interactions')
 
-    if (target_account.locked? && !bypass_locked) || source_account.silenced? || target_account.activitypub?
+    if (@target_account.locked? && !@options[:bypass_locked]) || @source_account.silenced? || @target_account.activitypub?
-      request_follow(source_account, target_account, reblogs: reblogs)
+      request_follow!
-    elsif target_account.local?
+    elsif @target_account.local?
-      direct_follow(source_account, target_account, reblogs: reblogs)
+      direct_follow!
     end
   end
 
   private
 
-  def request_follow(source_account, target_account, reblogs: true)
+  def following_not_possible?
-    follow_request = FollowRequest.create!(account: source_account, target_account: target_account, show_reblogs: reblogs)
+    @target_account.nil? || @target_account.id == @source_account.id || @target_account.suspended?
+  end
+
+  def following_not_allowed?
+    @target_account.blocking?(@source_account) || @source_account.blocking?(@target_account) || @target_account.moved? || (!@target_account.local? && @target_account.ostatus?) || @source_account.domain_blocking?(@target_account.domain)
+  end
+
+  def change_follow_options!
+    @source_account.follow!(@target_account, reblogs: @options[:reblogs])
+  end
+
+  def change_follow_request_options!
+    @source_account.request_follow!(@target_account, reblogs: @options[:reblogs])
+  end
+
+  def request_follow!
+    follow_request = @source_account.request_follow!(@target_account, reblogs: @options[:reblogs], rate_limit: @options[:with_rate_limit])
 
-    if target_account.local?
+    if @target_account.local?
-      LocalNotificationWorker.perform_async(target_account.id, follow_request.id, follow_request.class.name)
+      LocalNotificationWorker.perform_async(@target_account.id, follow_request.id, follow_request.class.name)
-    elsif target_account.activitypub?
+    elsif @target_account.activitypub?
-      ActivityPub::DeliveryWorker.perform_async(build_json(follow_request), source_account.id, target_account.inbox_url)
+      ActivityPub::DeliveryWorker.perform_async(build_json(follow_request), @source_account.id, @target_account.inbox_url)
     end
 
     follow_request
   end
 
-  def direct_follow(source_account, target_account, reblogs: true)
+  def direct_follow!
-    follow = source_account.follow!(target_account, reblogs: reblogs)
+    follow = @source_account.follow!(@target_account, reblogs: @options[:reblogs], rate_limit: @options[:with_rate_limit])
 
-    LocalNotificationWorker.perform_async(target_account.id, follow.id, follow.class.name)
+    LocalNotificationWorker.perform_async(@target_account.id, follow.id, follow.class.name)
-    MergeWorker.perform_async(target_account.id, source_account.id)
+    MergeWorker.perform_async(@target_account.id, @source_account.id)
 
     follow
   end

app/services/post_status_service.rb

@@ -19,6 +19,7 @@ class PostStatusService < BaseService
   # @option [Enumerable] :media_ids Optional array of media IDs to attach
   # @option [Doorkeeper::Application] :application
   # @option [String] :idempotency Optional idempotency key
+  # @option [Boolean] :with_rate_limit
   # @return [Status]
   def call(account, options = {})
     @account     = account
@@ -160,6 +161,7 @@ class PostStatusService < BaseService
       visibility: @visibility,
       language: language_from_option(@options[:language]) || @account.user&.setting_default_language&.presence || LanguageDetector.instance.detect(@text, @account),
       application: @options[:application],
+      rate_limit: @options[:with_rate_limit],
     }.compact
   end
 
@@ -183,6 +185,7 @@ class PostStatusService < BaseService
       options_hash[:application_id]  = options_hash.delete(:application)&.id
       options_hash[:scheduled_at]    = nil
       options_hash[:idempotency]     = nil
+      options_hash[:with_rate_limit] = false
     end
   end
 end

app/services/reblog_service.rb

@@ -8,6 +8,8 @@ class ReblogService < BaseService
   # @param [Account] account Account to reblog from
   # @param [Status] reblogged_status Status to be reblogged
   # @param [Hash] options
+  # @option [String]  :visibility
+  # @option [Boolean] :with_rate_limit
   # @return [Status]
   def call(account, reblogged_status, options = {})
     reblogged_status = reblogged_status.reblog if reblogged_status.reblog?
@@ -18,9 +20,15 @@ class ReblogService < BaseService
 
     return reblog unless reblog.nil?
 
-    visibility = options[:visibility] || account.user&.setting_default_privacy
+    visibility = begin
-    visibility = reblogged_status.visibility if reblogged_status.hidden?
+      if reblogged_status.hidden?
-    reblog = account.statuses.create!(reblog: reblogged_status, text: '', visibility: visibility)
+        reblogged_status.visibility
+      else
+        options[:visibility] || account.user&.setting_default_privacy
+      end
+    end
+
+    reblog = account.statuses.create!(reblog: reblogged_status, text: '', visibility: visibility, rate_limit: options[:with_rate_limit])
 
     DistributionWorker.perform_async(reblog.id)
     ActivityPub::DistributionWorker.perform_async(reblog.id)
@@ -45,7 +53,9 @@ class ReblogService < BaseService
 
   def bump_potential_friendship(account, reblog)
     ActivityTracker.increment('activity:interactions')
+
     return if account.following?(reblog.reblog.account_id)
+
     PotentialFriendshipTracker.record(account.id, reblog.reblog.account_id, :reblog)
   end
 

app/views/errors/429.html.haml

@@ -0,0 +1,5 @@
+- content_for :page_title do
+  = t('errors.429')
+
+- content_for :content do
+  = t('errors.429')

config/initializers/rack_attack.rb

@@ -70,7 +70,6 @@ class Rack::Attack
     req.remote_ip if req.post? && req.path == '/api/v1/accounts'
   end
 
-  # Throttle paging, as it is mainly used for public pages and AP collections
   throttle('throttle_authenticated_paging', limit: 300, period: 15.minutes) do |req|
     req.authenticated_user_id if req.paging_request?
   end

config/locales/en.yml

@@ -725,7 +725,7 @@ en:
     '422':
       content: Security verification failed. Are you blocking cookies?
       title: Security verification failed
-    '429': Throttled
+    '429': Too many requests
     '500':
       content: We're sorry, but something went wrong on our end.
       title: This page is not correct

spec/controllers/account_follow_controller_spec.rb

@@ -25,7 +25,7 @@ describe AccountFollowController do
       sign_in(user)
       subject
 
-      expect(service).to have_received(:call).with(user.account, 'alice')
+      expect(service).to have_received(:call).with(user.account, alice, with_rate_limit: true)
       expect(response).to redirect_to(account_path(alice))
     end
   end

spec/controllers/api/v1/statuses_controller_spec.rb

@@ -39,6 +39,7 @@ RSpec.describe Api::V1::StatusesController, type: :controller do
     describe 'POST #create' do
       let(:scopes) { 'write:statuses' }
 
+      context do
         before do
           post :create, params: { status: 'Hello world' }
         end
@@ -46,6 +47,43 @@ RSpec.describe Api::V1::StatusesController, type: :controller do
         it 'returns http success' do
           expect(response).to have_http_status(200)
         end
+
+        it 'returns rate limit headers' do
+          expect(response.headers['X-RateLimit-Limit']).to eq RateLimiter::FAMILIES[:statuses][:limit].to_s
+          expect(response.headers['X-RateLimit-Remaining']).to eq (RateLimiter::FAMILIES[:statuses][:limit] - 1).to_s
+        end
+      end
+
+      context 'with missing parameters' do
+        before do
+          post :create, params: {}
+        end
+
+        it 'returns http unprocessable entity' do
+          expect(response).to have_http_status(422)
+        end
+
+        it 'returns rate limit headers' do
+          expect(response.headers['X-RateLimit-Limit']).to eq RateLimiter::FAMILIES[:statuses][:limit].to_s
+        end
+      end
+
+      context 'when exceeding rate limit' do
+        before do
+          rate_limiter = RateLimiter.new(user.account, family: :statuses)
+          300.times { rate_limiter.record! }
+          post :create, params: { status: 'Hello world' }
+        end
+
+        it 'returns http too many requests' do
+          expect(response).to have_http_status(429)
+        end
+
+        it 'returns rate limit headers' do
+          expect(response.headers['X-RateLimit-Limit']).to eq RateLimiter::FAMILIES[:statuses][:limit].to_s
+          expect(response.headers['X-RateLimit-Remaining']).to eq '0'
+        end
+      end
     end
 
     describe 'DELETE #destroy' do