From ac8bf0f3dcb915749f12c887c95f9ab855553ffe Mon Sep 17 00:00:00 2001
From: Ike Kottlowski <ikottlowski@bitwarden.com>
Date: Fri, 21 Mar 2025 10:50:04 -0400
Subject: [PATCH] feat : add feature flag to grant validator; fix : authed user
 flag stays in sessions for 5 minutes to account for 2FA

---
 .../OpaqueKeyExchangeService.cs               |  6 +++++-
 .../OpaqueKeyExchangeGrantValidator.cs        | 20 ++++++++++---------
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/Core/Auth/Services/Implementations/OpaqueKeyExchangeService.cs b/src/Core/Auth/Services/Implementations/OpaqueKeyExchangeService.cs
index d6e5ffc915..60722690d8 100644
--- a/src/Core/Auth/Services/Implementations/OpaqueKeyExchangeService.cs
+++ b/src/Core/Auth/Services/Implementations/OpaqueKeyExchangeService.cs
@@ -197,7 +197,11 @@ public class OpaqueKeyExchangeService : IOpaqueKeyExchangeService
             await _distributedCache.SetAsync(
                 string.Format(LOGIN_SESSION_KEY, sessionId),
                 Encoding.ASCII.GetBytes(JsonSerializer.Serialize(loginSession)),
-                _distributedCacheEntryOptions);
+                new DistributedCacheEntryOptions()
+                {
+                    // Our login sessions are 5 minutes long so if a user needs to accomplish 2FA this ensures the user has time to do so.
+                    AbsoluteExpirationRelativeToNow = TimeSpan.FromMinutes(5)
+                });
 
             return true;
         }
diff --git a/src/Identity/IdentityServer/RequestValidators/OpaqueKeyExchangeGrantValidator.cs b/src/Identity/IdentityServer/RequestValidators/OpaqueKeyExchangeGrantValidator.cs
index eb12eeeaac..5016d1f4bd 100644
--- a/src/Identity/IdentityServer/RequestValidators/OpaqueKeyExchangeGrantValidator.cs
+++ b/src/Identity/IdentityServer/RequestValidators/OpaqueKeyExchangeGrantValidator.cs
@@ -1,7 +1,6 @@
 using System.Security.Claims;
 using Bit.Core;
 using Bit.Core.AdminConsole.Services;
-using Bit.Core.Auth.Models.Business.Tokenables;
 using Bit.Core.Auth.Repositories;
 using Bit.Core.Auth.Services;
 using Bit.Core.Context;
@@ -9,7 +8,6 @@ using Bit.Core.Entities;
 using Bit.Core.Repositories;
 using Bit.Core.Services;
 using Bit.Core.Settings;
-using Bit.Core.Tokens;
 using Duende.IdentityServer.Models;
 using Duende.IdentityServer.Validation;
 using Microsoft.AspNetCore.Identity;
@@ -19,8 +17,8 @@ namespace Bit.Identity.IdentityServer.RequestValidators;
 public class OpaqueKeyExchangeGrantValidator : BaseRequestValidator<ExtensionGrantValidationContext>, IExtensionGrantValidator
 {
     public const string GrantType = "opaque-ke";
-    private IUserRepository userRepository;
-    private IOpaqueKeyExchangeService opaqueKeyExchangeService;
+    private readonly IOpaqueKeyExchangeService _opaqueKeyExchangeService;
+    private readonly IFeatureService _featureService;
 
     public OpaqueKeyExchangeGrantValidator(
         UserManager<User> userManager,
@@ -36,7 +34,6 @@ public class OpaqueKeyExchangeGrantValidator : BaseRequestValidator<ExtensionGra
         ISsoConfigRepository ssoConfigRepository,
         IUserRepository userRepository,
         IPolicyService policyService,
-        IDataProtectorTokenFactory<WebAuthnLoginAssertionOptionsTokenable> assertionOptionsDataProtector,
         IFeatureService featureService,
         IUserDecryptionOptionsBuilder userDecryptionOptionsBuilder,
         IOpaqueKeyExchangeService opaqueKeyExchangeService)
@@ -57,14 +54,19 @@ public class OpaqueKeyExchangeGrantValidator : BaseRequestValidator<ExtensionGra
             ssoConfigRepository,
             userDecryptionOptionsBuilder)
     {
-        this.userRepository = userRepository;
-        this.opaqueKeyExchangeService = opaqueKeyExchangeService;
+        _opaqueKeyExchangeService = opaqueKeyExchangeService;
     }
 
     string IExtensionGrantValidator.GrantType => "opaque-ke";
 
     public async Task ValidateAsync(ExtensionGrantValidationContext context)
     {
+        if (!_featureService.IsEnabled(FeatureFlagKeys.OpaqueKeyExchange))
+        {
+            context.Result = new GrantValidationResult(TokenRequestErrors.InvalidGrant);
+            return;
+        }
+
         var sessionId = context.Request.Raw.Get("sessionId");
         if (string.IsNullOrWhiteSpace(sessionId))
         {
@@ -72,7 +74,7 @@ public class OpaqueKeyExchangeGrantValidator : BaseRequestValidator<ExtensionGra
             return;
         }
 
-        var user = await opaqueKeyExchangeService.GetUserForAuthenticatedSession(Guid.Parse(sessionId));
+        var user = await _opaqueKeyExchangeService.GetUserForAuthenticatedSession(Guid.Parse(sessionId));
         if (user == null)
         {
             context.Result = new GrantValidationResult(TokenRequestErrors.InvalidGrant);
@@ -104,7 +106,7 @@ public class OpaqueKeyExchangeGrantValidator : BaseRequestValidator<ExtensionGra
             identityProvider: Constants.IdentityProvider,
             claims: claims.Count > 0 ? claims : null,
             customResponse: customResponse);
-        await opaqueKeyExchangeService.ClearAuthenticationSession(Guid.Parse(context.Request.Raw.Get("sessionId")));
+        await _opaqueKeyExchangeService.ClearAuthenticationSession(Guid.Parse(context.Request.Raw.Get("sessionId")));
     }
 
     protected override ClaimsPrincipal GetSubject(ExtensionGrantValidationContext context)