From ef6829382aa32c03cf8536422537a9c219bd0035 Mon Sep 17 00:00:00 2001
From: Ekaterina Vaartis <vaartis@cock.li>
Date: Sun, 16 Dec 2018 23:41:11 +0300
Subject: [PATCH] Clean captchas up periodically, not schedule it after theyre
 created

---
 lib/pleroma/captcha/captcha.ex         | 20 +++++++++++---------
 lib/pleroma/captcha/captcha_service.ex |  2 +-
 lib/pleroma/captcha/kocaptcha.ex       | 26 +++++++++++++++++---------
 test/support/captcha_mock.ex           |  2 +-
 4 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/lib/pleroma/captcha/captcha.ex b/lib/pleroma/captcha/captcha.ex
index 26477a214..5630f6b57 100644
--- a/lib/pleroma/captcha/captcha.ex
+++ b/lib/pleroma/captcha/captcha.ex
@@ -14,6 +14,10 @@ defmodule Pleroma.Captcha do
     ets_name = Module.concat(method(), Ets)
     ^ets_name = :ets.new(Module.concat(method(), Ets), @ets_options)
 
+    # Clean up old captchas every few minutes
+    seconds_retained = Pleroma.Config.get!([__MODULE__, :seconds_retained])
+    Process.send_after(self(), :cleanup, 1000 * seconds_retained)
+
     {:ok, nil}
   end
 
@@ -38,13 +42,7 @@ defmodule Pleroma.Captcha do
     if !enabled do
       {:reply, %{type: :none}, state}
     else
-      new_captcha = method().new()
-
-      seconds_retained = Pleroma.Config.get!([__MODULE__, :seconds_retained])
-      # Wait several minutes and if the captcha is still there, delete it
-      Process.send_after(self(), {:cleanup, new_captcha.token}, 1000 * seconds_retained)
-
-      {:reply, new_captcha, state}
+      {:reply, method().new(), state}
     end
   end
 
@@ -54,8 +52,12 @@ defmodule Pleroma.Captcha do
   end
 
   @doc false
-  def handle_info({:cleanup, token}, state) do
-    method().cleanup(token)
+  def handle_info(:cleanup, state) do
+    :ok = method().cleanup()
+
+    seconds_retained = Pleroma.Config.get!([__MODULE__, :seconds_retained])
+    # Schedule the next clenup
+    Process.send_after(self(), :cleanup, 1000 * seconds_retained)
 
     {:noreply, state}
   end
diff --git a/lib/pleroma/captcha/captcha_service.ex b/lib/pleroma/captcha/captcha_service.ex
index fe5a6bf66..8d0b76f86 100644
--- a/lib/pleroma/captcha/captcha_service.ex
+++ b/lib/pleroma/captcha/captcha_service.ex
@@ -24,5 +24,5 @@ defmodule Pleroma.Captcha.Service do
   @doc """
   This function is called periodically to clean up old captchas
   """
-  @callback cleanup(token :: String.t()) :: :ok
+  @callback cleanup() :: :ok
 end
diff --git a/lib/pleroma/captcha/kocaptcha.ex b/lib/pleroma/captcha/kocaptcha.ex
index 9891d4031..7f9637ad0 100644
--- a/lib/pleroma/captcha/kocaptcha.ex
+++ b/lib/pleroma/captcha/kocaptcha.ex
@@ -1,4 +1,6 @@
 defmodule Pleroma.Captcha.Kocaptcha do
+  alias Calendar.DateTime
+
   alias Pleroma.Captcha.Service
   @behaviour Service
 
@@ -17,7 +19,7 @@ defmodule Pleroma.Captcha.Kocaptcha do
 
         token = json_resp["token"]
 
-        true = :ets.insert(@ets, {token, json_resp["md5"]})
+        true = :ets.insert(@ets, {token, json_resp["md5"], DateTime.now_utc()})
 
         %{type: :kocaptcha, token: token, url: endpoint <> json_resp["url"]}
     end
@@ -26,10 +28,10 @@ defmodule Pleroma.Captcha.Kocaptcha do
   @impl Service
   def validate(token, captcha) do
     with false <- is_nil(captcha),
-         [{^token, saved_md5}] <- :ets.lookup(@ets, token),
+         [{^token, saved_md5, _}] <- :ets.lookup(@ets, token),
          true <- :crypto.hash(:md5, captcha) |> Base.encode16() == String.upcase(saved_md5) do
       # Clear the saved value
-      cleanup(token)
+      :ets.delete(@ets, token)
 
       true
     else
@@ -38,11 +40,17 @@ defmodule Pleroma.Captcha.Kocaptcha do
   end
 
   @impl Service
-  def cleanup(token) do
-    # Only delete the entry if it exists in the table, because ets:delete raises an exception if it does not
-    case :ets.lookup(@ets, token) do
-      [{^token, _}] -> :ets.delete(@ets, token)
-      _ -> true
-    end
+  def cleanup() do
+    seconds_retained = Pleroma.Config.get!([Pleroma.Captcha, :seconds_retained])
+
+    # Go through captchas and remove expired ones
+    :ets.tab2list(@ets)
+    |> Enum.each(fn {token, _, time_inserted} ->
+      # time created + expiration time = time when the captcha should be removed
+      remove_time = DateTime.add!(time_inserted, seconds_retained)
+      if DateTime.after?(DateTime.now_utc(), remove_time), do: :ets.delete(@ets, token)
+    end)
+
+    :ok
   end
 end
diff --git a/test/support/captcha_mock.ex b/test/support/captcha_mock.ex
index 560d6c457..898aa17b8 100644
--- a/test/support/captcha_mock.ex
+++ b/test/support/captcha_mock.ex
@@ -9,5 +9,5 @@ defmodule Pleroma.Captcha.Mock do
   def validate(_token, _captcha), do: true
 
   @impl Service
-  def cleanup(_token), do: true
+  def cleanup(), do: :ok
 end