From 54edb4b853522b322922b65989ddfc5fdf928014 Mon Sep 17 00:00:00 2001
From: Eugen Rochko <eugen@zeonfederated.com>
Date: Sat, 16 Sep 2017 03:01:45 +0200
Subject: [PATCH] When accessing uncached media attachment, redownload it
 (#4955)

* When accessing uncached media attachment, redownload it

* Prevent re-download of rejected media
---
 app/controllers/media_proxy_controller.rb     | 40 +++++++++++++++++++
 app/models/media_attachment.rb                | 10 ++++-
 .../rest/media_attachment_serializer.rb       | 12 +++++-
 config/routes.rb                              |  2 +
 lib/tasks/mastodon.rake                       |  1 -
 5 files changed, 60 insertions(+), 5 deletions(-)
 create mode 100644 app/controllers/media_proxy_controller.rb

diff --git a/app/controllers/media_proxy_controller.rb b/app/controllers/media_proxy_controller.rb
new file mode 100644
index 000000000..6a83cf9dc
--- /dev/null
+++ b/app/controllers/media_proxy_controller.rb
@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+class MediaProxyController < ApplicationController
+  include RoutingHelper
+
+  def show
+    RedisLock.acquire(lock_options) do |lock|
+      if lock.acquired?
+        @media_attachment = MediaAttachment.remote.find(params[:id])
+        redownload! if @media_attachment.needs_redownload? && !reject_media?
+      end
+    end
+
+    redirect_to full_asset_url(@media_attachment.file.url(version))
+  end
+
+  private
+
+  def redownload!
+    @media_attachment.file_remote_url = @media_attachment.remote_url
+    @media_attachment.touch(:created_at)
+    @media_attachment.save!
+  end
+
+  def version
+    if request.path.ends_with?('/small')
+      :small
+    else
+      :original
+    end
+  end
+
+  def lock_options
+    { redis: Redis.current, key: "media_download:#{params[:id]}" }
+  end
+
+  def reject_media?
+    DomainBlock.find_by(domain: @media_attachment.account.domain)&.reject_media?
+  end
+end
diff --git a/app/models/media_attachment.rb b/app/models/media_attachment.rb
index d83ca44f1..d913e7372 100644
--- a/app/models/media_attachment.rb
+++ b/app/models/media_attachment.rb
@@ -56,15 +56,21 @@ class MediaAttachment < ApplicationRecord
 
   validates :account, presence: true
 
-  scope :attached, -> { where.not(status_id: nil) }
+  scope :attached,   -> { where.not(status_id: nil) }
   scope :unattached, -> { where(status_id: nil) }
-  scope :local, -> { where(remote_url: '') }
+  scope :local,      -> { where(remote_url: '') }
+  scope :remote,     -> { where.not(remote_url: '') }
+
   default_scope { order(id: :asc) }
 
   def local?
     remote_url.blank?
   end
 
+  def needs_redownload?
+    file.blank? && remote_url.present?
+  end
+
   def to_param
     shortcode
   end
diff --git a/app/serializers/rest/media_attachment_serializer.rb b/app/serializers/rest/media_attachment_serializer.rb
index 9055b8db4..31189406a 100644
--- a/app/serializers/rest/media_attachment_serializer.rb
+++ b/app/serializers/rest/media_attachment_serializer.rb
@@ -7,11 +7,19 @@ class REST::MediaAttachmentSerializer < ActiveModel::Serializer
              :remote_url, :text_url, :meta
 
   def url
-    full_asset_url(object.file.url(:original))
+    if object.needs_redownload?
+      media_proxy_url(object.id, :original)
+    else
+      full_asset_url(object.file.url(:original))
+    end
   end
 
   def preview_url
-    full_asset_url(object.file.url(:small))
+    if object.needs_redownload?
+      media_proxy_url(object.id, :small)
+    else
+      full_asset_url(object.file.url(:small))
+    end
   end
 
   def text_url
diff --git a/config/routes.rb b/config/routes.rb
index 63e94fb49..bf5428869 100644
--- a/config/routes.rb
+++ b/config/routes.rb
@@ -98,6 +98,8 @@ Rails.application.routes.draw do
   resources :media, only: [:show]
   resources :tags,  only: [:show]
 
+  get '/media_proxy/:id/(*any)', to: 'media_proxy#show', as: :media_proxy
+
   # Remote follow
   resource :authorize_follow, only: [:show, :create]
   resource :share, only: [:show, :create]
diff --git a/lib/tasks/mastodon.rake b/lib/tasks/mastodon.rake
index 2aef752ca..5614ddf48 100644
--- a/lib/tasks/mastodon.rake
+++ b/lib/tasks/mastodon.rake
@@ -83,7 +83,6 @@ namespace :mastodon do
 
       MediaAttachment.where.not(remote_url: '').where('created_at < ?', time_ago).find_each do |media|
         media.file.destroy
-        media.type = :unknown
         media.save
       end
     end