7 files changed, 53 insertions, 6 deletions
diff --git a/app/controllers/api/activitypub/activities_controller.rb b/app/controllers/api/activitypub/activities_controller.rb
index ba03738fc..025ab960e 100644
--- a/app/controllers/api/activitypub/activities_controller.rb
+++ b/app/controllers/api/activitypub/activities_controller.rb
@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 class Api::Activitypub::ActivitiesController < ApiController
+  include Authorization
+
   # before_action :set_follow, only: [:show_follow]
   before_action :set_status, only: [:show_status]
 
@@ -8,7 +10,7 @@ class Api::Activitypub::ActivitiesController < ApiController
 
   # Show a status in AS2 format, as either an Announce (reblog) or a Create (post) activity.
   def show_status
-    return forbidden unless @status.permitted?
+    authorize @status, :show?
 
     if @status.reblog?
       render :show_status_announce
diff --git a/app/controllers/api/activitypub/notes_controller.rb b/app/controllers/api/activitypub/notes_controller.rb
index 6489243dc..ff9383413 100644
--- a/app/controllers/api/activitypub/notes_controller.rb
+++ b/app/controllers/api/activitypub/notes_controller.rb
@@ -1,12 +1,14 @@
 # frozen_string_literal: true
 
 class Api::Activitypub::NotesController < ApiController
+  include Authorization
+
   before_action :set_status
 
   respond_to :activitystreams2
 
   def show
-    forbidden unless @status.permitted?
+    authorize @status, :show?
   end
 
   private
diff --git a/app/controllers/api/v1/statuses_controller.rb b/app/controllers/api/v1/statuses_controller.rb
index 852ffc3ab..592540f45 100644
--- a/app/controllers/api/v1/statuses_controller.rb
+++ b/app/controllers/api/v1/statuses_controller.rb
@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 class Api::V1::StatusesController < ApiController
+  include Authorization
+
   before_action :authorize_if_got_token, except:            [:create, :destroy, :reblog, :unreblog, :favourite, :unfavourite, :mute, :unmute]
   before_action -> { doorkeeper_authorize! :write }, only:  [:create, :destroy, :reblog, :unreblog, :favourite, :unfavourite, :mute, :unmute]
   before_action :require_user!, except:  [:show, :context, :card, :reblogged_by, :favourited_by]
@@ -130,7 +132,10 @@ class Api::V1::StatusesController < ApiController
 
   def set_status
     @status = Status.find(params[:id])
-    raise ActiveRecord::RecordNotFound unless @status.permitted?(current_account)
+    authorize @status, :show?
+  rescue Mastodon::NotPermittedError
+    # Reraise in order to get a 404 instead of a 403 error code
+    raise ActiveRecord::RecordNotFound
   end
 
   def set_conversation
diff --git a/app/controllers/concerns/authorization.rb b/app/controllers/concerns/authorization.rb
new file mode 100644
index 000000000..7828fe48d
--- /dev/null
+++ b/app/controllers/concerns/authorization.rb
@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module Authorization
+  extend ActiveSupport::Concern
+  include Pundit
+
+  def pundit_user
+    current_account
+  end
+
+  def authorize(*)
+    super
+  rescue Pundit::NotAuthorizedError
+    raise Mastodon::NotPermittedError
+  end
+
+  def authorize_with(user, record, query)
+    Pundit.authorize(user, record, query)
+  rescue Pundit::NotAuthorizedError
+    raise Mastodon::NotPermittedError
+  end
+end
diff --git a/app/controllers/media_controller.rb b/app/controllers/media_controller.rb
index fa1daf012..f652f5ace 100644
--- a/app/controllers/media_controller.rb
+++ b/app/controllers/media_controller.rb
@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 class MediaController < ApplicationController
+  include Authorization
+
   before_action :verify_permitted_status
 
   def show
@@ -14,6 +16,9 @@ class MediaController < ApplicationController
   end
 
   def verify_permitted_status
-    raise ActiveRecord::RecordNotFound unless media_attachment.status.permitted?(current_account)
+    authorize media_attachment.status, :show?
+  rescue Mastodon::NotPermittedError
+    # Reraise in order to get a 404 instead of a 403 error code
+    raise ActiveRecord::RecordNotFound
   end
 end
diff --git a/app/controllers/statuses_controller.rb b/app/controllers/statuses_controller.rb
index 696bb4f52..59c9d0a87 100644
--- a/app/controllers/statuses_controller.rb
+++ b/app/controllers/statuses_controller.rb
@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 class StatusesController < ApplicationController
+  include Authorization
+
   layout 'public'
 
   before_action :set_account
@@ -30,7 +32,10 @@ class StatusesController < ApplicationController
     @stream_entry = @status.stream_entry
     @type         = @stream_entry.activity_type.downcase
 
-    raise ActiveRecord::RecordNotFound unless @status.permitted?(current_account)
+    authorize @status, :show?
+  rescue Mastodon::NotPermittedError
+    # Reraise in order to get a 404
+    raise ActiveRecord::RecordNotFound
   end
 
   def check_account_suspension
diff --git a/app/controllers/stream_entries_controller.rb b/app/controllers/stream_entries_controller.rb
index baff4317a..314d59619 100644
--- a/app/controllers/stream_entries_controller.rb
+++ b/app/controllers/stream_entries_controller.rb
@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 class StreamEntriesController < ApplicationController
+  include Authorization
+
   layout 'public'
 
   before_action :set_account
@@ -42,7 +44,11 @@ class StreamEntriesController < ApplicationController
     @stream_entry = @account.stream_entries.where(activity_type: 'Status').find(params[:id])
     @type         = @stream_entry.activity_type.downcase
 
-    raise ActiveRecord::RecordNotFound if @stream_entry.activity.nil? || (@stream_entry.hidden? && !@stream_entry.activity.permitted?(current_account))
+    raise ActiveRecord::RecordNotFound if @stream_entry.activity.nil?
+    authorize @stream_entry.activity, :show? if @stream_entry.hidden?
+  rescue Mastodon::NotPermittedError
+    # Reraise in order to get a 404
+    raise ActiveRecord::RecordNotFound
   end
 
   def check_account_suspension