From 7a52965300f3b4f74e27d3e640b681446350959b Mon Sep 17 00:00:00 2001 From: Uncled1023 Date: Sun, 29 May 2022 14:21:24 -0700 Subject: [PATCH] Fixed disposal of encrypted streams --- .../Paste/Controllers/PasteController.cs | 4 +-- .../Upload/Controllers/UploadController.cs | 6 +++++ .../Vault/Controllers/VaultController.cs | 2 +- Utilities/Cryptography/AesCounterMode.cs | 27 ++++++++++++++----- Utilities/Cryptography/AesCounterStream.cs | 6 +++-- Utilities/PooledArray.cs | 20 +++++++++++++- Utilities/ResponseHelper.cs | 3 +-- 7 files changed, 54 insertions(+), 14 deletions(-) diff --git a/Teknik/Areas/Paste/Controllers/PasteController.cs b/Teknik/Areas/Paste/Controllers/PasteController.cs index 30837f2..f9c660c 100644 --- a/Teknik/Areas/Paste/Controllers/PasteController.cs +++ b/Teknik/Areas/Paste/Controllers/PasteController.cs @@ -146,7 +146,7 @@ namespace Teknik.Areas.Paste.Controllers if (string.IsNullOrEmpty(fileName)) return new StatusCodeResult(StatusCodes.Status404NotFound); var storageService = StorageServiceFactory.GetStorageService(_config.PasteConfig.StorageConfig); - var fileStream = storageService.GetFile(fileName); + using var fileStream = storageService.GetFile(fileName); if (fileStream == null) return new StatusCodeResult(StatusCodes.Status404NotFound); @@ -317,7 +317,7 @@ namespace Teknik.Areas.Paste.Controllers if (string.IsNullOrEmpty(paste.FileName)) return new StatusCodeResult(StatusCodes.Status404NotFound); var storageService = StorageServiceFactory.GetStorageService(_config.PasteConfig.StorageConfig); - var fileStream = storageService.GetFile(paste.FileName); + using var fileStream = storageService.GetFile(paste.FileName); if (fileStream == null) return new StatusCodeResult(StatusCodes.Status404NotFound); diff --git a/Teknik/Areas/Upload/Controllers/UploadController.cs b/Teknik/Areas/Upload/Controllers/UploadController.cs index 321ead0..7cdefe3 100644 --- a/Teknik/Areas/Upload/Controllers/UploadController.cs +++ b/Teknik/Areas/Upload/Controllers/UploadController.cs @@ -310,6 +310,8 @@ namespace Teknik.Areas.Upload.Controllers long length = contentLength; if (fileStream != null) { + Response.RegisterForDispose(fileStream); + #region Range Calculation // Are they downloading it by range? bool byRange = !string.IsNullOrEmpty(Request.Headers["Range"]); // We do not support ranges @@ -428,6 +430,8 @@ namespace Teknik.Areas.Upload.Controllers var fileStream = storageService.GetFile(upload.FileName); if (fileStream != null) { + Response.RegisterForDispose(fileStream); + // Notify the client the content length we'll be outputting Response.Headers.Add("Content-Length", upload.ContentLength.ToString()); @@ -527,6 +531,8 @@ namespace Teknik.Areas.Upload.Controllers Response.RegisterForDispose(ivArray); var aesStream = new AesCounterStream(fileStream, false, keyArray, ivArray); + Response.RegisterForDispose(aesStream); + //return File(aesStream, contentType, true); return new BufferedFileStreamResult(contentType, async (response) => await ResponseHelper.StreamToOutput(response, aesStream, length, _config.UploadConfig.ChunkSize), false); } diff --git a/Teknik/Areas/Vault/Controllers/VaultController.cs b/Teknik/Areas/Vault/Controllers/VaultController.cs index 119c27d..871d468 100644 --- a/Teknik/Areas/Vault/Controllers/VaultController.cs +++ b/Teknik/Areas/Vault/Controllers/VaultController.cs @@ -110,7 +110,7 @@ namespace Teknik.Areas.Vault.Controllers { // Read in the file var storageService = StorageServiceFactory.GetStorageService(_config.PasteConfig.StorageConfig); - var fileStream = storageService.GetFile(paste.Paste.FileName); + using var fileStream = storageService.GetFile(paste.Paste.FileName); if (fileStream == null) continue; diff --git a/Utilities/Cryptography/AesCounterMode.cs b/Utilities/Cryptography/AesCounterMode.cs index c0ccf01..f5aacb7 100644 --- a/Utilities/Cryptography/AesCounterMode.cs +++ b/Utilities/Cryptography/AesCounterMode.cs @@ -26,7 +26,7 @@ namespace Teknik.Utilities.Cryptography _Algo.Padding = PaddingMode.None; // Set the internal variables - _InitialCounter = initialCounter; + _InitialCounter = new PooledArray(initialCounter); } public override ICryptoTransform CreateEncryptor(byte[] key, byte[] iv) @@ -51,7 +51,11 @@ namespace Teknik.Utilities.Cryptography protected override void Dispose(bool disposed) { - _Algo.Dispose(); + if (!disposed) + { + _Algo.Dispose(); + _InitialCounter.Dispose(); + } } } @@ -229,10 +233,21 @@ namespace Teknik.Utilities.Cryptography public void Dispose() { - _CounterEncryptor.Dispose(); - _IV.Dispose(); - _Counter.Dispose(); - _EncryptedCounter.Dispose(); + Dispose(true); + + // Suppress finalization. + GC.SuppressFinalize(this); + } + + public void Dispose(bool disposing) + { + if (disposing) + { + _CounterEncryptor.Dispose(); + _IV.Dispose(); + _Counter.Dispose(); + _EncryptedCounter.Dispose(); + } } } } diff --git a/Utilities/Cryptography/AesCounterStream.cs b/Utilities/Cryptography/AesCounterStream.cs index 7c80dc0..9edeb2f 100644 --- a/Utilities/Cryptography/AesCounterStream.cs +++ b/Utilities/Cryptography/AesCounterStream.cs @@ -251,8 +251,10 @@ namespace Teknik.Utilities.Cryptography protected override void Dispose(bool disposing) { - _Inner.Dispose(); - _Cipher.Dispose(); + if (disposing) + { + _Cipher.Dispose(); + } base.Dispose(disposing); } diff --git a/Utilities/PooledArray.cs b/Utilities/PooledArray.cs index 1871a18..8da6b5b 100644 --- a/Utilities/PooledArray.cs +++ b/Utilities/PooledArray.cs @@ -23,9 +23,16 @@ namespace Teknik.Utilities public PooledArray(byte[] array) { + Length = array.Length; Array = _arrayPool.Rent(array.Length); array.CopyTo(Array, 0); + } + + public PooledArray(PooledArray array) + { Length = array.Length; + Array = _arrayPool.Rent(array.Length); + array.CopyTo(Array); } public void CopyTo(byte[] destination) @@ -40,7 +47,18 @@ namespace Teknik.Utilities public void Dispose() { - _arrayPool.Return(Array); + Dispose(true); + + // Suppress finalization. + GC.SuppressFinalize(this); + } + + public void Dispose(bool disposing) + { + if (disposing) + { + _arrayPool.Return(Array); + } } } } diff --git a/Utilities/ResponseHelper.cs b/Utilities/ResponseHelper.cs index 1e3ecbf..28acdde 100644 --- a/Utilities/ResponseHelper.cs +++ b/Utilities/ResponseHelper.cs @@ -18,14 +18,13 @@ namespace Teknik.Utilities bufferSize = length; var pooledArray = new PooledArray(bufferSize); - response.RegisterForDispose(stream); response.RegisterForDispose(pooledArray); try { int processedBytes; do { - processedBytes = await stream.ReadAsync(pooledArray.Array, 0, pooledArray.Length); + processedBytes = stream.Read(pooledArray.Array, 0, pooledArray.Length); if (processedBytes > 0) { await response.Body.WriteAsync(pooledArray.Array, 0, processedBytes);