From e0b231a2203736f8ae71b42d986c9b2f61e7ddef Mon Sep 17 00:00:00 2001
From: Justin Baur <19896123+justindbaur@users.noreply.github.com>
Date: Mon, 26 Jun 2023 20:17:39 -0400
Subject: [PATCH] [PM-2697] Return `UserDecryptionOptions` Always (#3032)

* Add Comments to UserDecryptionOptions

* Move Feature Flag Check

* Remove SSO Config Check

* Move UserDecryptionOptions Creation

- Put logic in BaseRequestValidator

* Remove 'async'
---
 .../Api/Response/UserDecryptionOptions.cs     |  6 +-
 .../IdentityServer/BaseRequestValidator.cs    | 63 ++++++++++++-
 .../CustomTokenRequestValidator.cs            | 94 ++++---------------
 .../ResourceOwnerPasswordValidator.cs         | 15 ++-
 .../Endpoints/IdentityServerSsoTests.cs       | 57 ++++++++++-
 .../Endpoints/IdentityServerTests.cs          | 11 +++
 6 files changed, 161 insertions(+), 85 deletions(-)

diff --git a/src/Core/Auth/Models/Api/Response/UserDecryptionOptions.cs b/src/Core/Auth/Models/Api/Response/UserDecryptionOptions.cs
index 28f876035..8cf2586ce 100644
--- a/src/Core/Auth/Models/Api/Response/UserDecryptionOptions.cs
+++ b/src/Core/Auth/Models/Api/Response/UserDecryptionOptions.cs
@@ -12,18 +12,18 @@ public class UserDecryptionOptions : ResponseModel
     }
 
     /// <summary>
-    /// 
+    /// Gets or sets whether the current user has a master password that can be used to decrypt their vault.
     /// </summary>
     public bool HasMasterPassword { get; set; }
 
     /// <summary>
-    /// 
+    /// Gets or sets information regarding this users trusted device decryption setup.
     /// </summary>
     [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
     public TrustedDeviceUserDecryptionOption? TrustedDeviceOption { get; set; }
 
     /// <summary>
-    /// 
+    /// Gets or set information about the current users KeyConnector setup.
     /// </summary>
     [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
     public KeyConnectorUserDecryptionOption? KeyConnectorOption { get; set; }
diff --git a/src/Identity/IdentityServer/BaseRequestValidator.cs b/src/Identity/IdentityServer/BaseRequestValidator.cs
index 271b9769c..9cfc19c13 100644
--- a/src/Identity/IdentityServer/BaseRequestValidator.cs
+++ b/src/Identity/IdentityServer/BaseRequestValidator.cs
@@ -6,7 +6,10 @@ using Bit.Core;
 using Bit.Core.Auth.Enums;
 using Bit.Core.Auth.Identity;
 using Bit.Core.Auth.Models;
+using Bit.Core.Auth.Models.Api.Response;
 using Bit.Core.Auth.Models.Business.Tokenables;
+using Bit.Core.Auth.Models.Data;
+using Bit.Core.Auth.Repositories;
 using Bit.Core.Context;
 using Bit.Core.Entities;
 using Bit.Core.Enums;
@@ -43,6 +46,8 @@ public abstract class BaseRequestValidator<T> where T : class
 
     protected ICurrentContext CurrentContext { get; }
     protected IPolicyService PolicyService { get; }
+    protected IFeatureService FeatureService { get; }
+    protected ISsoConfigRepository SsoConfigRepository { get; }
 
     public BaseRequestValidator(
         UserManager<User> userManager,
@@ -58,10 +63,11 @@ public abstract class BaseRequestValidator<T> where T : class
         ILogger logger,
         ICurrentContext currentContext,
         GlobalSettings globalSettings,
-        IPolicyRepository policyRepository,
         IUserRepository userRepository,
         IPolicyService policyService,
-        IDataProtectorTokenFactory<SsoEmail2faSessionTokenable> tokenDataFactory)
+        IDataProtectorTokenFactory<SsoEmail2faSessionTokenable> tokenDataFactory,
+        IFeatureService featureService,
+        ISsoConfigRepository ssoConfigRepository)
     {
         _userManager = userManager;
         _deviceRepository = deviceRepository;
@@ -79,6 +85,8 @@ public abstract class BaseRequestValidator<T> where T : class
         PolicyService = policyService;
         _userRepository = userRepository;
         _tokenDataFactory = tokenDataFactory;
+        FeatureService = featureService;
+        SsoConfigRepository = ssoConfigRepository;
     }
 
     protected async Task ValidateAsync(T context, ValidatedTokenRequest request,
@@ -199,6 +207,7 @@ public abstract class BaseRequestValidator<T> where T : class
         customResponse.Add("KdfIterations", user.KdfIterations);
         customResponse.Add("KdfMemory", user.KdfMemory);
         customResponse.Add("KdfParallelism", user.KdfParallelism);
+        customResponse.Add("UserDecryptionOptions", await CreateUserDecryptionOptionsAsync(user, GetSubject(context)));
 
         if (sendRememberToken)
         {
@@ -300,6 +309,7 @@ public abstract class BaseRequestValidator<T> where T : class
         Dictionary<string, object> customResponse);
 
     protected abstract void SetErrorResult(T context, Dictionary<string, object> customResponse);
+    protected abstract ClaimsPrincipal GetSubject(T context);
 
     private async Task<Tuple<bool, Organization>> RequiresTwoFactorAsync(User user, ValidatedTokenRequest request)
     {
@@ -572,4 +582,53 @@ public abstract class BaseRequestValidator<T> where T : class
 
         return new MasterPasswordPolicyResponseModel(await PolicyService.GetMasterPasswordPolicyForUserAsync(user));
     }
+
+#nullable enable
+    /// <summary>
+    /// Used to create a list of all possible ways the newly authenticated user can decrypt their vault contents
+    /// </summary>
+    private async Task<UserDecryptionOptions> CreateUserDecryptionOptionsAsync(User user, ClaimsPrincipal subject)
+    {
+        var ssoConfigurationData = await GetSsoConfigurationDataAsync(subject);
+
+        var userDecryptionOption = new UserDecryptionOptions
+        {
+            HasMasterPassword = !string.IsNullOrEmpty(user.MasterPassword)
+        };
+
+        if (ssoConfigurationData is { MemberDecryptionType: MemberDecryptionType.KeyConnector } && !string.IsNullOrEmpty(ssoConfigurationData.KeyConnectorUrl))
+        {
+            // KeyConnector makes it mutually exclusive
+            userDecryptionOption.KeyConnectorOption = new KeyConnectorUserDecryptionOption(ssoConfigurationData.KeyConnectorUrl);
+            return userDecryptionOption;
+        }
+
+        // Only add the trusted device specific option when the flag is turned on
+        if (FeatureService.IsEnabled(FeatureFlagKeys.TrustedDeviceEncryption, CurrentContext) && ssoConfigurationData is { MemberDecryptionType: MemberDecryptionType.TrustedDeviceEncryption })
+        {
+            var hasAdminApproval = await PolicyService.AnyPoliciesApplicableToUserAsync(user.Id, PolicyType.ResetPassword);
+            // TrustedDeviceEncryption only exists for SSO, but if that ever changes this value won't always be true
+            userDecryptionOption.TrustedDeviceOption = new TrustedDeviceUserDecryptionOption(hasAdminApproval);
+        }
+
+        return userDecryptionOption;
+    }
+
+    private async Task<SsoConfigurationData?> GetSsoConfigurationDataAsync(ClaimsPrincipal subject)
+    {
+        var organizationClaim = subject?.FindFirstValue("organizationId");
+
+        if (organizationClaim == null || !Guid.TryParse(organizationClaim, out var organizationId))
+        {
+            return null;
+        }
+
+        var ssoConfig = await SsoConfigRepository.GetByOrganizationIdAsync(organizationId);
+        if (ssoConfig == null)
+        {
+            return null;
+        }
+
+        return ssoConfig.GetData();
+    }
 }
diff --git a/src/Identity/IdentityServer/CustomTokenRequestValidator.cs b/src/Identity/IdentityServer/CustomTokenRequestValidator.cs
index 7096cb038..e6e39bade 100644
--- a/src/Identity/IdentityServer/CustomTokenRequestValidator.cs
+++ b/src/Identity/IdentityServer/CustomTokenRequestValidator.cs
@@ -1,14 +1,10 @@
 using System.Security.Claims;
-using Bit.Core;
-using Bit.Core.Auth.Enums;
 using Bit.Core.Auth.Identity;
 using Bit.Core.Auth.Models.Api.Response;
 using Bit.Core.Auth.Models.Business.Tokenables;
-using Bit.Core.Auth.Models.Data;
 using Bit.Core.Auth.Repositories;
 using Bit.Core.Context;
 using Bit.Core.Entities;
-using Bit.Core.Enums;
 using Bit.Core.IdentityServer;
 using Bit.Core.Repositories;
 using Bit.Core.Services;
@@ -27,8 +23,6 @@ public class CustomTokenRequestValidator : BaseRequestValidator<CustomTokenReque
     ICustomTokenRequestValidator
 {
     private readonly UserManager<User> _userManager;
-    private readonly ISsoConfigRepository _ssoConfigRepository;
-    private readonly IFeatureService _featureService;
 
     public CustomTokenRequestValidator(
         UserManager<User> userManager,
@@ -44,7 +38,6 @@ public class CustomTokenRequestValidator : BaseRequestValidator<CustomTokenReque
         ILogger<CustomTokenRequestValidator> logger,
         ICurrentContext currentContext,
         GlobalSettings globalSettings,
-        IPolicyRepository policyRepository,
         ISsoConfigRepository ssoConfigRepository,
         IUserRepository userRepository,
         IPolicyService policyService,
@@ -52,12 +45,10 @@ public class CustomTokenRequestValidator : BaseRequestValidator<CustomTokenReque
         IFeatureService featureService)
         : base(userManager, deviceRepository, deviceService, userService, eventService,
             organizationDuoWebTokenProvider, organizationRepository, organizationUserRepository,
-            applicationCacheService, mailService, logger, currentContext, globalSettings, policyRepository,
-            userRepository, policyService, tokenDataFactory)
+            applicationCacheService, mailService, logger, currentContext, globalSettings,
+            userRepository, policyService, tokenDataFactory, featureService, ssoConfigRepository)
     {
         _userManager = userManager;
-        _ssoConfigRepository = ssoConfigRepository;
-        _featureService = featureService;
     }
 
     public async Task ValidateAsync(CustomTokenRequestValidationContext context)
@@ -96,7 +87,7 @@ public class CustomTokenRequestValidator : BaseRequestValidator<CustomTokenReque
         return validatorContext.User != null;
     }
 
-    protected override async Task SetSuccessResult(CustomTokenRequestValidationContext context, User user,
+    protected override Task SetSuccessResult(CustomTokenRequestValidationContext context, User user,
         List<Claim> claims, Dictionary<string, object> customResponse)
     {
         context.Result.CustomResponse = customResponse;
@@ -110,23 +101,9 @@ public class CustomTokenRequestValidator : BaseRequestValidator<CustomTokenReque
             }
         }
 
-        // Attempts to find ssoConfigData for a given validate request subject
-        // this is actually guarenteed to pretty often be null, because more than just sso login requests will come
-        // through here
-        var ssoConfigData = await GetSsoConfigurationDataAsync(context.Result.ValidatedRequest.Subject);
-
-        // You can't put this below the user.MasterPassword != null check because TDE users can still have a MasterPassword
-        // It's worth noting that CurrentContext here will build a user in LaunchDarkly that is anonymous but DOES belong
-        // to an organization. So we will not be able to turn this feature on for only a single user, only for an entire 
-        // organization at a time.
-        if (ssoConfigData != null && _featureService.IsEnabled(FeatureFlagKeys.TrustedDeviceEncryption, CurrentContext))
-        {
-            context.Result.CustomResponse["UserDecryptionOptions"] = await CreateUserDecryptionOptionsAsync(ssoConfigData, user);
-        }
-
         if (context.Result.CustomResponse == null || user.MasterPassword != null)
         {
-            return;
+            return Task.CompletedTask;
         }
 
         // KeyConnector responses below
@@ -141,36 +118,30 @@ public class CustomTokenRequestValidator : BaseRequestValidator<CustomTokenReque
                 context.Result.CustomResponse["ResetMasterPassword"] = false;
             }
 
-            return;
+            return Task.CompletedTask;
         }
 
-        // SSO login
-        // This does a double check, that ssoConfigData is not null and that it has the KeyConnector member decryption type
-        if (ssoConfigData is { MemberDecryptionType: MemberDecryptionType.KeyConnector } && !string.IsNullOrEmpty(ssoConfigData.KeyConnectorUrl))
+        // Key connector data should have already been set in the decryption options
+        // for backwards compatibility we set them this way too. We can eventually get rid of this
+        // when all clients don't read them from the existing locations.
+        if (!context.Result.CustomResponse.TryGetValue("UserDecryptionOptions", out var userDecryptionOptionsObj) ||
+            userDecryptionOptionsObj is not UserDecryptionOptions userDecryptionOptions)
         {
-            // TODO: Can be removed in the future
-            context.Result.CustomResponse["KeyConnectorUrl"] = ssoConfigData.KeyConnectorUrl;
-            // Prevent clients redirecting to set-password
+            return Task.CompletedTask;
+        }
+
+        if (userDecryptionOptions is { KeyConnectorOption: { } })
+        {
+            context.Result.CustomResponse["KeyConnectorUrl"] = userDecryptionOptions.KeyConnectorOption.KeyConnectorUrl;
             context.Result.CustomResponse["ResetMasterPassword"] = false;
         }
+
+        return Task.CompletedTask;
     }
 
-    private async Task<SsoConfigurationData?> GetSsoConfigurationDataAsync(ClaimsPrincipal? subject)
+    protected override ClaimsPrincipal GetSubject(CustomTokenRequestValidationContext context)
     {
-        var organizationClaim = subject?.FindFirstValue("organizationId");
-
-        if (organizationClaim == null || !Guid.TryParse(organizationClaim, out var organizationId))
-        {
-            return null;
-        }
-
-        var ssoConfig = await _ssoConfigRepository.GetByOrganizationIdAsync(organizationId);
-        if (ssoConfig == null)
-        {
-            return null;
-        }
-
-        return ssoConfig.GetData();
+        return context.Result.ValidatedRequest.Subject;
     }
 
     protected override void SetTwoFactorResult(CustomTokenRequestValidationContext context,
@@ -198,29 +169,4 @@ public class CustomTokenRequestValidator : BaseRequestValidator<CustomTokenReque
         context.Result.IsError = true;
         context.Result.CustomResponse = customResponse;
     }
-
-    /// <summary>
-    /// Used to create a list of all possible ways the newly authenticated user can decrypt their vault contents
-    /// </summary>
-    private async Task<UserDecryptionOptions> CreateUserDecryptionOptionsAsync(SsoConfigurationData ssoConfigurationData, User user)
-    {
-        var userDecryptionOption = new UserDecryptionOptions();
-        if (ssoConfigurationData is { MemberDecryptionType: MemberDecryptionType.KeyConnector } && !string.IsNullOrEmpty(ssoConfigurationData.KeyConnectorUrl))
-        {
-            // KeyConnector makes it mutually exclusive
-            userDecryptionOption.KeyConnectorOption = new KeyConnectorUserDecryptionOption(ssoConfigurationData.KeyConnectorUrl);
-            return userDecryptionOption;
-        }
-
-        if (ssoConfigurationData is { MemberDecryptionType: MemberDecryptionType.TrustedDeviceEncryption })
-        {
-            var hasAdminApproval = await PolicyService.AnyPoliciesApplicableToUserAsync(user.Id, PolicyType.ResetPassword);
-            // TrustedDeviceEncryption only exists for SSO, but if that ever changes this value won't always be true
-            userDecryptionOption.TrustedDeviceOption = new TrustedDeviceUserDecryptionOption(hasAdminApproval);
-        }
-
-        userDecryptionOption.HasMasterPassword = !string.IsNullOrEmpty(user.MasterPassword);
-
-        return userDecryptionOption;
-    }
 }
diff --git a/src/Identity/IdentityServer/ResourceOwnerPasswordValidator.cs b/src/Identity/IdentityServer/ResourceOwnerPasswordValidator.cs
index c36f39b45..df17a474a 100644
--- a/src/Identity/IdentityServer/ResourceOwnerPasswordValidator.cs
+++ b/src/Identity/IdentityServer/ResourceOwnerPasswordValidator.cs
@@ -1,6 +1,7 @@
 using System.Security.Claims;
 using Bit.Core.Auth.Identity;
 using Bit.Core.Auth.Models.Business.Tokenables;
+using Bit.Core.Auth.Repositories;
 using Bit.Core.Auth.Services;
 using Bit.Core.Context;
 using Bit.Core.Entities;
@@ -37,16 +38,17 @@ public class ResourceOwnerPasswordValidator : BaseRequestValidator<ResourceOwner
         ILogger<ResourceOwnerPasswordValidator> logger,
         ICurrentContext currentContext,
         GlobalSettings globalSettings,
-        IPolicyRepository policyRepository,
         ICaptchaValidationService captchaValidationService,
         IAuthRequestRepository authRequestRepository,
         IUserRepository userRepository,
         IPolicyService policyService,
-        IDataProtectorTokenFactory<SsoEmail2faSessionTokenable> tokenDataFactory)
+        IDataProtectorTokenFactory<SsoEmail2faSessionTokenable> tokenDataFactory,
+        IFeatureService featureService,
+        ISsoConfigRepository ssoConfigRepository)
         : base(userManager, deviceRepository, deviceService, userService, eventService,
               organizationDuoWebTokenProvider, organizationRepository, organizationUserRepository,
-              applicationCacheService, mailService, logger, currentContext, globalSettings, policyRepository,
-              userRepository, policyService, tokenDataFactory)
+              applicationCacheService, mailService, logger, currentContext, globalSettings, userRepository, policyService,
+              tokenDataFactory, featureService, ssoConfigRepository)
     {
         _userManager = userManager;
         _userService = userService;
@@ -166,6 +168,11 @@ public class ResourceOwnerPasswordValidator : BaseRequestValidator<ResourceOwner
         context.Result = new GrantValidationResult(TokenRequestErrors.InvalidGrant, customResponse: customResponse);
     }
 
+    protected override ClaimsPrincipal GetSubject(ResourceOwnerPasswordValidationContext context)
+    {
+        return context.Result.Subject;
+    }
+
     private bool AuthEmailHeaderIsValid(ResourceOwnerPasswordValidationContext context)
     {
         if (!_currentContext.HttpContext.Request.Headers.ContainsKey("Auth-Email"))
diff --git a/test/Identity.IntegrationTest/Endpoints/IdentityServerSsoTests.cs b/test/Identity.IntegrationTest/Endpoints/IdentityServerSsoTests.cs
index 54cfa5d77..da23a2acd 100644
--- a/test/Identity.IntegrationTest/Endpoints/IdentityServerSsoTests.cs
+++ b/test/Identity.IntegrationTest/Endpoints/IdentityServerSsoTests.cs
@@ -245,6 +245,59 @@ public class IdentityServerSsoTests
         AssertHelper.AssertJsonProperty(trustedDeviceOption, "HasAdminApproval", JsonValueKind.False);
     }
 
+    [Fact]
+    public async Task SsoLogin_TrustedDeviceEncryption_FlagTurnedOff_DoesNotReturnOption()
+    {
+        // Arrange
+        var challenge = new string('c', 50);
+
+        // This creates SsoConfig that HAS enabled trusted device encryption which should have only been
+        // done with the feature flag turned on but we are testing that even if they have done that, this will turn off
+        // if returning as an option if the flag has later been turned off.  We should be very careful turning the flag
+        // back off.
+        var factory = await CreateFactoryAsync(new SsoConfigurationData
+        {
+            MemberDecryptionType = MemberDecryptionType.TrustedDeviceEncryption,
+        }, challenge, trustedDeviceEnabled: false);
+
+        await UpdateUserAsync(factory, user => user.MasterPassword = null);
+
+        // Act
+        var context = await factory.Server.PostAsync("/connect/token", new FormUrlEncodedContent(new Dictionary<string, string>
+        {
+            { "scope", "api offline_access" },
+            { "client_id", "web" },
+            { "deviceType", "10" },
+            { "deviceIdentifier", "test_id" },
+            { "deviceName", "firefox" },
+            { "twoFactorToken", "TEST"},
+            { "twoFactorProvider", "5" }, // RememberMe Provider
+            { "twoFactorRemember", "0" },
+            { "grant_type", "authorization_code" },
+            { "code", "test_code" },
+            { "code_verifier", challenge },
+            { "redirect_uri", "https://localhost:8080/sso-connector.html" }
+        }));
+
+        // Assert
+        // If the organization has selected TrustedDeviceEncryption but the user still has their master password
+        // they can decrypt with either option
+        Assert.Equal(StatusCodes.Status200OK, context.Response.StatusCode);
+        using var responseBody = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
+        var root = responseBody.RootElement;
+        AssertHelper.AssertJsonProperty(root, "access_token", JsonValueKind.String);
+        var userDecryptionOptions = AssertHelper.AssertJsonProperty(root, "UserDecryptionOptions", JsonValueKind.Object);
+
+        // Expected to look like:
+        // "UserDecryptionOptions": {
+        //   "Object": "userDecryptionOptions"
+        //   "HasMasterPassword": false
+        // }
+
+        // Should only have 2 properties
+        Assert.Equal(2, userDecryptionOptions.EnumerateObject().Count());
+    }
+
     [Fact]
     public async Task SsoLogin_KeyConnector_ReturnsOptions()
     {
@@ -301,7 +354,7 @@ public class IdentityServerSsoTests
         Assert.Equal("https://key_connector.com", keyConnectorUrl);
     }
 
-    private static async Task<IdentityApplicationFactory> CreateFactoryAsync(SsoConfigurationData ssoConfigurationData, string challenge)
+    private static async Task<IdentityApplicationFactory> CreateFactoryAsync(SsoConfigurationData ssoConfigurationData, string challenge, bool trustedDeviceEnabled = true)
     {
         var factory = new IdentityApplicationFactory();
 
@@ -327,7 +380,7 @@ public class IdentityServerSsoTests
         factory.SubstitueService<IFeatureService>(service =>
         {
             service.IsEnabled(FeatureFlagKeys.TrustedDeviceEncryption, Arg.Any<ICurrentContext>())
-                .Returns(true);
+                .Returns(trustedDeviceEnabled);
         });
 
         // This starts the server and finalizes services
diff --git a/test/Identity.IntegrationTest/Endpoints/IdentityServerTests.cs b/test/Identity.IntegrationTest/Endpoints/IdentityServerTests.cs
index 316f89647..f29a09c6e 100644
--- a/test/Identity.IntegrationTest/Endpoints/IdentityServerTests.cs
+++ b/test/Identity.IntegrationTest/Endpoints/IdentityServerTests.cs
@@ -65,6 +65,7 @@ public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
         Assert.Equal(0, kdf);
         var kdfIterations = AssertHelper.AssertJsonProperty(root, "KdfIterations", JsonValueKind.Number).GetInt32();
         Assert.Equal(5000, kdfIterations);
+        AssertUserDecryptionOptions(root);
     }
 
     [Theory, BitAutoData]
@@ -635,6 +636,16 @@ public class IdentityServerTests : IClassFixture<IdentityApplicationFactory>
         Assert.StartsWith("sso authentication", errorDescription.ToLowerInvariant());
     }
 
+    private static void AssertUserDecryptionOptions(JsonElement tokenResponse)
+    {
+        var userDecryptionOptions = AssertHelper.AssertJsonProperty(tokenResponse, "UserDecryptionOptions", JsonValueKind.Object)
+            .EnumerateObject();
+
+        Assert.Collection(userDecryptionOptions,
+            (prop) => { Assert.Equal("HasMasterPassword", prop.Name); Assert.Equal(JsonValueKind.True, prop.Value.ValueKind); },
+            (prop) => { Assert.Equal("Object", prop.Name); Assert.Equal("userDecryptionOptions", prop.Value.GetString()); });
+    }
+
     private void ReinitializeDbForTests()
     {
         var databaseContext = _factory.GetDatabaseContext();