Fixed some bugs

StorageService/IStorageService.cs

@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.IO;
 using System.Threading.Tasks;
 using Teknik.Configuration;
+using Teknik.Utilities;
 
 namespace Teknik.StorageService
 {
@@ -12,7 +13,7 @@ namespace Teknik.StorageService
         public Stream GetFile(string fileName);
         public List<string> GetFileNames();
         public Task SaveFile(string fileName, Stream file);
-        public Task SaveEncryptedFile(string fileName, Stream file, byte[] key, byte[] iv);
+        public Task SaveEncryptedFile(string fileName, Stream file, PooledArray key, PooledArray iv);
         public void DeleteFile(string fileName);
     }
 }

StorageService/LocalStorageService.cs

@@ -38,7 +38,7 @@ namespace Teknik.StorageService
             return null;
         }
 
-        public override async Task SaveEncryptedFile(string fileName, Stream file, byte[] key, byte[] iv)
+        public override async Task SaveEncryptedFile(string fileName, Stream file, PooledArray key, PooledArray iv)
         {
             if (!Directory.Exists(_config.LocalDirectory))
                 Directory.CreateDirectory(_config.LocalDirectory);

StorageService/MemoryStorageService.cs

@@ -57,7 +57,7 @@ namespace Teknik.StorageService
             return new MemoryStream(Files[fileName]);
         }
 
-        public override async Task SaveEncryptedFile(string fileName, Stream file, byte[] key, byte[] iv)
+        public override async Task SaveEncryptedFile(string fileName, Stream file, PooledArray key, PooledArray iv)
         {
             if (file == null ||
                 Files.ContainsKey(fileName))

StorageService/StorageService.cs

@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.IO;
 using System.Threading.Tasks;
 using Teknik.Configuration;
+using Teknik.Utilities;
 
 namespace Teknik.StorageService
 {
@@ -19,7 +20,7 @@ namespace Teknik.StorageService
         public abstract Stream GetFile(string fileName);
         public abstract List<string> GetFileNames();
         public abstract Task SaveFile(string fileName, Stream file);
-        public abstract Task SaveEncryptedFile(string fileName, Stream file, byte[] key, byte[] iv);
+        public abstract Task SaveEncryptedFile(string fileName, Stream file, PooledArray key, PooledArray iv);
         public abstract void DeleteFile(string fileName);
     }
 }

Teknik/Areas/Error/Controllers/ErrorController.cs

@@ -148,7 +148,7 @@ namespace Teknik.Areas.Error.Controllers
             model.StatusCode = StatusCodes.Status500InternalServerError;
             model.Exception = exception;
 
-            return GenerateActionResult(CreateErrorObj("Http", StatusCodes.Status500InternalServerError, exception.Message), View("~/Areas/Error/Views/Error/Http500.cshtml", model));
+            return GenerateActionResult(CreateErrorObj("Http", StatusCodes.Status500InternalServerError, exception?.Message), View("~/Areas/Error/Views/Error/Http500.cshtml", model));
         }
 
         [HttpPost]

Teknik/Areas/Paste/Controllers/PasteController.cs

@@ -85,8 +85,25 @@ namespace Teknik.Areas.Paste.Controllers
                     }
                 }
 
-                byte[] ivBytes = (string.IsNullOrEmpty(iv)) ? new byte[blockSize] : Encoding.Unicode.GetBytes(iv);
+                PooledArray ivArray;
-                byte[] keyBytes = (string.IsNullOrEmpty(key)) ? new byte[keySize] : AesCounterManaged.CreateKey(key, ivBytes, keySize);
+                if (!string.IsNullOrEmpty(iv))
+                {
+                    var ivBytes = Encoding.Unicode.GetBytes(iv);
+                    ivArray = new PooledArray(ivBytes);
+                }
+                else
+                    ivArray = new PooledArray(blockSize);
+                Response.RegisterForDispose(ivArray);
+
+                PooledArray keyArray;
+                if (!string.IsNullOrEmpty(key))
+                {
+                    var keyBytes = AesCounterManaged.CreateKey(key, ivArray.Array, keySize);
+                    keyArray = new PooledArray(keyBytes);
+                }
+                else
+                    keyArray = new PooledArray(keySize);
+                Response.RegisterForDispose(keyArray);
 
                 // The paste has a password set
                 if (!string.IsNullOrEmpty(hashedPass))
@@ -100,7 +117,7 @@ namespace Teknik.Areas.Paste.Controllers
                     if (!string.IsNullOrEmpty(password))
                     {
                         hash = Crypto.HashPassword(key, password);
-                        keyBytes = AesCounterManaged.CreateKey(password, ivBytes, keySize);
+                        AesCounterManaged.CreateKey(password, ivArray.Array, keySize).CopyTo(keyArray.Array, 0);
                     }
                     if (string.IsNullOrEmpty(password) || hash != hashedPass)
                     {
@@ -136,7 +153,7 @@ namespace Teknik.Areas.Paste.Controllers
                 // Only load the model content if we aren't downloading it.
                 if (type.ToLower() != "download")
                 {
-                    using (AesCounterStream cs = new AesCounterStream(fileStream, false, keyBytes, ivBytes))
+                    using (AesCounterStream cs = new AesCounterStream(fileStream, false, keyArray, ivArray))
                     using (StreamReader sr = new StreamReader(cs, Encoding.Unicode))
                     {
                         model.Content = await sr.ReadToEndAsync();
@@ -161,7 +178,7 @@ namespace Teknik.Areas.Paste.Controllers
 
                         Response.Headers.Add("Content-Disposition", cd.ToString());
 
-                        return new BufferedFileStreamResult("application/octet-stream", async (response) => await ResponseHelper.StreamToOutput(response, new AesCounterStream(fileStream, false, keyBytes, ivBytes), contentSize, _config.PasteConfig.ChunkSize), false);
+                        return new BufferedFileStreamResult("application/octet-stream", async (response) => await ResponseHelper.StreamToOutput(response, new AesCounterStream(fileStream, false, keyArray, ivArray), contentSize, _config.PasteConfig.ChunkSize), false);
                     default:
                         return View("~/Areas/Paste/Views/Paste/Full.cshtml", model);
                 }
@@ -240,8 +257,25 @@ namespace Teknik.Areas.Paste.Controllers
                 model.DatePosted = paste.DatePosted;
                 model.Username = paste.User?.Username;
 
-                byte[] ivBytes = (string.IsNullOrEmpty(paste.IV)) ? new byte[paste.BlockSize] : Encoding.Unicode.GetBytes(paste.IV);
+                PooledArray ivArray;
-                byte[] keyBytes = (string.IsNullOrEmpty(paste.Key)) ? new byte[paste.KeySize] : AesCounterManaged.CreateKey(paste.Key, ivBytes, paste.KeySize);
+                if (!string.IsNullOrEmpty(paste.IV))
+                {
+                    var ivBytes = Encoding.Unicode.GetBytes(paste.IV);
+                    ivArray = new PooledArray(ivBytes);
+                }
+                else
+                    ivArray = new PooledArray(paste.BlockSize);
+                Response.RegisterForDispose(ivArray);
+
+                PooledArray keyArray;
+                if (!string.IsNullOrEmpty(paste.Key))
+                {
+                    var keyBytes = AesCounterManaged.CreateKey(paste.Key, ivArray.Array, paste.KeySize);
+                    keyArray = new PooledArray(keyBytes);
+                }
+                else
+                    keyArray = new PooledArray(paste.KeySize);
+                Response.RegisterForDispose(keyArray);
 
                 // The paste has a password set
                 if (!string.IsNullOrEmpty(paste.HashedPassword))
@@ -255,7 +289,7 @@ namespace Teknik.Areas.Paste.Controllers
                     if (!string.IsNullOrEmpty(password))
                     {
                         hash = Crypto.HashPassword(paste.Key, password);
-                        keyBytes = AesCounterManaged.CreateKey(password, ivBytes, paste.KeySize);
+                        AesCounterManaged.CreateKey(password, ivArray.Array, paste.KeySize).CopyTo(keyArray.Array, 0);
                     }
                     if (string.IsNullOrEmpty(password) || hash != paste.HashedPassword)
                     {
@@ -285,7 +319,7 @@ namespace Teknik.Areas.Paste.Controllers
                 if (fileStream == null)
                     return new StatusCodeResult(StatusCodes.Status404NotFound);
 
-                using (AesCounterStream cs = new AesCounterStream(fileStream, false, keyBytes, ivBytes))
+                using (AesCounterStream cs = new AesCounterStream(fileStream, false, keyArray, ivArray))
                 using (StreamReader sr = new StreamReader(cs, Encoding.Unicode))
                 {
                     model.Content = await sr.ReadToEndAsync();

Teknik/Areas/Paste/PasteHelper.cs

@@ -120,8 +120,10 @@ namespace Teknik.Areas.Paste
             // Encrypt Content
             byte[] data = Encoding.Unicode.GetBytes(content);
             using (MemoryStream ms = new MemoryStream(data))
+            using (var keyArray = new PooledArray(keyBytes))
+            using (var ivArray = new PooledArray(ivBytes))
             {
-                await storageService.SaveEncryptedFile(fileName, ms, keyBytes, ivBytes);
+                await storageService.SaveEncryptedFile(fileName, ms, keyArray, ivArray);
             }
         }
 

Teknik/Areas/Upload/Controllers/UploadController.cs

@@ -517,8 +517,14 @@ namespace Teknik.Areas.Upload.Controllers
                 {
                     byte[] keyBytes = Encoding.UTF8.GetBytes(key);
                     byte[] ivBytes = Encoding.UTF8.GetBytes(iv);
+                    
+                    var keyArray = new PooledArray(keyBytes);
+                    var ivArray = new PooledArray(ivBytes);
 
-                    var aesStream = new AesCounterStream(fileStream, false, keyBytes, ivBytes);
+                    Response.RegisterForDispose(keyArray);
+                    Response.RegisterForDispose(ivArray);
+
+                    var aesStream = new AesCounterStream(fileStream, false, keyArray, ivArray);
                     //return File(aesStream, contentType, true);
                     return new BufferedFileStreamResult(contentType, async (response) => await ResponseHelper.StreamToOutput(response, aesStream, length, _config.UploadConfig.ChunkSize), false);
                 }

Teknik/Areas/Upload/UploadHelper.cs

@@ -61,7 +61,11 @@ namespace Teknik.Areas.Upload
                 byte[] keyBytes = Encoding.UTF8.GetBytes(key);
                 byte[] ivBytes = Encoding.UTF8.GetBytes(iv);
 
-                await storageService.SaveEncryptedFile(fileName, file, keyBytes, ivBytes);
+                using (var keyArray = new PooledArray(keyBytes))
+                using (var ivArray = new PooledArray(ivBytes))
+                {
+                    await storageService.SaveEncryptedFile(fileName, file, keyArray, ivArray);
+                }
             }
             else
             {

Teknik/Areas/Vault/Controllers/VaultController.cs

@@ -116,7 +116,10 @@ namespace Teknik.Areas.Vault.Controllers
 
                                 byte[] ivBytes = Encoding.Unicode.GetBytes(paste.Paste.IV);
                                 byte[] keyBytes = AesCounterManaged.CreateKey(paste.Paste.Key, ivBytes, paste.Paste.KeySize);
-                                using (AesCounterStream cs = new AesCounterStream(fileStream, false, keyBytes, ivBytes))
+
+                                using (var keyArray = new PooledArray(keyBytes))
+                                using (var ivArray = new PooledArray(ivBytes))
+                                using (AesCounterStream cs = new AesCounterStream(fileStream, false, keyArray, ivArray))
                                 using (StreamReader sr = new StreamReader(cs, Encoding.Unicode))
                                 {
                                     pasteModel.Content = await sr.ReadToEndAsync();

Utilities/Cryptography/AesCounterManaged.cs

@@ -9,7 +9,7 @@ namespace Teknik.Utilities.Cryptography
 {
     public class AesCounterManaged
     {
-        public static async Task EncryptToFile(Stream input, string filePath, byte[] key, byte[] iv)
+        public static async Task EncryptToFile(Stream input, string filePath, PooledArray key, PooledArray iv)
         {
             using (FileStream fileStream = new FileStream(filePath, FileMode.Create, FileAccess.Write))
             {
@@ -17,7 +17,7 @@ namespace Teknik.Utilities.Cryptography
             }
         }
 
-        public static async Task EncryptToStream(Stream input, Stream output, byte[] key, byte[] iv)
+        public static async Task EncryptToStream(Stream input, Stream output, PooledArray key, PooledArray iv)
         {
             // Make sure the input stream is at the beginning
             if (input.CanSeek)

Utilities/Cryptography/AesCounterMode.cs

@@ -7,17 +7,17 @@ namespace Teknik.Utilities.Cryptography
     {
         // Internal Variables
         private const int _BlockSize = 16;
-        private readonly byte[] _InitialCounter;
+        private readonly PooledArray _InitialCounter;
         private readonly Aes _Algo;
 
-        public AesCounterMode() : this(new byte[_BlockSize]) { }
+        public AesCounterMode() : this(new PooledArray(_BlockSize)) { }
 
-        public AesCounterMode(byte[] initialCounter)
+        public AesCounterMode(PooledArray initialCounter)
         {
             if (initialCounter == null) throw new ArgumentNullException("counter");
-            if (initialCounter.Length != _BlockSize)
+            if (initialCounter.Array.Length != _BlockSize)
                 throw new ArgumentException(String.Format("Counter size must be same as block size (actual: {0}, expected: {1})",
-                    initialCounter.Length, _BlockSize));
+                    initialCounter.Array.Length, _BlockSize));
 
             // Generate a new instance of the Aes Algorithm in ECB mode with no padding
             _Algo = Aes.Create();
@@ -57,9 +57,9 @@ namespace Teknik.Utilities.Cryptography
     public class CounterModeCryptoTransform : ICryptoTransform, IDisposable
     {
         private readonly int _BlockSize;
-        private readonly byte[] _IV;
+        private readonly PooledArray _IV;
-        private readonly byte[] _Counter;
+        private readonly PooledArray _Counter;
-        private readonly byte[] _EncryptedCounter;
+        private readonly PooledArray _EncryptedCounter;
         private readonly ICryptoTransform _CounterEncryptor;
 
         // Stateful Fields
@@ -82,14 +82,14 @@ namespace Teknik.Utilities.Cryptography
             }
             set
             {
-                if (value >= 0 && value < _EncryptedCounter.Length)
+                if (value >= 0 && value < _EncryptedCounter.Array.Length)
                 {
                     _CounterPosition = value;
                 }
             }
         }
 
-        public CounterModeCryptoTransform(SymmetricAlgorithm symmetricAlgorithm, byte[] key, byte[] iv, byte[] initialCounter)
+        public CounterModeCryptoTransform(SymmetricAlgorithm symmetricAlgorithm, byte[] key, byte[] iv, PooledArray initialCounter)
         {
             if (symmetricAlgorithm == null) throw new ArgumentNullException("symmetricAlgorithm");
             if (key == null) throw new ArgumentNullException("key");
@@ -97,22 +97,22 @@ namespace Teknik.Utilities.Cryptography
             if (initialCounter == null) throw new ArgumentNullException("counter");
 
             // Check lengths
-            if (initialCounter.Length != symmetricAlgorithm.BlockSize / 8)
+            if (initialCounter.Array.Length != symmetricAlgorithm.BlockSize / 8)
                 throw new ArgumentException(String.Format("Counter size must be same as block size (actual: {0}, expected: {1})",
-                    initialCounter.Length, symmetricAlgorithm.BlockSize / 8));
+                    initialCounter.Array.Length, symmetricAlgorithm.BlockSize / 8));
 
             _BlockSize = symmetricAlgorithm.BlockSize;
 
             // Initialize Counter
-            _Counter = new byte[initialCounter.Length];
+            _Counter = new PooledArray(initialCounter.Array.Length);
-            initialCounter.CopyTo(_Counter, 0);
+            initialCounter.Array.CopyTo(_Counter.Array, 0);
 
             // Initialize the encrypted counter
-            _EncryptedCounter = new byte[_BlockSize / 8];
+            _EncryptedCounter = new PooledArray(_BlockSize / 8);
 
             // Initialize IV
-            _IV = new byte[iv.Length];
+            _IV = new PooledArray(iv.Length);
-            iv.CopyTo(_IV, 0);
+            iv.CopyTo(_IV.Array, 0);
 
             _CounterEncryptor = symmetricAlgorithm.CreateEncryptor(key, iv);
 
@@ -144,7 +144,7 @@ namespace Teknik.Utilities.Cryptography
             for (var i = 0; i < inputCount; i++)
             {
                 // Encrypt the counter if we have reached the end, or 
-                if (_CounterPosition >= _EncryptedCounter.Length)
+                if (_CounterPosition >= _EncryptedCounter.Array.Length)
                 {
                     //Reset current counter position
                     _CounterPosition = 0;
@@ -157,7 +157,7 @@ namespace Teknik.Utilities.Cryptography
                 }
 
                 // XOR the encrypted counter with the input plain text
-                inputBuffer[inputOffset + i] = (byte)(_EncryptedCounter[_CounterPosition] ^ inputBuffer[inputOffset + i]);
+                inputBuffer[inputOffset + i] = (byte)(_EncryptedCounter.Array[_CounterPosition] ^ inputBuffer[inputOffset + i]);
 
                 // Move the counter position
                 _CounterPosition++;
@@ -178,7 +178,7 @@ namespace Teknik.Utilities.Cryptography
             for (var i = 0; i < inputCount; i++)
             {
                 // Encrypt the counter if we have reached the end, or 
-                if (_CounterPosition >= _EncryptedCounter.Length)
+                if (_CounterPosition >= _EncryptedCounter.Array.Length)
                 {
                     //Reset current counter position
                     _CounterPosition = 0;
@@ -191,7 +191,7 @@ namespace Teknik.Utilities.Cryptography
                 }
 
                 // XOR the encrypted counter with the input plain text
-                outputBuffer[outputOffset + i] = (byte)(_EncryptedCounter[_CounterPosition] ^ inputBuffer[inputOffset + i]);
+                outputBuffer[outputOffset + i] = (byte)(_EncryptedCounter.Array[_CounterPosition] ^ inputBuffer[inputOffset + i]);
 
                 // Move the counter position
                 _CounterPosition++;
@@ -203,19 +203,19 @@ namespace Teknik.Utilities.Cryptography
         public void EncryptCounter()
         {
             // Encrypt the current counter to the encrypted counter
-            _CounterEncryptor.TransformBlock(_Counter, 0, _Counter.Length, _EncryptedCounter, 0);
+            _CounterEncryptor.TransformBlock(_Counter.Array, 0, _Counter.Array.Length, _EncryptedCounter.Array, 0);
         }
 
         public void ResetCounter()
         {
-            _IV.CopyTo(_Counter, 0);
+            _IV.Array.CopyTo(_Counter.Array, 0);
             _Iterations = 0;
         }
 
         public void IncrementCounter()
         {
-            int j = _Counter.Length;
+            int j = _Counter.Array.Length;
-            while (--j >= 0 && ++_Counter[j] == 0)
+            while (--j >= 0 && ++_Counter.Array[j] == 0)
             {
             }
             _Iterations++;
@@ -229,6 +229,9 @@ namespace Teknik.Utilities.Cryptography
         public void Dispose()
         {
             _CounterEncryptor.Dispose();
+            _IV.Dispose();
+            _Counter.Dispose();
+            _EncryptedCounter.Dispose();
         }
     }
 }

Utilities/Cryptography/AesCounterStream.cs

@@ -19,7 +19,7 @@ namespace Teknik.Utilities.Cryptography
         /// <param name="encrypt"></param>
         /// <param name="key"></param>
         /// <param name="iv"></param>
-        public AesCounterStream(Stream stream, bool encrypt, byte[] key, byte[] iv)
+        public AesCounterStream(Stream stream, bool encrypt, PooledArray key, PooledArray iv)
         {
             _Inner = stream;
 
@@ -27,11 +27,11 @@ namespace Teknik.Utilities.Cryptography
             using AesCounterMode aes = new AesCounterMode(iv);
             if (encrypt)
             {
-                _Cipher = (CounterModeCryptoTransform)aes.CreateEncryptor(key, iv); // Encrypt
+                _Cipher = (CounterModeCryptoTransform)aes.CreateEncryptor(key.Array, iv.Array); // Encrypt
             }
             else
             {
-                _Cipher = (CounterModeCryptoTransform)aes.CreateDecryptor(key, iv); // Decrypt
+                _Cipher = (CounterModeCryptoTransform)aes.CreateDecryptor(key.Array, iv.Array); // Decrypt
             }
 
             // Sync the counter

Utilities/PooledArray.cs

@@ -18,6 +18,12 @@ namespace Teknik.Utilities
             Array = _arrayPool.Rent(size);
         }
 
+        public PooledArray(byte[] array)
+        {
+            Array = _arrayPool.Rent(array.Length);
+            array.CopyTo(Array, 0);
+        }
+
         public void Dispose()
         {
             _arrayPool.Return(Array);