From aba05f7970edd3a2aa00449098cff4d673dfae8f Mon Sep 17 00:00:00 2001
From: Maciej Zieniuk <mzieniuk@bitwarden.com>
Date: Thu, 27 Mar 2025 12:39:58 +0000
Subject: [PATCH] Move Key Connector controller endpoints into Key Management
 team ownership

---
 .../Auth/Controllers/AccountsController.cs    |  46 -------
 .../AccountsKeyManagementController.cs        |  56 +++++++-
 .../SetKeyConnectorKeyRequestModel.cs         |   2 +-
 .../Helpers/OrganizationTestHelpers.cs        |   5 +-
 .../AccountsKeyManagementControllerTests.cs   |  97 ++++++++++++-
 .../AccountsKeyManagementControllerTests.cs   | 129 ++++++++++++++++++
 6 files changed, 282 insertions(+), 53 deletions(-)
 rename src/Api/{Auth/Models/Request/Accounts => KeyManagement/Models/Requests}/SetKeyConnectorKeyRequestModel.cs (94%)

diff --git a/src/Api/Auth/Controllers/AccountsController.cs b/src/Api/Auth/Controllers/AccountsController.cs
index 2555a6fe2d..255993e8eb 100644
--- a/src/Api/Auth/Controllers/AccountsController.cs
+++ b/src/Api/Auth/Controllers/AccountsController.cs
@@ -284,52 +284,6 @@ public class AccountsController : Controller
         throw new BadRequestException(ModelState);
     }
 
-    [HttpPost("set-key-connector-key")]
-    public async Task PostSetKeyConnectorKeyAsync([FromBody] SetKeyConnectorKeyRequestModel model)
-    {
-        var user = await _userService.GetUserByPrincipalAsync(User);
-        if (user == null)
-        {
-            throw new UnauthorizedAccessException();
-        }
-
-        var result = await _userService.SetKeyConnectorKeyAsync(model.ToUser(user), model.Key, model.OrgIdentifier);
-        if (result.Succeeded)
-        {
-            return;
-        }
-
-        foreach (var error in result.Errors)
-        {
-            ModelState.AddModelError(string.Empty, error.Description);
-        }
-
-        throw new BadRequestException(ModelState);
-    }
-
-    [HttpPost("convert-to-key-connector")]
-    public async Task PostConvertToKeyConnector()
-    {
-        var user = await _userService.GetUserByPrincipalAsync(User);
-        if (user == null)
-        {
-            throw new UnauthorizedAccessException();
-        }
-
-        var result = await _userService.ConvertToKeyConnectorAsync(user);
-        if (result.Succeeded)
-        {
-            return;
-        }
-
-        foreach (var error in result.Errors)
-        {
-            ModelState.AddModelError(string.Empty, error.Description);
-        }
-
-        throw new BadRequestException(ModelState);
-    }
-
     [HttpPost("kdf")]
     public async Task PostKdf([FromBody] KdfRequestModel model)
     {
diff --git a/src/Api/KeyManagement/Controllers/AccountsKeyManagementController.cs b/src/Api/KeyManagement/Controllers/AccountsKeyManagementController.cs
index 85e0981f22..6ac0bb6448 100644
--- a/src/Api/KeyManagement/Controllers/AccountsKeyManagementController.cs
+++ b/src/Api/KeyManagement/Controllers/AccountsKeyManagementController.cs
@@ -23,7 +23,7 @@ using Microsoft.AspNetCore.Mvc;
 
 namespace Bit.Api.KeyManagement.Controllers;
 
-[Route("accounts/key-management")]
+[Route("accounts")]
 [Authorize("Application")]
 public class AccountsKeyManagementController : Controller
 {
@@ -73,7 +73,7 @@ public class AccountsKeyManagementController : Controller
         _webauthnKeyValidator = webAuthnKeyValidator;
     }
 
-    [HttpPost("regenerate-keys")]
+    [HttpPost("key-management/regenerate-keys")]
     public async Task RegenerateKeysAsync([FromBody] KeyRegenerationRequestModel request)
     {
         if (!_featureService.IsEnabled(FeatureFlagKeys.PrivateKeyRegeneration))
@@ -89,7 +89,7 @@ public class AccountsKeyManagementController : Controller
     }
 
 
-    [HttpPost("rotate-user-account-keys")]
+    [HttpPost("key-management/rotate-user-account-keys")]
     public async Task RotateUserAccountKeysAsync([FromBody] RotateUserAccountKeysAndDataRequestModel model)
     {
         var user = await _userService.GetUserByPrincipalAsync(User);
@@ -128,4 +128,54 @@ public class AccountsKeyManagementController : Controller
 
         throw new BadRequestException(ModelState);
     }
+
+    [HttpPost("key-management/set-key-connector-key")]
+    // Backwards compatibility, to be deleted in the future
+    [HttpPost("set-key-connector-key")]
+    public async Task PostSetKeyConnectorKeyAsync([FromBody] SetKeyConnectorKeyRequestModel model)
+    {
+        var user = await _userService.GetUserByPrincipalAsync(User);
+        if (user == null)
+        {
+            throw new UnauthorizedAccessException();
+        }
+
+        var result = await _userService.SetKeyConnectorKeyAsync(model.ToUser(user), model.Key, model.OrgIdentifier);
+        if (result.Succeeded)
+        {
+            return;
+        }
+
+        foreach (var error in result.Errors)
+        {
+            ModelState.AddModelError(string.Empty, error.Description);
+        }
+
+        throw new BadRequestException(ModelState);
+    }
+
+    [HttpPost("key-management/convert-to-key-connector")]
+    // Backwards compatibility, to be deleted in the future
+    [HttpPost("convert-to-key-connector")]
+    public async Task PostConvertToKeyConnectorAsync()
+    {
+        var user = await _userService.GetUserByPrincipalAsync(User);
+        if (user == null)
+        {
+            throw new UnauthorizedAccessException();
+        }
+
+        var result = await _userService.ConvertToKeyConnectorAsync(user);
+        if (result.Succeeded)
+        {
+            return;
+        }
+
+        foreach (var error in result.Errors)
+        {
+            ModelState.AddModelError(string.Empty, error.Description);
+        }
+
+        throw new BadRequestException(ModelState);
+    }
 }
diff --git a/src/Api/Auth/Models/Request/Accounts/SetKeyConnectorKeyRequestModel.cs b/src/Api/KeyManagement/Models/Requests/SetKeyConnectorKeyRequestModel.cs
similarity index 94%
rename from src/Api/Auth/Models/Request/Accounts/SetKeyConnectorKeyRequestModel.cs
rename to src/Api/KeyManagement/Models/Requests/SetKeyConnectorKeyRequestModel.cs
index 25d543b916..bac42bc302 100644
--- a/src/Api/Auth/Models/Request/Accounts/SetKeyConnectorKeyRequestModel.cs
+++ b/src/Api/KeyManagement/Models/Requests/SetKeyConnectorKeyRequestModel.cs
@@ -3,7 +3,7 @@ using Bit.Core.Auth.Models.Api.Request.Accounts;
 using Bit.Core.Entities;
 using Bit.Core.Enums;
 
-namespace Bit.Api.Auth.Models.Request.Accounts;
+namespace Bit.Api.KeyManagement.Models.Requests;
 
 public class SetKeyConnectorKeyRequestModel
 {
diff --git a/test/Api.IntegrationTest/Helpers/OrganizationTestHelpers.cs b/test/Api.IntegrationTest/Helpers/OrganizationTestHelpers.cs
index 9370948a85..f2bc9f4bac 100644
--- a/test/Api.IntegrationTest/Helpers/OrganizationTestHelpers.cs
+++ b/test/Api.IntegrationTest/Helpers/OrganizationTestHelpers.cs
@@ -59,7 +59,8 @@ public static class OrganizationTestHelpers
         string userEmail,
         OrganizationUserType type,
         bool accessSecretsManager = false,
-        Permissions? permissions = null
+        Permissions? permissions = null,
+        OrganizationUserStatusType userStatusType = OrganizationUserStatusType.Confirmed
     ) where T : class
     {
         var userRepository = factory.GetService<IUserRepository>();
@@ -74,7 +75,7 @@ public static class OrganizationTestHelpers
             UserId = user.Id,
             Key = null,
             Type = type,
-            Status = OrganizationUserStatusType.Confirmed,
+            Status = userStatusType,
             ExternalId = null,
             AccessSecretsManager = accessSecretsManager,
         };
diff --git a/test/Api.IntegrationTest/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs b/test/Api.IntegrationTest/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs
index 7c05e1d680..054507976d 100644
--- a/test/Api.IntegrationTest/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs
+++ b/test/Api.IntegrationTest/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs
@@ -1,4 +1,5 @@
-using System.Net;
+#nullable enable
+using System.Net;
 using Bit.Api.IntegrationTest.Factories;
 using Bit.Api.IntegrationTest.Helpers;
 using Bit.Api.KeyManagement.Models.Requests;
@@ -30,6 +31,7 @@ public class AccountsKeyManagementControllerTests : IClassFixture<ApiApplication
     private readonly LoginHelper _loginHelper;
     private readonly IUserRepository _userRepository;
     private readonly IPasswordHasher<User> _passwordHasher;
+    private readonly IOrganizationRepository _organizationRepository;
     private string _ownerEmail = null!;
 
     public AccountsKeyManagementControllerTests(ApiApplicationFactory factory)
@@ -43,6 +45,7 @@ public class AccountsKeyManagementControllerTests : IClassFixture<ApiApplication
         _emergencyAccessRepository = _factory.GetService<IEmergencyAccessRepository>();
         _organizationUserRepository = _factory.GetService<IOrganizationUserRepository>();
         _passwordHasher = _factory.GetService<IPasswordHasher<User>>();
+        _organizationRepository = _factory.GetService<IOrganizationRepository>();
     }
 
     public async Task InitializeAsync()
@@ -252,4 +255,96 @@ public class AccountsKeyManagementControllerTests : IClassFixture<ApiApplication
         Assert.Equal(request.AccountUnlockData.MasterPasswordUnlockData.KdfMemory, userNewState.KdfMemory);
         Assert.Equal(request.AccountUnlockData.MasterPasswordUnlockData.KdfParallelism, userNewState.KdfParallelism);
     }
+
+    [Theory]
+    [BitAutoData("/accounts/key-management/set-key-connector-key")]
+    [BitAutoData("/accounts/set-key-connector-key")]
+    public async Task PostSetKeyConnectorKeyAsync_NotLoggedIn_Unauthorized(string uri,
+        SetKeyConnectorKeyRequestModel request)
+    {
+        var response = await _client.PostAsJsonAsync(uri, request);
+
+        Assert.Equal(HttpStatusCode.Unauthorized, response.StatusCode);
+    }
+
+    [Theory]
+    [BitAutoData("/accounts/key-management/set-key-connector-key")]
+    [BitAutoData("/accounts/set-key-connector-key")]
+    public async Task PostSetKeyConnectorKeyAsync_Success(string uri, string organizationSsoIdentifier,
+        SetKeyConnectorKeyRequestModel request)
+    {
+        var (organization, _) = await OrganizationTestHelpers.SignUpAsync(_factory,
+            PlanType.EnterpriseAnnually, _ownerEmail, passwordManagerSeats: 10,
+            paymentMethod: PaymentMethodType.Card);
+        organization.UseKeyConnector = true;
+        organization.UseSso = true;
+        organization.Identifier = organizationSsoIdentifier;
+        await _organizationRepository.ReplaceAsync(organization);
+
+        var ssoUserEmail = $"integration-test{Guid.NewGuid()}@bitwarden.com";
+        await _factory.LoginWithNewAccount(ssoUserEmail);
+        await _loginHelper.LoginAsync(ssoUserEmail);
+
+        await OrganizationTestHelpers.CreateUserAsync(_factory, organization.Id, ssoUserEmail,
+            OrganizationUserType.User, userStatusType: OrganizationUserStatusType.Invited);
+
+        request.Key = _mockEncryptedString;
+        request.OrgIdentifier = organizationSsoIdentifier;
+
+        var response = await _client.PostAsJsonAsync(uri, request);
+        response.EnsureSuccessStatusCode();
+
+        var user = await _userRepository.GetByEmailAsync(ssoUserEmail);
+        Assert.NotNull(user);
+        Assert.Equal(request.Key, user.Key);
+        Assert.True(user.UsesKeyConnector);
+        Assert.Equal(DateTime.UtcNow, user.RevisionDate, TimeSpan.FromMinutes(1));
+        Assert.Equal(DateTime.UtcNow, user.AccountRevisionDate, TimeSpan.FromMinutes(1));
+        var ssoOrganizationUser =
+            await _organizationUserRepository.GetByOrganizationAsync(organization.Id, user.Id);
+        Assert.NotNull(ssoOrganizationUser);
+        Assert.Equal(OrganizationUserStatusType.Accepted, ssoOrganizationUser.Status);
+        Assert.Equal(user.Id, ssoOrganizationUser.UserId);
+        Assert.Null(ssoOrganizationUser.Email);
+    }
+
+    [Theory]
+    [BitAutoData("/accounts/key-management/convert-to-key-connector")]
+    [BitAutoData("/accounts/convert-to-key-connector")]
+    public async Task PostConvertToKeyConnectorAsync_NotLoggedIn_Unauthorized(string uri)
+    {
+        var response = await _client.PostAsJsonAsync(uri, new { });
+
+        Assert.Equal(HttpStatusCode.Unauthorized, response.StatusCode);
+    }
+
+    [Theory]
+    [BitAutoData("/accounts/key-management/convert-to-key-connector")]
+    [BitAutoData("/accounts/convert-to-key-connector")]
+    public async Task PostConvertToKeyConnectorAsync_Success(string uri)
+    {
+        var (organization, _) = await OrganizationTestHelpers.SignUpAsync(_factory,
+            PlanType.EnterpriseAnnually, _ownerEmail, passwordManagerSeats: 10,
+            paymentMethod: PaymentMethodType.Card);
+        organization.UseKeyConnector = true;
+        organization.UseSso = true;
+        await _organizationRepository.ReplaceAsync(organization);
+
+        var ssoUserEmail = $"integration-test{Guid.NewGuid()}@bitwarden.com";
+        await _factory.LoginWithNewAccount(ssoUserEmail);
+        await _loginHelper.LoginAsync(ssoUserEmail);
+
+        await OrganizationTestHelpers.CreateUserAsync(_factory, organization.Id, ssoUserEmail,
+            OrganizationUserType.User, userStatusType: OrganizationUserStatusType.Accepted);
+
+        var response = await _client.PostAsJsonAsync(uri, new { });
+        response.EnsureSuccessStatusCode();
+
+        var user = await _userRepository.GetByEmailAsync(ssoUserEmail);
+        Assert.NotNull(user);
+        Assert.Null(user.MasterPassword);
+        Assert.True(user.UsesKeyConnector);
+        Assert.Equal(DateTime.UtcNow, user.RevisionDate, TimeSpan.FromMinutes(1));
+        Assert.Equal(DateTime.UtcNow, user.AccountRevisionDate, TimeSpan.FromMinutes(1));
+    }
 }
diff --git a/test/Api.Test/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs b/test/Api.Test/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs
index 49c4f88cb4..d2775762e8 100644
--- a/test/Api.Test/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs
+++ b/test/Api.Test/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs
@@ -178,4 +178,133 @@ public class AccountsKeyManagementControllerTests
             Assert.NotEmpty(ex.ModelState.Values);
         }
     }
+
+    [Theory]
+    [BitAutoData]
+    public async Task PostSetKeyConnectorKeyAsync_UserNull_Throws(
+        SutProvider<AccountsKeyManagementController> sutProvider,
+        SetKeyConnectorKeyRequestModel data)
+    {
+        sutProvider.GetDependency<IUserService>().GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>()).ReturnsNull();
+
+        await Assert.ThrowsAsync<UnauthorizedAccessException>(() => sutProvider.Sut.PostSetKeyConnectorKeyAsync(data));
+
+        await sutProvider.GetDependency<IUserService>().ReceivedWithAnyArgs(0)
+            .SetKeyConnectorKeyAsync(Arg.Any<User>(), Arg.Any<string>(), Arg.Any<string>());
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task PostSetKeyConnectorKeyAsync_SetKeyConnectorKeyFails_ThrowsBadRequestWithErrorResponse(
+        SutProvider<AccountsKeyManagementController> sutProvider,
+        SetKeyConnectorKeyRequestModel data, User expectedUser)
+    {
+        expectedUser.PublicKey = null;
+        expectedUser.PrivateKey = null;
+        sutProvider.GetDependency<IUserService>().GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>())
+            .Returns(expectedUser);
+        sutProvider.GetDependency<IUserService>()
+            .SetKeyConnectorKeyAsync(Arg.Any<User>(), Arg.Any<string>(), Arg.Any<string>())
+            .Returns(IdentityResult.Failed(new IdentityError { Description = "set key connector key error" }));
+
+        var badRequestException =
+            await Assert.ThrowsAsync<BadRequestException>(() => sutProvider.Sut.PostSetKeyConnectorKeyAsync(data));
+
+        Assert.Equal(1, badRequestException.ModelState.ErrorCount);
+        Assert.Equal("set key connector key error", badRequestException.ModelState.Root.Errors[0].ErrorMessage);
+        await sutProvider.GetDependency<IUserService>().Received(1)
+            .SetKeyConnectorKeyAsync(Arg.Do<User>(user =>
+            {
+                Assert.Equal(expectedUser.Id, user.Id);
+                Assert.Equal(data.Key, user.Key);
+                Assert.Equal(data.Kdf, user.Kdf);
+                Assert.Equal(data.KdfIterations, user.KdfIterations);
+                Assert.Equal(data.KdfMemory, user.KdfMemory);
+                Assert.Equal(data.KdfParallelism, user.KdfParallelism);
+                Assert.Equal(data.Keys.PublicKey, user.PublicKey);
+                Assert.Equal(data.Keys.EncryptedPrivateKey, user.PrivateKey);
+            }), Arg.Is(data.Key), Arg.Is(data.OrgIdentifier));
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task PostSetKeyConnectorKeyAsync_SetKeyConnectorKeySucceeds_OkResponse(
+        SutProvider<AccountsKeyManagementController> sutProvider,
+        SetKeyConnectorKeyRequestModel data, User expectedUser)
+    {
+        expectedUser.PublicKey = null;
+        expectedUser.PrivateKey = null;
+        sutProvider.GetDependency<IUserService>().GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>())
+            .Returns(expectedUser);
+        sutProvider.GetDependency<IUserService>()
+            .SetKeyConnectorKeyAsync(Arg.Any<User>(), Arg.Any<string>(), Arg.Any<string>())
+            .Returns(IdentityResult.Success);
+
+        await sutProvider.Sut.PostSetKeyConnectorKeyAsync(data);
+
+        await sutProvider.GetDependency<IUserService>().Received(1)
+            .SetKeyConnectorKeyAsync(Arg.Do<User>(user =>
+            {
+                Assert.Equal(expectedUser.Id, user.Id);
+                Assert.Equal(data.Key, user.Key);
+                Assert.Equal(data.Kdf, user.Kdf);
+                Assert.Equal(data.KdfIterations, user.KdfIterations);
+                Assert.Equal(data.KdfMemory, user.KdfMemory);
+                Assert.Equal(data.KdfParallelism, user.KdfParallelism);
+                Assert.Equal(data.Keys.PublicKey, user.PublicKey);
+                Assert.Equal(data.Keys.EncryptedPrivateKey, user.PrivateKey);
+            }), Arg.Is(data.Key), Arg.Is(data.OrgIdentifier));
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task PostConvertToKeyConnectorAsync_UserNull_Throws(
+        SutProvider<AccountsKeyManagementController> sutProvider)
+    {
+        sutProvider.GetDependency<IUserService>().GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>()).ReturnsNull();
+
+        await Assert.ThrowsAsync<UnauthorizedAccessException>(() => sutProvider.Sut.PostConvertToKeyConnectorAsync());
+
+        await sutProvider.GetDependency<IUserService>().ReceivedWithAnyArgs(0)
+            .ConvertToKeyConnectorAsync(Arg.Any<User>());
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task PostConvertToKeyConnectorAsync_ConvertToKeyConnectorFails_ThrowsBadRequestWithErrorResponse(
+        SutProvider<AccountsKeyManagementController> sutProvider,
+        User expectedUser)
+    {
+        sutProvider.GetDependency<IUserService>().GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>())
+            .Returns(expectedUser);
+        sutProvider.GetDependency<IUserService>()
+            .ConvertToKeyConnectorAsync(Arg.Any<User>())
+            .Returns(IdentityResult.Failed(new IdentityError { Description = "convert to key connector error" }));
+
+        var badRequestException =
+            await Assert.ThrowsAsync<BadRequestException>(() => sutProvider.Sut.PostConvertToKeyConnectorAsync());
+
+        Assert.Equal(1, badRequestException.ModelState.ErrorCount);
+        Assert.Equal("convert to key connector error", badRequestException.ModelState.Root.Errors[0].ErrorMessage);
+        await sutProvider.GetDependency<IUserService>().Received(1)
+            .ConvertToKeyConnectorAsync(Arg.Is(expectedUser));
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task PostConvertToKeyConnectorAsync_ConvertToKeyConnectorSucceeds_OkResponse(
+        SutProvider<AccountsKeyManagementController> sutProvider,
+        User expectedUser)
+    {
+        sutProvider.GetDependency<IUserService>().GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>())
+            .Returns(expectedUser);
+        sutProvider.GetDependency<IUserService>()
+            .ConvertToKeyConnectorAsync(Arg.Any<User>())
+            .Returns(IdentityResult.Success);
+
+        await sutProvider.Sut.PostConvertToKeyConnectorAsync();
+
+        await sutProvider.GetDependency<IUserService>().Received(1)
+            .ConvertToKeyConnectorAsync(Arg.Is(expectedUser));
+    }
 }