From a19ae0159fe0f898ab644a1f6c501a3239871518 Mon Sep 17 00:00:00 2001
From: Ike <137194738+ike-kottlowski@users.noreply.github.com>
Date: Fri, 9 Feb 2024 12:08:22 -0800
Subject: [PATCH] [PM-5424] fix TDE provider user (#3771)

* Add Test Asserting Problem

* Fix Test

---------

Co-authored-by: Justin Baur <19896123+justindbaur@users.noreply.github.com>
---
 src/Core/Context/CurrentContext.cs            |  8 +++
 .../Endpoints/IdentityServerSsoTests.cs       | 71 +++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/src/Core/Context/CurrentContext.cs b/src/Core/Context/CurrentContext.cs
index cc74e60a8c..90ad275d02 100644
--- a/src/Core/Context/CurrentContext.cs
+++ b/src/Core/Context/CurrentContext.cs
@@ -489,6 +489,10 @@ public class CurrentContext : ICurrentContext
     {
         if (Organizations == null)
         {
+            // If we haven't had our user id set, take the one passed in since we are about to get information
+            // for them anyways.
+            UserId ??= userId;
+
             var userOrgs = await organizationUserRepository.GetManyDetailsByUserAsync(userId);
             Organizations = userOrgs.Where(ou => ou.Status == OrganizationUserStatusType.Confirmed)
                 .Select(ou => new CurrentContextOrganization(ou)).ToList();
@@ -501,6 +505,10 @@ public class CurrentContext : ICurrentContext
     {
         if (Providers == null)
         {
+            // If we haven't had our user id set, take the one passed in since we are about to get information
+            // for them anyways.
+            UserId ??= userId;
+
             var userProviders = await providerUserRepository.GetManyByUserAsync(userId);
             Providers = userProviders.Where(ou => ou.Status == ProviderUserStatusType.Confirmed)
                 .Select(ou => new CurrentContextProvider(ou)).ToList();
diff --git a/test/Identity.IntegrationTest/Endpoints/IdentityServerSsoTests.cs b/test/Identity.IntegrationTest/Endpoints/IdentityServerSsoTests.cs
index 70f6e5e1af..c775fce5eb 100644
--- a/test/Identity.IntegrationTest/Endpoints/IdentityServerSsoTests.cs
+++ b/test/Identity.IntegrationTest/Endpoints/IdentityServerSsoTests.cs
@@ -1,6 +1,9 @@
 using System.Security.Claims;
 using System.Text.Json;
 using Bit.Core.AdminConsole.Entities;
+using Bit.Core.AdminConsole.Entities.Provider;
+using Bit.Core.AdminConsole.Enums.Provider;
+using Bit.Core.AdminConsole.Repositories;
 using Bit.Core.Auth.Entities;
 using Bit.Core.Auth.Enums;
 using Bit.Core.Auth.Models.Api.Request.Accounts;
@@ -380,6 +383,74 @@ public class IdentityServerSsoTests
 
     }
 
+    [Fact]
+    public async Task SsoLogin_TrustedDeviceEncryption_ProviderUserHasManageResetPassword_ReturnsCorrectOptions()
+    {
+        var challenge = new string('c', 50);
+
+        var factory = await CreateFactoryAsync(new SsoConfigurationData
+        {
+            MemberDecryptionType = MemberDecryptionType.TrustedDeviceEncryption,
+        }, challenge);
+
+        var user = await factory.Services.GetRequiredService<IUserRepository>().GetByEmailAsync(TestEmail);
+        var providerRepository = factory.Services.GetRequiredService<IProviderRepository>();
+        var provider = await providerRepository.CreateAsync(new Provider
+        {
+            Name = "Test Provider",
+        });
+
+        var providerUserRepository = factory.Services.GetRequiredService<IProviderUserRepository>();
+        await providerUserRepository.CreateAsync(new ProviderUser
+        {
+            ProviderId = provider.Id,
+            UserId = user.Id,
+            Status = ProviderUserStatusType.Confirmed,
+            Permissions = CoreHelpers.ClassToJsonData(new Permissions
+            {
+                ManageResetPassword = true,
+            }),
+        });
+
+        var organizationUserRepository = factory.Services.GetRequiredService<IOrganizationUserRepository>();
+        var organizationUser = (await organizationUserRepository.GetManyByUserAsync(user.Id)).Single();
+
+        var providerOrganizationRepository = factory.Services.GetRequiredService<IProviderOrganizationRepository>();
+        await providerOrganizationRepository.CreateAsync(new ProviderOrganization
+        {
+            ProviderId = provider.Id,
+            OrganizationId = organizationUser.OrganizationId,
+        });
+
+        // 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.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);
+
+        var trustedDeviceOption = AssertHelper.AssertJsonProperty(userDecryptionOptions, "TrustedDeviceOption", JsonValueKind.Object);
+        AssertHelper.AssertJsonProperty(trustedDeviceOption, "HasAdminApproval", JsonValueKind.False);
+        AssertHelper.AssertJsonProperty(trustedDeviceOption, "HasManageResetPasswordPermission", JsonValueKind.True);
+    }
+
     [Fact]
     public async Task SsoLogin_KeyConnector_ReturnsOptions()
     {