From 5cf951e908d30214a62152564416a3f0482c651a Mon Sep 17 00:00:00 2001
From: Laura Hausmann <laura@hausmann.dev>
Date: Thu, 21 Nov 2024 22:04:13 +0100
Subject: [PATCH] [backend/razor] Add TOTP 2FA support to OAuth page

---
 .../Pages/OAuth/Authorize.cshtml              | 130 +++++++++++++-----
 .../Pages/OAuth/Authorize.cshtml.cs           |  49 ++++++-
 .../Pages/OAuth/Authorize.cshtml.css          |  11 ++
 3 files changed, 155 insertions(+), 35 deletions(-)

diff --git a/Iceshrimp.Backend/Pages/OAuth/Authorize.cshtml b/Iceshrimp.Backend/Pages/OAuth/Authorize.cshtml
index a91fe33b..9798fe93 100644
--- a/Iceshrimp.Backend/Pages/OAuth/Authorize.cshtml
+++ b/Iceshrimp.Backend/Pages/OAuth/Authorize.cshtml
@@ -17,52 +17,114 @@
             }
         </ul>
 
-        <form method="post">
-            <div class="margin-bottom-5px">
-                Feature flags:
-            </div>
-            <div>
-                <input type="checkbox" name="supportsHtmlFormatting" id="supportsHtmlFormatting" value="1"/>
-                <label for="supportsHtmlFormatting">This app supports HTML formatting</label>
-            </div>
-            <div>
-                <input type="checkbox" name="autoDetectQuotes" id="autoDetectQuotes" value="1"/>
-                <label for="autoDetectQuotes">Automatically detect quotes</label>
-            </div>
-            <div>
-                <input type="checkbox" name="isPleroma" id="isPleroma" value="1"/>
-                <label for="isPleroma">This app is intended for Pleroma or Akkoma</label>
-            </div>
-            @if (Model.AuthenticatedUsers.Count > 0)
-            {
-                <div class="margin-top-5px">
-                    @foreach (var user in Model.AuthenticatedUsers)
+        @if (Model.TwoFactorFormData is { } data)
+        {
+            <form method="post">
+                <div class="margin-bottom-5px">
+                    Feature flags:
+                </div>
+                <div>
+                    @if (data.SupportsHtmlFormatting)
                     {
-                        <button type="submit" name="userId" value="@user.Id">Log in as @@@user.Username</button>
+                        <input type="checkbox" id="supportsHtmlFormatting" checked disabled/>
+                        <input type="hidden" name="supportsHtmlFormatting" value="1"/>
                     }
+                    else
+                    {
+                        <input type="checkbox" id="supportsHtmlFormatting" disabled/>
+                    }
+                    <label for="supportsHtmlFormatting">This app supports HTML formatting</label>
                 </div>
-                <div class="margin-bottom-5px margin-top-5px">
-                    Alternatively, sign in with to a different account below:
+                <div>
+                    @if (data.AutoDetectQuotes)
+                    {
+                        <input type="checkbox" id="autoDetectQuotes" checked disabled/>
+                        <input type="hidden" name="autoDetectQuotes" value="1"/>
+                    }
+                    else
+                    {
+                        <input type="checkbox" id="autoDetectQuotes" disabled/>
+                    }
+                    <label for="autoDetectQuotes">Automatically detect quotes</label>
+                </div>
+                <div>
+                    @if (data.IsPleroma)
+                    {
+                        <input type="checkbox" id="isPleroma" checked disabled/>
+                        <input type="hidden" name="isPleroma" value="1"/>
+                    }
+                    else
+                    {
+                        <input type="checkbox" id="isPleroma" disabled/>
+                    }
+                    <label for="isPleroma">This app is intended for Pleroma or Akkoma</label>
                 </div>
-                <input type="text" placeholder="Username" name="username"/>
-                <input type="password" placeholder="Password" name="password"/>
-            }
-            else
-            {
                 <div class="margin-bottom-5px margin-top-10px">
                     Log in below to confirm this:
                 </div>
-                <input type="text" placeholder="Username" name="username" required/>
-                <input type="password" placeholder="Password" name="password" required/>
-            }
-            <button type="submit">Submit</button>
-        </form>
+                <input type="text" disabled value="@data.Username"
+                       autocomplete="username"/>
+                <input type="hidden" name="username" value="@data.Username"/>
+                <input type="password" disabled value="@data.Password"/>
+                <input type="hidden" name="password" value="@data.Password"/>
+                <input type="text" inputmode="numeric" pattern="[0-9]{6}" autocomplete="one-time-code"
+                       placeholder="TOTP" name="totp" required/>
+                <button type="submit">Submit</button>
+            </form>
+        }
+        else
+        {
+            <form method="post">
+                <div class="margin-bottom-5px">
+                    Feature flags:
+                </div>
+                <div>
+                    <input type="checkbox" name="supportsHtmlFormatting" id="supportsHtmlFormatting" value="1"/>
+                    <label for="supportsHtmlFormatting">This app supports HTML formatting</label>
+                </div>
+                <div>
+                    <input type="checkbox" name="autoDetectQuotes" id="autoDetectQuotes" value="1"/>
+                    <label for="autoDetectQuotes">Automatically detect quotes</label>
+                </div>
+                <div>
+                    <input type="checkbox" name="isPleroma" id="isPleroma" value="1"/>
+                    <label for="isPleroma">This app is intended for Pleroma or Akkoma</label>
+                </div>
+                @if (Model.AuthenticatedUsers.Count > 0)
+                {
+                    <div class="margin-top-5px">
+                        @foreach (var user in Model.AuthenticatedUsers)
+                        {
+                            <button type="submit" name="userId" value="@user.Id">Log in as @@@user.Username</button>
+                        }
+                    </div>
+                    <div class="margin-bottom-5px margin-top-5px">
+                        Alternatively, sign in with to a different account below:
+                    </div>
+                    <input type="text" placeholder="Username" name="username" autocomplete="username"/>
+                    <input type="password" placeholder="Password" name="password" autocomplete="current-password"/>
+                }
+                else
+                {
+                    <div class="margin-bottom-5px margin-top-10px">
+                        Log in below to confirm this:
+                    </div>
+                    <input type="text" placeholder="Username" name="username" autocomplete="username" required/>
+                    <input type="password" placeholder="Password" name="password" autocomplete="current-password"
+                           required/>
+                }
+                <input type="text" autocomplete="one-time-code" placeholder="TOTP" name="totp" class="hidden-input"
+                       tabindex="-1" aria-hidden="true"/>
+                <button type="submit">Submit</button>
+            </form>
+        }
     </div>
 }
 else if (Model.Token.RedirectUri == "urn:ietf:wg:oauth:2.0:oob")
 {
     <div>
-        Your code is: <pre>@Model.Token.Code</pre>
+        Your code is:
+        <pre>@Model.Token.Code</pre>
     </div>
 }
 else
diff --git a/Iceshrimp.Backend/Pages/OAuth/Authorize.cshtml.cs b/Iceshrimp.Backend/Pages/OAuth/Authorize.cshtml.cs
index e82b9d61..9930223d 100644
--- a/Iceshrimp.Backend/Pages/OAuth/Authorize.cshtml.cs
+++ b/Iceshrimp.Backend/Pages/OAuth/Authorize.cshtml.cs
@@ -3,6 +3,7 @@ using Iceshrimp.Backend.Core.Database;
 using Iceshrimp.Backend.Core.Database.Tables;
 using Iceshrimp.Backend.Core.Helpers;
 using Iceshrimp.Backend.Core.Middleware;
+using Iceshrimp.Shared.Schemas.Web;
 using Microsoft.AspNetCore.Mvc;
 using Microsoft.AspNetCore.Mvc.RazorPages;
 using Microsoft.EntityFrameworkCore;
@@ -30,6 +31,8 @@ public class AuthorizeModel(DatabaseContext db) : PageModel
 		[MemberNotNull(nameof(Scopes))] set => Scopes = value.Split(' ').ToList();
 	}
 
+	public LoginData? TwoFactorFormData = null;
+
 	public async Task OnGet()
 	{
 		if (ResponseType == null || ClientId == null || RedirectUri == null)
@@ -53,7 +56,7 @@ public class AuthorizeModel(DatabaseContext db) : PageModel
 	}
 
 	public async Task OnPost(
-		[FromForm] string? username, [FromForm] string? password, [FromForm] string? userId,
+		[FromForm] string? username, [FromForm] string? password, [FromForm] string? totp, [FromForm] string? userId,
 		[FromForm] bool supportsHtmlFormatting, [FromForm] bool autoDetectQuotes, [FromForm] bool isPleroma
 	)
 	{
@@ -77,6 +80,25 @@ public class AuthorizeModel(DatabaseContext db) : PageModel
 				throw Forbidden();
 			if (AuthHelpers.ComparePassword(password, userSettings.Password) == false)
 				throw Forbidden();
+
+			if (userSettings.TwoFactorEnabled)
+			{
+				if (!string.IsNullOrWhiteSpace(totp))
+				{
+					if (userSettings.TwoFactorSecret == null)
+						throw new Exception("2FA is enabled but secret is null");
+					if (!TotpHelper.Validate(userSettings.TwoFactorSecret, totp))
+					{
+						SetTwoFactorFormData();
+						return;
+					}
+				}
+				else
+				{
+					SetTwoFactorFormData();
+					return;
+				}
+			}
 		}
 
 		var token = new OauthToken
@@ -99,5 +121,30 @@ public class AuthorizeModel(DatabaseContext db) : PageModel
 		await db.SaveChangesAsync();
 
 		Token = token;
+		return;
+
+		void SetTwoFactorFormData()
+		{
+			TwoFactorFormData = new LoginData
+			{
+				Username               = username,
+				Password               = password,
+				AutoDetectQuotes       = autoDetectQuotes,
+				SupportsHtmlFormatting = supportsHtmlFormatting,
+				IsPleroma              = isPleroma
+			};
+
+			Response.Headers.CacheControl = "private, no-store, no-cache";
+			Response.Headers.Pragma       = "no-cache";
+		}
+	}
+
+	public class LoginData
+	{
+		public required string Username;
+		public required string Password;
+		public required bool   SupportsHtmlFormatting;
+		public required bool   AutoDetectQuotes;
+		public required bool   IsPleroma;
 	}
 }
\ No newline at end of file
diff --git a/Iceshrimp.Backend/Pages/OAuth/Authorize.cshtml.css b/Iceshrimp.Backend/Pages/OAuth/Authorize.cshtml.css
index 656e0dc2..5f99e055 100644
--- a/Iceshrimp.Backend/Pages/OAuth/Authorize.cshtml.css
+++ b/Iceshrimp.Backend/Pages/OAuth/Authorize.cshtml.css
@@ -12,4 +12,15 @@
 
 .margin-top-10px {
     margin-top: 10px;
+}
+
+.hidden-input {
+    position: absolute;
+    top: -10px;
+    left: -10px;
+    height: 0;
+    width: 0;
+    padding: 0;
+    margin: 0;
+    pointer-events: none;
 }
\ No newline at end of file