From 7cfee1dc94f9a0bbf61ff294a4cd5fc659de16ed Mon Sep 17 00:00:00 2001
From: syeopite <syeopite@syeopite.dev>
Date: Thu, 15 Jul 2021 01:27:27 -0700
Subject: [PATCH] Fixes + add 2fa to pass change and acc delete

---
 src/invidious.cr                 | 10 +++++++---
 src/invidious/helpers/utils.cr   |  7 +++++++
 src/invidious/routes/accounts.cr |  7 ++++++-
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/src/invidious.cr b/src/invidious.cr
index dbbf67a0..10ae25e9 100644
--- a/src/invidious.cr
+++ b/src/invidious.cr
@@ -858,9 +858,8 @@ get "/change_password" do |env|
 
   user = user.as(User)
   sid = sid.as(String)
-  if user.totp_secret && env.response.cookies["2faVerified"]?.try &.value != "1" || nil
-    csrf_token = generate_response(sid, {":validate_2fa"}, HMAC_KEY, PG_DB)
-    next templated "account/validate_2fa?referer=#{env.get?("current_page")}"
+  if user.totp_secret && env.request.cookies["2faVerified"]?.try &.value != "1" || nil
+    next call_totp_validator(env, user, sid, locale)
   end
 
   csrf_token = generate_response(sid, {":change_password"}, HMAC_KEY, PG_DB)
@@ -937,6 +936,11 @@ get "/delete_account" do |env|
 
   user = user.as(User)
   sid = sid.as(String)
+
+  if user.totp_secret && env.request.cookies["2faVerified"]?.try &.value != "1" || nil
+    next call_totp_validator(env, user, sid, locale)
+  end
+
   csrf_token = generate_response(sid, {":delete_account"}, HMAC_KEY, PG_DB)
 
   templated "account/delete_account"
diff --git a/src/invidious/helpers/utils.cr b/src/invidious/helpers/utils.cr
index 9b667b1e..ebf8e185 100644
--- a/src/invidious/helpers/utils.cr
+++ b/src/invidious/helpers/utils.cr
@@ -546,3 +546,10 @@ def totp_validator(env)
     end
   end
 end
+
+def call_totp_validator(env, user, sid, locale)
+  referer = URI.decode_www_form(env.get?("current_page").to_s)
+  csrf_token = generate_response(sid, {":validate_2fa"}, HMAC_KEY, PG_DB)
+  email, password = {user.email, nil}
+  return templated "account/validate_2fa"
+end
diff --git a/src/invidious/routes/accounts.cr b/src/invidious/routes/accounts.cr
index 479d1c46..1715bd76 100644
--- a/src/invidious/routes/accounts.cr
+++ b/src/invidious/routes/accounts.cr
@@ -29,6 +29,8 @@ class Invidious::Routes::Accounts < Invidious::Routes::BaseRoute
     sid = env.get? "sid"
     referer = get_referer(env, unroll: false)
 
+    puts referer
+
     if !user
       return env.redirect referer
     end
@@ -60,11 +62,12 @@ class Invidious::Routes::Accounts < Invidious::Routes::BaseRoute
   # Validate 2fa code endpoint
   def validate_2fa(env)
     locale = LOCALES[env.get("preferences").as(Preferences).locale]?
-    referer = get_referer(env)
+    referer = get_referer(env, unroll: false)
 
     email = env.params.body["email"]?.try &.downcase.byte_slice(0, 254)
     password = env.params.body["password"]?
     totp_code = env.params.body["totp_code"]?
+
     # This endpoint is only called when the user has a totp_secret.
     user = PG_DB.query_one?("SELECT * FROM users WHERE email = $1", email, as: User).not_nil!
 
@@ -131,5 +134,7 @@ class Invidious::Routes::Accounts < Invidious::Routes::BaseRoute
         env.response.cookies["2faVerified"] = HTTP::Cookie.new(name: "2faVerified", value: "1", expires: Time.utc + 1.hours, secure: secure, http_only: true)
       end
     end
+
+    env.redirect referer
   end
 end