rollback share if errors

src/Core/Repositories/IOrganizationRepository.cs

@@ -8,5 +8,6 @@ namespace Bit.Core.Repositories
     public interface IOrganizationRepository : IRepository<Organization, Guid>
     {
         Task<ICollection<Organization>> GetManyByUserIdAsync(Guid userId);
+        Task UpdateStorageAsync(Guid id, long storageIncrease);
     }
 }

src/Core/Repositories/IUserRepository.cs

@@ -9,5 +9,6 @@ namespace Bit.Core.Repositories
         Task<User> GetByEmailAsync(string email);
         Task<string> GetPublicKeyAsync(Guid id);
         Task<DateTime> GetAccountRevisionDateAsync(Guid id);
+        Task UpdateStorageAsync(Guid id, long storageIncrease);
     }
 }

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);
+            }
+        }
     }
 }

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);
+            }
+        }
     }
 }

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}");
+            var source = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/temp/{organizationId}/{attachmentId}");
-            if(!await source.ExistsAsync())
+            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 source = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/temp/{organizationId}/{attachmentId}");
-            var dest = _attachmentsContainer.GetBlockBlobReference($"{cipherId}/{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)

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(cipher.Id == default(Guid))
-            if(!org.MaxStorageGb.HasValue)
+                {
+                    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();
+                throw;
-            if(storageBytesRemaining < requestLength)
-            {
-                throw new BadRequestException("Not enough storage available for this organization.");
             }
-
-            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<Guid> 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(updatedCipher)
-                if(!hasAttachments)
+                {
+                    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);

src/Sql/dbo/Stored Procedures/Cipher_UpdateWithCollections.sql

@@ -22,13 +22,17 @@ BEGIN
     WHERE [Id] = @Id
 
     DECLARE @Size BIGINT
+    DECLARE @SizeDec BIGINT
 
-    SELECT 
+    IF @CipherAttachments IS NOT NULL
-        @Size = SUM(CAST(JSON_VALUE(value,'$.Size') AS BIGINT))
+    BEGIN
-    FROM
+        SELECT 
-        OPENJSON(@CipherAttachments)
+            @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
+    IF ISNULL(@Size, 0) > 0
-    EXEC [dbo].[User_UpdateStorage] @UserId, @SizeDec
+    BEGIN
+        EXEC [dbo].[Organization_UpdateStorage] @OrganizationId, @Size
+        EXEC [dbo].[User_UpdateStorage] @UserId, @SizeDec
+    END
+
     EXEC [dbo].[User_BumpAccountRevisionDateByOrganizationId] @OrganizationId
 END