From 10f590b4e75404388b4a01f317a6446c55e8c669 Mon Sep 17 00:00:00 2001
From: Thomas Rittson <31796059+eliykat@users.noreply.github.com>
Date: Thu, 25 Jan 2024 16:57:57 +1000
Subject: [PATCH] [AC-2026] Add flexible collections opt-in endpoint (#3643)

Stored procedure to be added in AC-1682
---
 .../Controllers/OrganizationsController.cs    | 33 ++++++++++++++
 .../Repositories/IOrganizationRepository.cs   |  1 +
 src/Core/Constants.cs                         |  9 +++-
 .../Repositories/OrganizationRepository.cs    | 12 +++++
 .../Repositories/OrganizationRepository.cs    |  5 +++
 .../OrganizationsControllerTests.cs           | 45 ++++++++++++++++++-
 6 files changed, 103 insertions(+), 2 deletions(-)

diff --git a/src/Api/AdminConsole/Controllers/OrganizationsController.cs b/src/Api/AdminConsole/Controllers/OrganizationsController.cs
index da0b8d4e2c..6acdace5e3 100644
--- a/src/Api/AdminConsole/Controllers/OrganizationsController.cs
+++ b/src/Api/AdminConsole/Controllers/OrganizationsController.cs
@@ -815,6 +815,39 @@ public class OrganizationsController : Controller
         return new OrganizationResponseModel(organization);
     }
 
+    /// <summary>
+    /// Migrates user, collection, and group data to the new Flexible Collections permissions scheme,
+    /// then sets organization.FlexibleCollections to true to enable these new features for the organization.
+    /// This is irreversible.
+    /// </summary>
+    /// <param name="organizationId"></param>
+    /// <exception cref="NotFoundException"></exception>
+    [HttpPost("{id}/enable-collection-enhancements")]
+    [RequireFeature(FeatureFlagKeys.FlexibleCollectionsMigration)]
+    public async Task EnableCollectionEnhancements(Guid id)
+    {
+        if (!await _currentContext.OrganizationOwner(id))
+        {
+            throw new NotFoundException();
+        }
+
+        var organization = await _organizationRepository.GetByIdAsync(id);
+        if (organization == null)
+        {
+            throw new NotFoundException();
+        }
+
+        if (organization.FlexibleCollections)
+        {
+            throw new BadRequestException("Organization has already been migrated to the new collection enhancements");
+        }
+
+        await _organizationRepository.EnableCollectionEnhancements(id);
+
+        organization.FlexibleCollections = true;
+        await _organizationService.ReplaceAndUpdateCacheAsync(organization);
+    }
+
     private async Task TryGrantOwnerAccessToSecretsManagerAsync(Guid organizationId, Guid userId)
     {
         var organizationUser = await _organizationUserRepository.GetByOrganizationAsync(organizationId, userId);
diff --git a/src/Core/AdminConsole/Repositories/IOrganizationRepository.cs b/src/Core/AdminConsole/Repositories/IOrganizationRepository.cs
index 4598a11fb9..36a445442b 100644
--- a/src/Core/AdminConsole/Repositories/IOrganizationRepository.cs
+++ b/src/Core/AdminConsole/Repositories/IOrganizationRepository.cs
@@ -15,4 +15,5 @@ public interface IOrganizationRepository : IRepository<Organization, Guid>
     Task<SelfHostedOrganizationDetails> GetSelfHostedOrganizationDetailsById(Guid id);
     Task<ICollection<Organization>> SearchUnassignedToProviderAsync(string name, string ownerEmail, int skip, int take);
     Task<IEnumerable<string>> GetOwnerEmailAddressesById(Guid organizationId);
+    Task EnableCollectionEnhancements(Guid organizationId);
 }
diff --git a/src/Core/Constants.cs b/src/Core/Constants.cs
index 07fe176381..d3e9b3f749 100644
--- a/src/Core/Constants.cs
+++ b/src/Core/Constants.cs
@@ -106,8 +106,15 @@ public static class FeatureFlagKeys
     public const string ItemShare = "item-share";
     public const string KeyRotationImprovements = "key-rotation-improvements";
     public const string DuoRedirect = "duo-redirect";
-    public const string FlexibleCollectionsMigration = "flexible-collections-migration";
+    /// <summary>
+    /// Enables flexible collections improvements for new organizations on creation
+    /// </summary>
     public const string FlexibleCollectionsSignup = "flexible-collections-signup";
+    /// <summary>
+    /// Exposes a migration button in the web vault which allows users to migrate an existing organization to
+    /// flexible collections
+    /// </summary>
+    public const string FlexibleCollectionsMigration = "flexible-collections-migration";
 
     public static List<string> GetAllKeys()
     {
diff --git a/src/Infrastructure.Dapper/AdminConsole/Repositories/OrganizationRepository.cs b/src/Infrastructure.Dapper/AdminConsole/Repositories/OrganizationRepository.cs
index f4c771adec..9080e17c3e 100644
--- a/src/Infrastructure.Dapper/AdminConsole/Repositories/OrganizationRepository.cs
+++ b/src/Infrastructure.Dapper/AdminConsole/Repositories/OrganizationRepository.cs
@@ -169,4 +169,16 @@ public class OrganizationRepository : Repository<Organization, Guid>, IOrganizat
             new { OrganizationId = organizationId },
             commandType: CommandType.StoredProcedure);
     }
+
+    public async Task EnableCollectionEnhancements(Guid organizationId)
+    {
+        using (var connection = new SqlConnection(ConnectionString))
+        {
+            await connection.ExecuteAsync(
+                "[dbo].[Organization_EnableCollectionEnhancements]",
+                new { OrganizationId = organizationId },
+                commandType: CommandType.StoredProcedure,
+                commandTimeout: 180);
+        }
+    }
 }
diff --git a/src/Infrastructure.EntityFramework/AdminConsole/Repositories/OrganizationRepository.cs b/src/Infrastructure.EntityFramework/AdminConsole/Repositories/OrganizationRepository.cs
index acc36c9449..7610f8dd15 100644
--- a/src/Infrastructure.EntityFramework/AdminConsole/Repositories/OrganizationRepository.cs
+++ b/src/Infrastructure.EntityFramework/AdminConsole/Repositories/OrganizationRepository.cs
@@ -267,4 +267,9 @@ public class OrganizationRepository : Repository<Core.AdminConsole.Entities.Orga
 
         return await query.ToListAsync();
     }
+
+    public Task EnableCollectionEnhancements(Guid organizationId)
+    {
+        throw new NotImplementedException("Collection enhancements migration is not yet supported for Entity Framework.");
+    }
 }
diff --git a/test/Api.Test/AdminConsole/Controllers/OrganizationsControllerTests.cs b/test/Api.Test/AdminConsole/Controllers/OrganizationsControllerTests.cs
index 0e4ae48877..947079eb57 100644
--- a/test/Api.Test/AdminConsole/Controllers/OrganizationsControllerTests.cs
+++ b/test/Api.Test/AdminConsole/Controllers/OrganizationsControllerTests.cs
@@ -21,10 +21,10 @@ using Bit.Core.OrganizationFeatures.OrganizationLicenses.Interfaces;
 using Bit.Core.OrganizationFeatures.OrganizationSubscriptions.Interface;
 using Bit.Core.Repositories;
 using Bit.Core.Services;
-using Bit.Core.Settings;
 using NSubstitute;
 using NSubstitute.ReturnsExtensions;
 using Xunit;
+using GlobalSettings = Bit.Core.Settings.GlobalSettings;
 
 namespace Bit.Api.Test.AdminConsole.Controllers;
 
@@ -353,4 +353,47 @@ public class OrganizationsControllerTests : IDisposable
             .SignUpAsync(organization, model.AdditionalSmSeats, model.AdditionalServiceAccounts);
         await _organizationUserRepository.DidNotReceiveWithAnyArgs().ReplaceAsync(Arg.Any<OrganizationUser>());
     }
+
+    [Theory, AutoData]
+    public async Task EnableCollectionEnhancements_Success(Organization organization)
+    {
+        organization.FlexibleCollections = false;
+        _currentContext.OrganizationOwner(organization.Id).Returns(true);
+        _organizationRepository.GetByIdAsync(organization.Id).Returns(organization);
+
+        await _sut.EnableCollectionEnhancements(organization.Id);
+
+        await _organizationRepository.Received(1).EnableCollectionEnhancements(organization.Id);
+        await _organizationService.Received(1).ReplaceAndUpdateCacheAsync(
+            Arg.Is<Organization>(o =>
+                o.Id == organization.Id &&
+                o.FlexibleCollections));
+    }
+
+    [Theory, AutoData]
+    public async Task EnableCollectionEnhancements_WhenNotOwner_Throws(Organization organization)
+    {
+        organization.FlexibleCollections = false;
+        _currentContext.OrganizationOwner(organization.Id).Returns(false);
+        _organizationRepository.GetByIdAsync(organization.Id).Returns(organization);
+
+        await Assert.ThrowsAsync<NotFoundException>(async () => await _sut.EnableCollectionEnhancements(organization.Id));
+
+        await _organizationRepository.DidNotReceiveWithAnyArgs().EnableCollectionEnhancements(Arg.Any<Guid>());
+        await _organizationService.DidNotReceiveWithAnyArgs().ReplaceAndUpdateCacheAsync(Arg.Any<Organization>());
+    }
+
+    [Theory, AutoData]
+    public async Task EnableCollectionEnhancements_WhenAlreadyMigrated_Throws(Organization organization)
+    {
+        organization.FlexibleCollections = true;
+        _currentContext.OrganizationOwner(organization.Id).Returns(true);
+        _organizationRepository.GetByIdAsync(organization.Id).Returns(organization);
+
+        var exception = await Assert.ThrowsAsync<BadRequestException>(async () => await _sut.EnableCollectionEnhancements(organization.Id));
+        Assert.Contains("has already been migrated", exception.Message);
+
+        await _organizationRepository.DidNotReceiveWithAnyArgs().EnableCollectionEnhancements(Arg.Any<Guid>());
+        await _organizationService.DidNotReceiveWithAnyArgs().ReplaceAndUpdateCacheAsync(Arg.Any<Organization>());
+    }
 }