From 2d47243c12c10879787a25a1d149b6d56b79ba03 Mon Sep 17 00:00:00 2001
From: syeopite <syeopite@syeopite.dev>
Date: Fri, 16 Jul 2021 14:25:50 -0700
Subject: [PATCH] Document 2fa functions

---
 src/invidious/helpers/utils.cr   |  4 ++++
 src/invidious/routes/accounts.cr | 26 +++++++++++++++++++-------
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/src/invidious/helpers/utils.cr b/src/invidious/helpers/utils.cr
index 5210cf73..4abd3502 100644
--- a/src/invidious/helpers/utils.cr
+++ b/src/invidious/helpers/utils.cr
@@ -472,6 +472,10 @@ def fetch_random_instance
   return filtered_instance_list.sample(1)[0]
 end
 
+# Templates the 2fa validator page.
+#
+# Requires the env, user, sid and locale variables for
+# generating a csrf_token and the required variables for the view.
 def call_totp_validator(env, user, sid, locale)
   referer = URI.decode_www_form(env.get?("current_page").to_s)
   csrf_token = generate_response(sid, {":2fa/validate"}, HMAC_KEY, PG_DB)
diff --git a/src/invidious/routes/accounts.cr b/src/invidious/routes/accounts.cr
index 55c2fc4c..7047c40f 100644
--- a/src/invidious/routes/accounts.cr
+++ b/src/invidious/routes/accounts.cr
@@ -3,7 +3,7 @@ require "./base_route"
 
 # Different routes relating to existing accounts and the control of their data.
 class Invidious::Routes::Accounts < Invidious::Routes::BaseRoute
-  # Setup 2fa page
+  # Templates the page to setup 2fa on an user account
   def setup_2fa_page(env)
     locale = LOCALES[env.get("preferences").as(Preferences).locale]?
 
@@ -11,6 +11,10 @@ class Invidious::Routes::Accounts < Invidious::Routes::BaseRoute
     sid = env.get? "sid"
     referer = get_referer(env, unroll: false)
 
+    if !user
+      return env.redirect referer
+    end
+
     user = user.as(User)
     sid = sid.as(String)
     csrf_token = generate_response(sid, {":2fa/setup"}, HMAC_KEY, PG_DB)
@@ -22,19 +26,27 @@ class Invidious::Routes::Accounts < Invidious::Routes::BaseRoute
     return templated "account/setup_2fa"
   end
 
-  # Remove 2fa page
+  # Templates the page to remove 2fa on an user account
   def remove_2fa_page(env)
     locale = LOCALES[env.get("preferences").as(Preferences).locale]?
     referer = get_referer(env)
 
-    user = env.get("user").as(User)
-    sid = env.get("sid").as(String)
+    user = env.get? "user"
+    sid = env.get? "sid"
+    referer = get_referer(env, unroll: false)
+
+    if !user
+      return env.redirect referer
+    end
+
+    user = user.as(User)
+    sid = sid.as(String)
     csrf_token = generate_response(sid, {":2fa/remove"}, HMAC_KEY, PG_DB)
 
     return templated "account/remove_2fa"
   end
 
-  # Remove 2fa post request.
+  # Handles requests to remove 2fa on an user account
   def remove_2fa(env)
     locale = LOCALES[env.get("preferences").as(Preferences).locale]?
 
@@ -59,7 +71,7 @@ class Invidious::Routes::Accounts < Invidious::Routes::BaseRoute
     PG_DB.exec("UPDATE users SET totp_secret = $1 WHERE email = $2", nil, user.email)
   end
 
-  # Setup 2fa post request.
+  # Handles requests to setup 2fa on an user account
   def setup_2fa(env)
     locale = LOCALES[env.get("preferences").as(Preferences).locale]?
 
@@ -96,7 +108,7 @@ class Invidious::Routes::Accounts < Invidious::Routes::BaseRoute
     env.redirect referer
   end
 
-  # Validate 2fa code endpoint
+  # Handles requests to validate a TOTP code on an user account
   def validate_2fa(env)
     locale = LOCALES[env.get("preferences").as(Preferences).locale]?
     referer = get_referer(env, unroll: false)