From 051a204d5588073db94311df4dd5d0a20d003a99 Mon Sep 17 00:00:00 2001
From: Dan Schultzer <dan@dreamconception.com>
Date: Sat, 18 Apr 2020 17:19:07 -0700
Subject: [PATCH] Add TOTP support in the invalidation session plug (#87)

---
 .../plugs/pow_invalidated_session_plug.ex     | 18 ++++++--
 .../pow_invalidated_session_plug_test.exs     | 45 +++++++++++++++++++
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/lib/philomena_web/plugs/pow_invalidated_session_plug.ex b/lib/philomena_web/plugs/pow_invalidated_session_plug.ex
index ba967709..f4b9913b 100644
--- a/lib/philomena_web/plugs/pow_invalidated_session_plug.ex
+++ b/lib/philomena_web/plugs/pow_invalidated_session_plug.ex
@@ -87,9 +87,14 @@ defmodule PhilomenaWeb.PowInvalidatedSessionPlug do
   defp maybe_put_cache(conn, user, old_token, opts) do
     fetch_fn = Keyword.fetch!(opts, :fetch_token)
 
+    metadata =
+      conn.private
+      |> Map.get(:pow_session_metadata, [])
+      |> Keyword.take([:valid_totp_at])
+
     case fetch_fn.(conn) do
       ^old_token -> conn
-      _token -> put_cache(conn, user, old_token, opts)
+      _token -> put_cache(conn, {user, metadata}, old_token, opts)
     end
   end
 
@@ -106,8 +111,15 @@ defmodule PhilomenaWeb.PowInvalidatedSessionPlug do
     {store, store_config} = invalidated_cache(conn, opts)
 
     case store.get(store_config, token) do
-      :not_found -> conn
-      user -> Plug.assign_current_user(conn, user, config)
+      :not_found ->
+        conn
+
+      {user, metadata} ->
+        metadata = Keyword.merge(metadata, conn.private[:pow_session_metadata] || [])
+
+        conn
+        |> Conn.put_private(:pow_session_metadata, metadata)
+        |> Plug.assign_current_user(user, config)
     end
   end
 
diff --git a/test/philomena_web/plug/pow_invalidated_session_plug_test.exs b/test/philomena_web/plug/pow_invalidated_session_plug_test.exs
index d8b05915..73b26192 100644
--- a/test/philomena_web/plug/pow_invalidated_session_plug_test.exs
+++ b/test/philomena_web/plug/pow_invalidated_session_plug_test.exs
@@ -43,12 +43,16 @@ defmodule PhilomenaWeb.PowInvalidatedSessionPlugTest do
 
     assert Pow.Plug.current_user(conn).id == user.id
     assert Conn.get_session(conn, @session_key) != session_id
+    assert metadata = conn.private[:pow_session_metadata]
+    refute metadata[:valid_totp_at]
 
     :timer.sleep(100)
     conn = run_plug(init_conn, config)
 
     assert Pow.Plug.current_user(conn).id == user.id
     assert Conn.get_session(conn, @session_key) == session_id
+    assert metadata = conn.private[:pow_session_metadata]
+    refute metadata[:valid_totp_at]
 
     :timer.sleep(@invalidated_ttl - 100)
     conn = run_plug(init_conn)
@@ -65,12 +69,16 @@ defmodule PhilomenaWeb.PowInvalidatedSessionPlugTest do
 
     assert Pow.Plug.current_user(conn).id == user.id
     assert conn.cookies[@cookie_key] != persistent_session_id
+    assert metadata = conn.private[:pow_session_metadata]
+    refute metadata[:valid_totp_at]
 
     :timer.sleep(100)
     conn = run_plug(init_conn)
 
     assert Pow.Plug.current_user(conn).id == user.id
     assert conn.cookies[@cookie_key] == persistent_session_id
+    assert metadata = conn.private[:pow_session_metadata]
+    refute metadata[:valid_totp_at]
 
     :timer.sleep(@invalidated_ttl - 100)
     conn = run_plug(init_conn)
@@ -79,6 +87,42 @@ defmodule PhilomenaWeb.PowInvalidatedSessionPlugTest do
     assert conn.cookies[@cookie_key] == persistent_session_id
   end
 
+  test "call/2 with TOTP turned on", %{conn: init_conn, user: user} do
+    user =
+      user
+      |> Ecto.Changeset.change(%{otp_required_for_login: true})
+      |> Repo.update!()
+
+    config = Keyword.put(@config, :session_ttl_renewal, 0)
+
+    no_otp_auth_conn =
+      init_conn
+      |> prepare_session_conn(user, config)
+      |> init_plug(config)
+
+    assert no_otp_auth_conn.halted
+
+    init_conn =
+      init_conn
+      |> Conn.put_private(:pow_session_metadata, valid_totp_at: DateTime.utc_now())
+      |> prepare_session_conn(user, config)
+
+    conn = run_plug(init_conn, config)
+
+    assert Pow.Plug.current_user(conn).id == user.id
+    assert metadata = conn.private[:pow_session_metadata]
+    assert metadata[:valid_totp_at]
+    assert metadata[:inserted_at]
+
+    :timer.sleep(100)
+    conn = run_plug(init_conn, config)
+
+    assert Pow.Plug.current_user(conn).id == user.id
+    assert metadata = conn.private[:pow_session_metadata]
+    assert metadata[:valid_totp_at]
+    refute metadata[:inserted_at]
+  end
+
   defp init_session_plug(conn) do
     conn
     |> Map.put(:secret_key_base, String.duplicate("abcdefghijklmnopqrstuvxyz0123456789", 2))
@@ -97,6 +141,7 @@ defmodule PhilomenaWeb.PowInvalidatedSessionPlugTest do
     |> Session.call(Session.init(config))
     |> Cookie.call(Cookie.init([]))
     |> PowInvalidatedSessionPlug.call(PowInvalidatedSessionPlug.init(:load))
+    |> PhilomenaWeb.TotpPlug.call(PhilomenaWeb.TotpPlug.init([]))
   end
 
   defp run_plug(conn, config \\ @config) do