From dce7e429286dfe8cb44a27c50713a03f0e696357 Mon Sep 17 00:00:00 2001
From: Tusooa Zhu <tusooa@kazv.moe>
Date: Sat, 23 Jul 2022 22:34:03 -0400
Subject: [PATCH] Make MediaProxyWarmingPolicy history-aware

---
 .../mrf/media_proxy_warming_policy.ex         |  9 ++--
 .../mrf/media_proxy_warming_policy_test.exs   | 44 +++++++++++++++++++
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/lib/pleroma/web/activity_pub/mrf/media_proxy_warming_policy.ex b/lib/pleroma/web/activity_pub/mrf/media_proxy_warming_policy.ex
index 0eac8f021..c95d35bb9 100644
--- a/lib/pleroma/web/activity_pub/mrf/media_proxy_warming_policy.ex
+++ b/lib/pleroma/web/activity_pub/mrf/media_proxy_warming_policy.ex
@@ -16,6 +16,9 @@ defmodule Pleroma.Web.ActivityPub.MRF.MediaProxyWarmingPolicy do
     recv_timeout: 10_000
   ]
 
+  @impl true
+  def history_awareness, do: :auto
+
   defp prefetch(url) do
     # Fetching only proxiable resources
     if MediaProxy.enabled?() and MediaProxy.url_proxiable?(url) do
@@ -54,10 +57,8 @@ defmodule Pleroma.Web.ActivityPub.MRF.MediaProxyWarmingPolicy do
   end
 
   @impl true
-  def filter(
-        %{"type" => "Create", "object" => %{"attachment" => attachments} = _object} = message
-      )
-      when is_list(attachments) and length(attachments) > 0 do
+  def filter(%{"type" => type, "object" => %{"attachment" => attachments} = _object} = message)
+      when type in ["Create", "Update"] and is_list(attachments) and length(attachments) > 0 do
     preload(message)
 
     {:ok, message}
diff --git a/test/pleroma/web/activity_pub/mrf/media_proxy_warming_policy_test.exs b/test/pleroma/web/activity_pub/mrf/media_proxy_warming_policy_test.exs
index 09301c6ca..6557c3a98 100644
--- a/test/pleroma/web/activity_pub/mrf/media_proxy_warming_policy_test.exs
+++ b/test/pleroma/web/activity_pub/mrf/media_proxy_warming_policy_test.exs
@@ -7,6 +7,7 @@ defmodule Pleroma.Web.ActivityPub.MRF.MediaProxyWarmingPolicyTest do
   use Pleroma.Tests.Helpers
 
   alias Pleroma.HTTP
+  alias Pleroma.Web.ActivityPub.MRF
   alias Pleroma.Web.ActivityPub.MRF.MediaProxyWarmingPolicy
 
   import Mock
@@ -22,6 +23,25 @@ defmodule Pleroma.Web.ActivityPub.MRF.MediaProxyWarmingPolicyTest do
     }
   }
 
+  @message_with_history %{
+    "type" => "Create",
+    "object" => %{
+      "type" => "Note",
+      "content" => "content",
+      "formerRepresentations" => %{
+        "orderedItems" => [
+          %{
+            "type" => "Note",
+            "content" => "content",
+            "attachment" => [
+              %{"url" => [%{"href" => "http://example.com/image.jpg"}]}
+            ]
+          }
+        ]
+      }
+    }
+  }
+
   setup do: clear_config([:media_proxy, :enabled], true)
 
   test "it prefetches media proxy URIs" do
@@ -50,4 +70,28 @@ defmodule Pleroma.Web.ActivityPub.MRF.MediaProxyWarmingPolicyTest do
       refute called(HTTP.get(:_, :_, :_))
     end
   end
+
+  test "history-aware" do
+    Tesla.Mock.mock(fn %{method: :get, url: "http://example.com/image.jpg"} ->
+      {:ok, %Tesla.Env{status: 200, body: ""}}
+    end)
+
+    with_mock HTTP, get: fn _, _, _ -> {:ok, []} end do
+      MRF.filter_one(MediaProxyWarmingPolicy, @message_with_history)
+
+      assert called(HTTP.get(:_, :_, :_))
+    end
+  end
+
+  test "works with Updates" do
+    Tesla.Mock.mock(fn %{method: :get, url: "http://example.com/image.jpg"} ->
+      {:ok, %Tesla.Env{status: 200, body: ""}}
+    end)
+
+    with_mock HTTP, get: fn _, _, _ -> {:ok, []} end do
+      MRF.filter_one(MediaProxyWarmingPolicy, @message_with_history |> Map.put("type", "Update"))
+
+      assert called(HTTP.get(:_, :_, :_))
+    end
+  end
 end