From 9ce6e01be17e5a95a2351c9c906d344340fc04ca Mon Sep 17 00:00:00 2001
From: Jared Goodwin <translucency@outlook.com>
Date: Fri, 31 May 2024 13:17:03 -0700
Subject: [PATCH] Fix two factor enforcement.

---
 Server/Auth/TwoFactorRequiredHandler.cs   | 42 +++++++++++++----------
 Server/Components/AuthorizedIndex.razor   | 11 ------
 Server/Components/Layout/MainLayout.razor | 12 +++++--
 3 files changed, 34 insertions(+), 31 deletions(-)

diff --git a/Server/Auth/TwoFactorRequiredHandler.cs b/Server/Auth/TwoFactorRequiredHandler.cs
index 0d62fb79..47fb2dfe 100644
--- a/Server/Auth/TwoFactorRequiredHandler.cs
+++ b/Server/Auth/TwoFactorRequiredHandler.cs
@@ -1,31 +1,21 @@
 using Microsoft.AspNetCore.Authorization;
-using Microsoft.AspNetCore.Identity;
+using Remotely.Server.Models;
 using Remotely.Server.Services;
-using Remotely.Shared.Entities;
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Threading.Tasks;
+using System.Security.Principal;
 
 namespace Remotely.Server.Auth;
 
-public class TwoFactorRequiredHandler : AuthorizationHandler<TwoFactorRequiredRequirement>
+public class TwoFactorRequiredHandler(
+    IHttpContextAccessor _contextAccessor,
+    IDataService _dataService) : AuthorizationHandler<TwoFactorRequiredRequirement>
 {
-    private readonly IDataService _dataService;
-
-    public TwoFactorRequiredHandler(IDataService dataService)
-    {
-        _dataService = dataService;
-    }
-
     protected override async Task HandleRequirementAsync(AuthorizationHandlerContext context, TwoFactorRequiredRequirement requirement)
     {
         var settings = await _dataService.GetSettings();
-        if (context.User.Identity?.IsAuthenticated == true &&
-            context.User.Identity.Name is not null &&
-            settings.Require2FA)
+        if (context.User?.Identity is { } identity &&
+            IsTwoFactorRequired(identity, settings))
         {
-            var userResult = await _dataService.GetUserByName(context.User.Identity.Name);
+            var userResult = await _dataService.GetUserByName(identity.Name!);
 
             if (!userResult.IsSuccess ||
                 !userResult.Value.TwoFactorEnabled)
@@ -36,4 +26,20 @@ public class TwoFactorRequiredHandler : AuthorizationHandler<TwoFactorRequiredRe
         }
         context.Succeed(requirement);
     }
+
+    private bool IsTwoFactorRequired(IIdentity identity, SettingsModel settings)
+    {
+        // Account management pages are exempt since they're required
+        // to set up 2FA.
+        var path = _contextAccessor.HttpContext?.Request.Path ?? "";
+        if (path.StartsWithSegments("/Account/Manage"))
+        {
+            return false;
+        }
+
+        return 
+            settings.Require2FA &&
+            identity.IsAuthenticated &&
+            identity.Name is not null;
+    }
 }
diff --git a/Server/Components/AuthorizedIndex.razor b/Server/Components/AuthorizedIndex.razor
index 2c8a1f36..20eb4ed9 100644
--- a/Server/Components/AuthorizedIndex.razor
+++ b/Server/Components/AuthorizedIndex.razor
@@ -19,17 +19,6 @@
 @code {
     private SettingsModel? _settings;
 
-    protected override async Task OnAfterRenderAsync(bool firstRender)
-    {
-        if (User is not null &&
-            _settings?.Require2FA == true &&
-            !User.TwoFactorEnabled)
-        {
-            NavManager.NavigateTo("/TwoFactorRequired");
-        }
-        await base.OnAfterRenderAsync(firstRender);
-    }
-
     protected override async Task OnInitializedAsync()
     {
         _settings = await DataService.GetSettings();
diff --git a/Server/Components/Layout/MainLayout.razor b/Server/Components/Layout/MainLayout.razor
index c8a88b43..c4acb02f 100644
--- a/Server/Components/Layout/MainLayout.razor
+++ b/Server/Components/Layout/MainLayout.razor
@@ -1,6 +1,7 @@
 @using Remotely.Server.Components 
 @using Remotely.Server.Auth 
 @inherits LayoutComponentBase
+@inject NavigationManager NavMan
 
 <AuthorizeView Policy="@PolicyNames.TwoFactorRequired">
     <Authorized>
@@ -25,7 +26,7 @@
                 <div class="col-sm-12">
                     <p>Two-factor authentication is required.  Click the button below to set up your authenticator app.</p>
                     <p>
-                        <a href="Account/Manage/TwoFactorAuthentication" class="btn btn-primary">Enable 2FA</a>
+                        <button class="btn btn-primary" @onclick="NavigateToTwoFactor">Enable 2FA</button>
                     </p>
                 </div>
             </div>
@@ -42,4 +43,11 @@
 
 <ToastHarness />
 <LoaderHarness />
-<ModalHarness />
\ No newline at end of file
+<ModalHarness />
+
+@code {
+    private void NavigateToTwoFactor()
+    {
+        NavMan.NavigateTo("/Account/Manage/TwoFactorAuthentication", true);
+    }
+}
\ No newline at end of file