diff --git a/src/Core/Repositories/IOrganizationRepository.cs b/src/Core/Repositories/IOrganizationRepository.cs index 2c2788c7a5..826084c9a8 100644 --- a/src/Core/Repositories/IOrganizationRepository.cs +++ b/src/Core/Repositories/IOrganizationRepository.cs @@ -8,5 +8,6 @@ namespace Bit.Core.Repositories public interface IOrganizationRepository : IRepository { Task> GetManyByUserIdAsync(Guid userId); + Task UpdateStorageAsync(Guid id, long storageIncrease); } } diff --git a/src/Core/Repositories/IUserRepository.cs b/src/Core/Repositories/IUserRepository.cs index cd96a165d8..57a945cf66 100644 --- a/src/Core/Repositories/IUserRepository.cs +++ b/src/Core/Repositories/IUserRepository.cs @@ -9,5 +9,6 @@ namespace Bit.Core.Repositories Task GetByEmailAsync(string email); Task GetPublicKeyAsync(Guid id); Task GetAccountRevisionDateAsync(Guid id); + Task UpdateStorageAsync(Guid id, long storageIncrease); } } diff --git a/src/Core/Repositories/SqlServer/OrganizationRepository.cs b/src/Core/Repositories/SqlServer/OrganizationRepository.cs index 4274d847c4..a5b1510024 100644 --- a/src/Core/Repositories/SqlServer/OrganizationRepository.cs +++ b/src/Core/Repositories/SqlServer/OrganizationRepository.cs @@ -31,5 +31,17 @@ namespace Bit.Core.Repositories.SqlServer return results.ToList(); } } + + public async Task UpdateStorageAsync(Guid id, long storageIncrease) + { + using(var connection = new SqlConnection(ConnectionString)) + { + await connection.ExecuteAsync( + "[dbo].[Organization_UpdateStorage]", + new { Id = id, StorageIncrease = storageIncrease }, + commandType: CommandType.StoredProcedure, + commandTimeout: 180); + } + } } } diff --git a/src/Core/Repositories/SqlServer/UserRepository.cs b/src/Core/Repositories/SqlServer/UserRepository.cs index 5963b4b7d1..33a489812d 100644 --- a/src/Core/Repositories/SqlServer/UserRepository.cs +++ b/src/Core/Repositories/SqlServer/UserRepository.cs @@ -78,5 +78,17 @@ namespace Bit.Core.Repositories.SqlServer commandTimeout: 180); } } + + public async Task UpdateStorageAsync(Guid id, long storageIncrease) + { + using(var connection = new SqlConnection(ConnectionString)) + { + await connection.ExecuteAsync( + $"[{Schema}].[{Table}_UpdateStorage]", + new { Id = id, StorageIncrease = storageIncrease }, + commandType: CommandType.StoredProcedure, + commandTimeout: 180); + } + } } } diff --git a/src/Core/Services/Implementations/AzureAttachmentStorageService.cs b/src/Core/Services/Implementations/AzureAttachmentStorageService.cs index 0a6e37c4a9..b021f623e2 100644 --- a/src/Core/Services/Implementations/AzureAttachmentStorageService.cs +++ b/src/Core/Services/Implementations/AzureAttachmentStorageService.cs @@ -27,20 +27,20 @@ namespace Bit.Core.Services public async Task UploadShareAttachmentAsync(Stream stream, Guid cipherId, Guid organizationId, string attachmentId) { - await UploadAttachmentAsync(stream, $"{cipherId}/share/{organizationId}/{attachmentId}"); + await UploadAttachmentAsync(stream, $"{cipherId}/temp/{organizationId}/{attachmentId}"); } public async Task StartShareAttachmentAsync(Guid cipherId, Guid organizationId, string attachmentId) { await InitAsync(); - var source = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/share/{organizationId}/{attachmentId}"); - if(!await source.ExistsAsync()) + var source = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/temp/{organizationId}/{attachmentId}"); + if(!(await source.ExistsAsync())) { return; } var dest = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/{attachmentId}"); - if(!await dest.ExistsAsync()) + if(!(await dest.ExistsAsync())) { return; } @@ -56,7 +56,7 @@ namespace Bit.Core.Services public async Task CommitShareAttachmentAsync(Guid cipherId, Guid organizationId, string attachmentId) { await InitAsync(); - var source = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/share/{organizationId}/{attachmentId}"); + var source = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/temp/{organizationId}/{attachmentId}"); var original = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/temp/{attachmentId}"); await original.DeleteIfExistsAsync(); await source.DeleteIfExistsAsync(); @@ -65,18 +65,19 @@ namespace Bit.Core.Services public async Task RollbackShareAttachmentAsync(Guid cipherId, Guid organizationId, string attachmentId) { await InitAsync(); - var source = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/share/{organizationId}/{attachmentId}"); - var dest = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/{attachmentId}"); + var source = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/temp/{organizationId}/{attachmentId}"); + await source.DeleteIfExistsAsync(); + var original = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/temp/{attachmentId}"); - if(!await original.ExistsAsync()) + if(!(await original.ExistsAsync())) { return; } + var dest = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/{attachmentId}"); await dest.DeleteIfExistsAsync(); await dest.StartCopyAsync(original); await original.DeleteIfExistsAsync(); - await source.DeleteIfExistsAsync(); } public async Task DeleteAttachmentAsync(Guid cipherId, string attachmentId) diff --git a/src/Core/Services/Implementations/CipherService.cs b/src/Core/Services/Implementations/CipherService.cs index 8cff364785..124602eca1 100644 --- a/src/Core/Services/Implementations/CipherService.cs +++ b/src/Core/Services/Implementations/CipherService.cs @@ -175,24 +175,46 @@ namespace Bit.Core.Services public async Task CreateAttachmentShareAsync(Cipher cipher, Stream stream, string fileName, long requestLength, string attachmentId, Guid organizationId) { - if(requestLength < 1) + try { - throw new BadRequestException("No data to attach."); - } + if(requestLength < 1) + { + throw new BadRequestException("No data to attach."); + } - var org = await _organizationRepository.GetByIdAsync(organizationId); - if(!org.MaxStorageGb.HasValue) + if(cipher.Id == default(Guid)) + { + throw new BadRequestException(nameof(cipher.Id)); + } + + if(cipher.OrganizationId.HasValue) + { + throw new BadRequestException("Cipher belongs to an organization already."); + } + + var org = await _organizationRepository.GetByIdAsync(organizationId); + if(org == null || !org.MaxStorageGb.HasValue) + { + throw new BadRequestException("This organization cannot use attachments."); + } + + var storageBytesRemaining = org.StorageBytesRemaining(); + if(storageBytesRemaining < requestLength) + { + throw new BadRequestException("Not enough storage available for this organization."); + } + + await _attachmentStorageService.UploadShareAttachmentAsync(stream, cipher.Id, organizationId, attachmentId); + } + catch { - throw new BadRequestException("This organization cannot use attachments."); - } + foreach(var attachment in cipher.GetAttachments()) + { + await _attachmentStorageService.RollbackShareAttachmentAsync(cipher.Id, organizationId, attachment.Key); + } - var storageBytesRemaining = org.StorageBytesRemaining(); - if(storageBytesRemaining < requestLength) - { - throw new BadRequestException("Not enough storage available for this organization."); + throw; } - - await _attachmentStorageService.UploadShareAttachmentAsync(stream, cipher.Id, organizationId, attachmentId); } public async Task DeleteAsync(Cipher cipher, Guid deletingUserId, bool orgAdmin = false) @@ -281,31 +303,47 @@ namespace Bit.Core.Services public async Task ShareAsync(Cipher originalCipher, Cipher cipher, Guid organizationId, IEnumerable collectionIds, Guid sharingUserId) { - if(cipher.Id == default(Guid)) - { - throw new BadRequestException(nameof(cipher.Id)); - } - - if(cipher.OrganizationId.HasValue) - { - throw new BadRequestException("Already belongs to an organization."); - } - - if(!cipher.UserId.HasValue || cipher.UserId.Value != sharingUserId) - { - throw new NotFoundException(); - } - var attachments = cipher.GetAttachments(); var hasAttachments = (attachments?.Count ?? 0) > 0; + var storageAdjustment = attachments?.Sum(a => a.Value.Size) ?? 0; + var updatedCipher = false; + var migratedAttachments = false; try { + if(cipher.Id == default(Guid)) + { + throw new BadRequestException(nameof(cipher.Id)); + } + + if(cipher.OrganizationId.HasValue) + { + throw new BadRequestException("Already belongs to an organization."); + } + + if(!cipher.UserId.HasValue || cipher.UserId.Value != sharingUserId) + { + throw new NotFoundException(); + } + + var org = await _organizationRepository.GetByIdAsync(organizationId); + if(!org.MaxStorageGb.HasValue) + { + throw new BadRequestException("This organization cannot use attachments."); + } + + var storageBytesRemaining = org.StorageBytesRemaining(); + if(storageBytesRemaining < storageAdjustment) + { + throw new BadRequestException("Not enough storage available for this organization."); + } + // Sproc will not save this UserId on the cipher. It is used limit scope of the collectionIds. cipher.UserId = sharingUserId; cipher.OrganizationId = organizationId; cipher.RevisionDate = DateTime.UtcNow; await _cipherRepository.ReplaceAsync(cipher, collectionIds); + updatedCipher = true; if(hasAttachments) { @@ -313,18 +351,29 @@ namespace Bit.Core.Services foreach(var attachment in attachments) { await _attachmentStorageService.StartShareAttachmentAsync(cipher.Id, organizationId, attachment.Key); + migratedAttachments = true; } } } catch { // roll everything back - await _cipherRepository.ReplaceAsync(originalCipher); - if(!hasAttachments) + if(updatedCipher) + { + await _cipherRepository.ReplaceAsync(originalCipher); + } + + if(!hasAttachments || !migratedAttachments) { throw; } + if(updatedCipher) + { + await _userRepository.UpdateStorageAsync(sharingUserId, storageAdjustment); + await _organizationRepository.UpdateStorageAsync(organizationId, -1 * storageAdjustment); + } + foreach(var attachment in attachments) { await _attachmentStorageService.RollbackShareAttachmentAsync(cipher.Id, organizationId, attachment.Key); diff --git a/src/Sql/dbo/Stored Procedures/Cipher_UpdateWithCollections.sql b/src/Sql/dbo/Stored Procedures/Cipher_UpdateWithCollections.sql index dd3f8c7984..d29f42a438 100644 --- a/src/Sql/dbo/Stored Procedures/Cipher_UpdateWithCollections.sql +++ b/src/Sql/dbo/Stored Procedures/Cipher_UpdateWithCollections.sql @@ -22,13 +22,17 @@ BEGIN WHERE [Id] = @Id DECLARE @Size BIGINT + DECLARE @SizeDec BIGINT - SELECT - @Size = SUM(CAST(JSON_VALUE(value,'$.Size') AS BIGINT)) - FROM - OPENJSON(@CipherAttachments) + IF @CipherAttachments IS NOT NULL + BEGIN + SELECT + @Size = SUM(CAST(JSON_VALUE(value,'$.Size') AS BIGINT)) + FROM + OPENJSON(@CipherAttachments) - DECLARE @SizeDec BIGINT = @Size * -1 + SET @SizeDec = @Size * -1 + END UPDATE [dbo].[Cipher] @@ -83,7 +87,11 @@ BEGIN WHERE [Id] IN (SELECT [Id] FROM [AvailableCollectionsCTE]) - EXEC [dbo].[Organization_UpdateStorage] @OrganizationId, @Size - EXEC [dbo].[User_UpdateStorage] @UserId, @SizeDec + IF ISNULL(@Size, 0) > 0 + BEGIN + EXEC [dbo].[Organization_UpdateStorage] @OrganizationId, @Size + EXEC [dbo].[User_UpdateStorage] @UserId, @SizeDec + END + EXEC [dbo].[User_BumpAccountRevisionDateByOrganizationId] @OrganizationId END \ No newline at end of file