Modified to pluck tag from end of stream automatically

BCryptAuthenticatedSymmetricCryptoTransform.cs

// Copyright (c) Microsoft Corporation.  All rights reserved.
//
// Originally from CLR Security licensed by Microsoft under MIT License
// Modified by Crypteron under MIT License
//
// changelog
// - Make decrypt symmetrical with encrypt by keeping tag at the end 
//         of ciphertext. This oddity requires some buffer management 
//         but better to do internally here than expect all users to 
//         perform tricky buffer management themselves.

using System;
using System.Diagnostics;
using System.IO;
using System.Runtime.ConstrainedExecution;
using System.Security;
using System.Security.Cryptography;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;

namespace Security.Cryptography
{
    /// <summary>
    ///     Generic crypto transform, which implements authenticated symmetric encryption and decryption for
    ///     algorithms implemented in the BCrypt layer of CNG.  This type is used as the workhorse for the
    ///     BCryptAuthenticatedSymmetricAlgorithm generic BCrypt authenticated symmetric algorithm
    ///     implementation.
    /// </summary>
    internal sealed class BCryptAuthenticatedSymmetricCryptoTransform : CriticalFinalizerObject,
                                                                        IAuthenticatedCryptoTransform,
                                                                        IDisposable
    {
        private SafeBCryptAlgorithmHandle m_algorithm;
        private BCryptNative.BCRYPT_AUTHENTICATED_CIPHER_MODE_INFO m_authInfo;
        private byte[] m_chainData;
        private bool m_chainingSupported;
        private SafeBCryptKeyHandle m_key;

        private MemoryStream m_inputBuffer;
        private bool m_encrypting;
        private bool m_transformedFinalBlock;
        private byte[] tailBuf;
        private int tailBufValidCount;
        private byte[] inputExcessBuf;
        private int inputExcessBufValid;


        /// <summary>
        ///     Create an encrypting authenticated symmetric algorithm transform.  This type takes ownership
        ///     of the incoming algorithm handle, which should no longer be used by the calling code after
        ///     it has called this constructor.
        /// </summary>
        [SecurityCritical]
        internal BCryptAuthenticatedSymmetricCryptoTransform(SafeBCryptAlgorithmHandle algorithm,
                                                             byte[] key,
                                                             byte[] nonce,
                                                             byte[] authenticatedData,
                                                             bool chainingSupported,
                                                             int tagSize) :
            this(algorithm, key, nonce, authenticatedData, new byte[tagSize / 8], chainingSupported)
        {
            m_encrypting = true;
        }

        /// <summary>
        ///     Create a decrypting authenticated symmetric algorithm transform.  This type takes ownership
        ///     of the incoming algorithm handle, which should no longer be used by the calling code after
        ///     it has called this constructor.
        /// </summary>
        [SecurityCritical]
        internal BCryptAuthenticatedSymmetricCryptoTransform(SafeBCryptAlgorithmHandle algorithm,
                                                             byte[] key,
                                                             byte[] nonce,
                                                             byte[] authenticatedData,
                                                             byte[] tag,
                                                             bool chainingSupported)
        {
            Debug.Assert(algorithm != null, "algorithm != null");
            Debug.Assert(!algorithm.IsClosed && !algorithm.IsInvalid, "!algorithm.IsClosed && !algorithm.IsInvalid");

            if (key == null)
                throw new ArgumentNullException("key");
            if (tag == null)
                throw new ArgumentNullException("tag");

            bool initializationComplete = false;

            RuntimeHelpers.PrepareConstrainedRegions();
            try
            {
                m_algorithm = algorithm;
                m_key = BCryptNative.ImportSymmetricKey(algorithm, key);

                // Initialize the padding info structure.
                m_authInfo = new BCryptNative.BCRYPT_AUTHENTICATED_CIPHER_MODE_INFO();
                BCryptNative.InitializeAuthnenticatedCipherModeInfo(ref m_authInfo);

                if (nonce != null)
                {
                    m_authInfo.cbNonce = nonce.Length;
                    m_authInfo.pbNonce = Marshal.AllocCoTaskMem(m_authInfo.cbNonce);
                    Marshal.Copy(nonce, 0, m_authInfo.pbNonce, m_authInfo.cbNonce);
                }

                if (authenticatedData != null)
                {
                    m_authInfo.cbAuthData = authenticatedData.Length;
                    m_authInfo.pbAuthData = Marshal.AllocCoTaskMem(m_authInfo.cbAuthData);
                    Marshal.Copy(authenticatedData, 0, m_authInfo.pbAuthData, m_authInfo.cbAuthData);
                }

                if (chainingSupported)
                {
                    m_chainingSupported = chainingSupported;

                    m_authInfo.cbMacContext = tag.Length;
                    m_authInfo.pbMacContext = Marshal.AllocCoTaskMem(m_authInfo.cbMacContext);

                    BCryptNative.BCRYPT_KEY_LENGTHS_STRUCT tagLengths =
                        BCryptNative.GetValueTypeProperty<SafeBCryptAlgorithmHandle, BCryptNative.BCRYPT_KEY_LENGTHS_STRUCT>(
                            algorithm,
                            BCryptNative.ObjectPropertyName.AuthTagLength);

                    m_chainData = new byte[tagLengths.dwMaxLength];
                }
                else
                {
                    m_inputBuffer = new MemoryStream();
                }

                m_authInfo.cbTag = tag.Length;
                m_authInfo.pbTag = Marshal.AllocCoTaskMem(m_authInfo.cbTag);
                Marshal.Copy(tag, 0, m_authInfo.pbTag, m_authInfo.cbTag);

                // Set chaining mode if supported.
                if (CanChainBlocks)
                {
                    m_authInfo.dwFlags |= BCryptNative.AuthenticatedCipherModeInfoFlags.ChainCalls;
                }

                tailBuf = new byte[tag.Length];
                tailBufValidCount = 0;
                inputExcessBuf = new byte[InputBlockSize];
                inputExcessBufValid = 0;

                initializationComplete = true;
            }
            finally
            {
                // If we failed to complete initialization we may have already allocated some native
                // resources.  Clean those up before leaving the constructor.
                if (!initializationComplete)
                {
                    Dispose();
                }
            }
        }

        ~BCryptAuthenticatedSymmetricCryptoTransform()
        {
            Dispose(false);
        }

        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        [SecurityCritical]
        [SecuritySafeCritical]
        [ReliabilityContract(Consistency.MayCorruptInstance, Cer.Success)]
        private void Dispose(bool disposing)
        {
            if (disposing)
            {
                if (m_key != null)
                {
                    m_key.Dispose();
                }

                if (m_algorithm != null)
                {
                    m_algorithm.Dispose();
                }

                if (m_inputBuffer != null)
                {
                    m_inputBuffer.Dispose();
                }
            }

            if (m_authInfo.pbAuthData != IntPtr.Zero)
            {
                Marshal.FreeCoTaskMem(m_authInfo.pbAuthData);
                m_authInfo.pbAuthData = IntPtr.Zero;
                m_authInfo.cbAuthData = 0;
            }

            if (m_authInfo.pbMacContext != IntPtr.Zero)
            {
                Marshal.FreeCoTaskMem(m_authInfo.pbMacContext);
                m_authInfo.pbMacContext = IntPtr.Zero;
                m_authInfo.cbMacContext = 0;
            }

            if (m_authInfo.pbNonce != IntPtr.Zero)
            {
                Marshal.FreeCoTaskMem(m_authInfo.pbNonce);
                m_authInfo.pbNonce = IntPtr.Zero;
                m_authInfo.cbNonce = 0;
            }

            if (m_authInfo.pbTag != IntPtr.Zero)
            {
                Marshal.FreeCoTaskMem(m_authInfo.pbTag);
                m_authInfo.pbTag = IntPtr.Zero;
                m_authInfo.cbTag = 0;
            }
        }

        /// <summary>
        ///     Can the transform chain multiple blocks of ciphertext, or must they all come at once.
        /// </summary>
        public bool CanChainBlocks
        {
            get { return m_chainingSupported; }
        }

        /// <summary>
        ///     Gets a value indicating whether the transform can be reused.
        /// </summary>
        public bool CanReuseTransform
        {
            get { return false; }
        }

        /// <summary>
        ///     Gets a value indicating whether the transform can process multiple blocks at once.
        /// </summary>
        public bool CanTransformMultipleBlocks
        {
            get { return true; }
        }

        /// <summary>
        ///     Gets the input block length in bytes.
        /// </summary>
        public int InputBlockSize
        {
            [SecurityCritical]
            [SecuritySafeCritical]
            get { return BCryptNative.GetInt32Property(m_algorithm, BCryptNative.ObjectPropertyName.BlockLength); }
        }

        /// <summary>
        ///     Gets the output block length in bytes.
        /// </summary>
        public int OutputBlockSize
        {
            [SecurityCritical]
            [SecuritySafeCritical]
            get { return BCryptNative.GetInt32Property(m_algorithm, BCryptNative.ObjectPropertyName.BlockLength); }
        }

        /// <summary>
        ///     Get the authentication tag generated from encryption.
        /// </summary>
        [SecurityCritical]
        [SecuritySafeCritical]
        public byte[] GetTag()
        {
            // Authentication tags are only generated for encryption operations - they are input to decryption
            // operations.  They are also only generated after all of the data has been transformed.
            if (!m_encrypting)
                throw new InvalidOperationException(Resources.TagIsOnlyGeneratedDuringEncryption);
            if (!m_transformedFinalBlock)
                throw new InvalidOperationException(Resources.TagIsOnlyGeneratedAfterFinalBlock);

            byte[] tag = new byte[m_authInfo.cbTag];
            Marshal.Copy(m_authInfo.pbTag, tag, 0, m_authInfo.cbTag);
            return tag;
        }

        [SecurityCritical]
        [SecuritySafeCritical]
        public void SetExpectedTag(byte[] tag)
        {
            if (m_transformedFinalBlock)
                throw new InvalidOperationException(Resources.TagMustBeSuppliedBeforeFinalBlock);
            if (m_encrypting)
                throw new InvalidOperationException("Setting the expected tag is valid only during decryption");

            m_authInfo.cbTag = tag.Length;
            m_authInfo.pbTag = Marshal.AllocCoTaskMem(m_authInfo.cbTag);
            Marshal.Copy(tag, 0, m_authInfo.pbTag, m_authInfo.cbTag);
        }

        /// <summary>
        /// Transforms some blocks of input data, but don't finalize the transform. We make this into a wrapper
        /// so that we can extract the tag from the end and process it separately. We basically flow the 
        /// data through, trapping the last extractByteCount bytes inside a tailBuf
        /// </summary>
        public int TransformBlock(byte[] inputBuffer, int inputOffset, int inputCount, byte[] outputBuffer, int outputOffset)
        {
            if (m_encrypting)
            {
                // Encryption is straightforward
                return TransformBlockInternal(inputBuffer, inputOffset, inputCount, outputBuffer, outputOffset);
            }
            else
            {
                // Decryption is handled differently since we have to extract tag from within the stream
                var tagByteCount = m_authInfo.cbTag;

                //////////////////////////////////////////////////////////////////////////////////////
                // Step1: Map: how much of what lives where. 
                // Transform everything except potential tag i.e. all but last tagByteCount bytes
                int totalTransformCount = Math.Max(0, tailBufValidCount + inputCount - tagByteCount);
                // constrain Transform count on tailBuf
                int tailBufTransformCount = Math.Min(tailBufValidCount, totalTransformCount);
                // whatever isn't transformed from total valid, goes to tailBuf
                int tailBufTailCount = tailBufValidCount - tailBufTransformCount;
                // from the total Transforms, whatever isn't coming from tailBuf's Transforms are inputBuffer's Transforms
                int inputBufferTransformCount = totalTransformCount - tailBufTransformCount;
                // whatever isn't transformed from input, goes to tailBuf
                int inputBufferTailCount = inputCount - inputBufferTransformCount;
                // double check our math - tailBuf can't have tag bytes if inputBuffer still has cipher-bytes
                if (tailBufTailCount > 0 && inputBufferTransformCount > 0)
                {
                    var errDetails = $"input:: Len:{inputBuffer.Length},Off:{inputOffset},Valid:{inputCount}[Transform:{inputBufferTransformCount} Tail:{inputBufferTailCount}] => " + 
                                   $"=> tailBuf:: Len:{tailBuf.Length}, Off:0, Valid:{tailBufValidCount}[Transform:{tailBufTransformCount} Tail:{tailBufTailCount}] => ";
                    throw new InvalidOperationException($"Error in buffer math. Details:{errDetails}");
                }

                //////////////////////////////////////////////////////////////////////////////////////
                // Step2: Gather ciphertext from end-to-end from various buffers, in-order
                // allocate buffer for that
                var tempDecTransformBuf = new byte[totalTransformCount + inputExcessBufValid];
                // First pack excess bytes trimmed off from last invocation (from inputExcessBuf)
                Array.Copy(inputExcessBuf, 0, tempDecTransformBuf, 0, inputExcessBufValid);
                // Then pack bytes from tailBuf (from last invocation), now known to be ciphertext, not tag
                Array.Copy(tailBuf, 0, tempDecTransformBuf, inputExcessBufValid, tailBufTransformCount);
                // Finally pack bytes just received from inputBuffer guaranteed to be ciphertext
                Array.Copy(inputBuffer, inputOffset, tempDecTransformBuf, tailBufTransformCount + inputExcessBufValid, inputBufferTransformCount);

                //////////////////////////////////////////////////////////////////////////////////////
                // Step 3: Transform ciphertext bytes in whole InputBlockSize
                int excessBytes = 0;
                if (tempDecTransformBuf.Length > 0)
                    excessBytes = tempDecTransformBuf.Length % InputBlockSize;
                int transformBytes = tempDecTransformBuf.Length - excessBytes;
                int processedBytes = 0;

                if (transformBytes > 0)
                    processedBytes = TransformBlockInternal(tempDecTransformBuf, 0, transformBytes, outputBuffer, outputOffset);

                //////////////////////////////////////////////////////////////////////////////////////
                // Step4: Cleanup: Copy leftover ciphertext bytes and pack potential tag bytes
                // 4a Copy leftover bytes guaranteed to be ciphertext (not tag) into inputExcessBuf 
                // for future Transforms 
                Array.Copy(tempDecTransformBuf, tempDecTransformBuf.Length - excessBytes, inputExcessBuf, 0, excessBytes);
                inputExcessBufValid = excessBytes;
                // Pack bytes into tailBuf that can potentially be tag bytes
                var tempTailBuf = new byte[tagByteCount];
                // 4b First pack oldest bytes from tailBuf that didn't get Transform'd above
                Array.Copy(tailBuf, tailBufTransformCount, tempTailBuf, 0, tailBufTailCount);
                // 4c Then pack ending bytes from inputBuffer that didn't get Transform'd above
                Array.Copy(inputBuffer, inputOffset + inputBufferTransformCount, tempTailBuf, tailBufTailCount, inputBufferTailCount);
                tailBufValidCount = tailBufTailCount + inputBufferTailCount;
                tailBuf = tempTailBuf; // rename buffer instead of copying

                return processedBytes;
            }
        }

        /// <summary>
        ///     Transforms some blocks of input data, but don't finalize the transform
        /// </summary>
        private int TransformBlockInternal(byte[] inputBuffer, int inputOffset, int inputCount, byte[] outputBuffer, int outputOffset)
        {
            if (m_transformedFinalBlock)
                throw new InvalidOperationException("This transform has already transformed its final block and can no longer transform additional data.");
            if (inputBuffer == null)
                throw new ArgumentNullException("inputBuffer");
            if (inputOffset < 0)
                throw new ArgumentOutOfRangeException("inputOffset");
            if (inputCount <= 0)
                throw new ArgumentOutOfRangeException("inputCount");
            if (inputCount % InputBlockSize != 0)
                throw new ArgumentOutOfRangeException("inputCount");
            if (inputCount > inputBuffer.Length - inputOffset)
                throw new ArgumentOutOfRangeException("inputCount");
            if (outputBuffer == null)
                throw new ArgumentNullException("outputBuffer");
            if (inputCount > outputBuffer.Length - outputOffset)
                throw new ArgumentOutOfRangeException("outputOffset");

            // If the transform can chain multiple blocks of data, then transform the input now.  Otherwise,
            // save it away to be transformed when TransformFinalBlock is called.
            if (CanChainBlocks)
            {
                byte[] transformed = null;
                try
                {
                    transformed = CngTransform(inputBuffer, inputOffset, inputCount);
                    Array.Copy(transformed, 0, outputBuffer, outputOffset, transformed.Length);
                    return transformed.Length;
                }
                finally
                {
                    if (transformed != null)
                    {
                        Array.Clear(transformed, 0, transformed.Length);
                    }
                }
            }
            else
            {
                m_inputBuffer.Write(inputBuffer, inputOffset, inputCount);
                return 0;
            }
        }

        /// <summary>
        ///     Transform the final block and finalize the encryption or decryption operation.
        /// </summary>
        public byte[] TransformFinalBlock(byte[] inputBuffer, int inputOffset, int inputCount)
        {
            if (inputBuffer == null)
                throw new ArgumentNullException("inputBuffer");
            if (inputOffset < 0)
                throw new ArgumentOutOfRangeException("inputOffset");
            if (inputCount < 0)
                throw new ArgumentOutOfRangeException("inputCount");
            if (inputCount > inputBuffer.Length - inputOffset)
                throw new ArgumentOutOfRangeException("inputCount");

            if (!m_transformedFinalBlock)
            {
                var finalInputBlockOutput = new byte[0];
                int finalInputBlockOutputValid = 0;

                // 4 main flows: encrypt+CanChainBlock, encrypt+!CanChainBlock, decrypt+CanChainBlock, decrypt+!CanChainBlock,
                // Decrypt
                if (!m_encrypting)
                {
                    // Transform Input buffer, pushing it through tailBuf
                    finalInputBlockOutput = new byte[OutputBlockSize];
                    finalInputBlockOutputValid = TransformBlock(inputBuffer, inputOffset, inputCount, finalInputBlockOutput, 0);

                    // At this point, we MAY have left over ciphertext-bytes in inputExcessBuf
                    if (CanChainBlocks)
                    {
                        // Transform those by variable reassignment and falling thru to the CngTransform below
                        inputBuffer = inputExcessBuf;
                        inputOffset = 0;
                        inputCount = inputExcessBufValid;
                    } 
                    else
                    {
                        // queue those for Transformation by pushing into m_inputBuffer
                        m_inputBuffer.Write(inputExcessBuf, 0, inputExcessBufValid);
                        // if !CanChainBlocks, we don't really have output bytes, we've only
                        // queued all bytes
                        finalInputBlockOutputValid = 0;
                    }

                    // At this point, tailBuf should only contain Tag, program it into the encryptor
                    if (tailBufValidCount == m_authInfo.cbTag)
                        SetExpectedTag(tailBuf);
                    else
                    {
                        string errorMsg =
                            String.Format("Tag buffer validity bytes ({0}) don't match total tag length ({1} in bytes)",
                                          tailBufValidCount, m_authInfo.cbTag);
                        throw new CryptographicException(errorMsg);
                    }
                }

                // Remove the chaining call flag, but retain the rest.
                m_authInfo.dwFlags &= ~BCryptNative.AuthenticatedCipherModeInfoFlags.ChainCalls;
                m_transformedFinalBlock = true;

                // If we cannot chain multiple blocks of data, then add the final block to the other input
                // blocks we've already collected, and use that to transform
                if (!CanChainBlocks)
                {
                    // This code path is for unsupported configurations
                    //throw new NotSupportedException("Chaining multiple blocks of data is disabled. Currently unsupported configuration, cannot proceed until block chaining is enabled.");
                    if (m_encrypting)
                    {
                        m_inputBuffer.Write(inputBuffer, inputOffset, inputCount);
                    }
                    // Reassign the input to be the full set of data that we've already gathered across all
                    // calls into the stream.
                    inputBuffer = m_inputBuffer.ToArray();
                    inputOffset = 0;
                    inputCount = inputBuffer.Length;
                }

                var retVal = CngTransform(inputBuffer, inputOffset, inputCount);

                if (m_encrypting)
                {
                    // We combine the tag here and add to end of the stream
                    var tag = GetTag();
                    var result = new byte[retVal.Length + tag.Length];
                    Array.Copy(retVal, result, retVal.Length);
                    Array.Copy(tag, 0, result, retVal.Length, tag.Length);

                    return result;
                }
                else
                {
                    // Decrypting, final block involves combining finalInputBlockOutput (decrypted at start of method)
                    // and retVal (decrypted just above)
                    var result = new byte[retVal.Length + finalInputBlockOutputValid];
                    Array.Copy(finalInputBlockOutput, result, finalInputBlockOutputValid);
                    Array.Copy(retVal, 0, result, finalInputBlockOutputValid, retVal.Length);

                    return retVal;
                }
            }
            else
            {
                // We don't want to throw if we're re-flushing the final block, because if the crypto stream
                // was used in a try/finally block, and CngTransform throws an exception we'll end up
                // flushing again in the dispose.  That will end up covering the orginal exception with a
                // less useful re-flushing the final block exception.  Instead, make the call a no-op.
                return new byte[0];
            }
        }

        /// <summary>
        ///     Transform given blocks of data
        /// </summary>
        [SecurityCritical]
        [SecuritySafeCritical]
        private byte[] CngTransform(byte[] input, int inputOffset, int inputCount)
        {
            Debug.Assert(m_key != null, "key != null");
            Debug.Assert(!m_key.IsClosed && !m_key.IsInvalid, "!m_key.IsClosed && !m_key.IsInvalid");
            Debug.Assert(input != null, "input != null");
            Debug.Assert(inputOffset >= 0, "inputOffset >= 0");
            Debug.Assert(inputCount >= 0, "inputCount >= 0");
            Debug.Assert(inputCount <= input.Length - inputOffset, "inputCount <= input.Length - inputOffset");

            byte[] inputBuffer = null;
            try
            {
                // Build up a buffer of the only portion of the input we should be transforming
                inputBuffer = input;
                if (inputOffset > 0 || inputCount != input.Length)
                {
                    inputBuffer = new byte[inputCount];
                    Array.Copy(input, inputOffset, inputBuffer, 0, inputBuffer.Length);
                }

                if (m_encrypting)
                {
                    return BCryptNative.SymmetricEncrypt(m_key, inputBuffer, m_chainData, ref m_authInfo);
                }
                else
                {
                    return BCryptNative.SymmetricDecrypt(m_key, inputBuffer, this.m_chainData, ref m_authInfo);
                }
            }
            finally
            {
                if (inputBuffer != input)
                {
                    // Zeroize the input buffer if we allocated one
                    Array.Clear(inputBuffer, 0, inputBuffer.Length);
                }
            }
        }
    }
}