src/MassTransit/Serialization/AesCryptoStreamProviderV2.cs - develop - vshmidt/masstransit

src/MassTransit/Serialization/AesCryptoStreamProviderV2.cs

85 строк2 KB

Chris Patterson

Added .NET 9 support (Fixed #5651)

14 ноя 2024, 17:30

14 ноя 2024, 17:30f9ae6a6

namespace MassTransit.Serialization
{
    using System;
    using System.IO;
    using System.Security.Cryptography;


    public class AesCryptoStreamProviderV2 :
        ICryptoStreamProviderV2
    {
        readonly PaddingMode _paddingMode;
        readonly ISecureKeyProvider _secureKeyProvider;

        public AesCryptoStreamProviderV2(ISecureKeyProvider secureKeyProvider, PaddingMode paddingMode = PaddingMode.PKCS7)
        {
            _secureKeyProvider = secureKeyProvider;
            _paddingMode = paddingMode;
        }

        public Stream GetDecryptStream(Stream stream, Headers headers)
        {
            var key = _secureKeyProvider.GetKey(headers);

            var iv = new byte[16];
            var read = stream.Read(iv, 0, iv.Length);
            if (read != iv.Length)
                throw new InvalidOperationException("The stream does not contain enough bytes to read the IV.");

            var encryptor = CreateDecryptor(key, iv);

            return new DisposingCryptoStream(stream, encryptor, CryptoStreamMode.Read);
        }

        public Stream GetEncryptStream(Stream stream, Headers headers)
        {
            var key = _secureKeyProvider.GetKey(headers);

            var iv = GenerateIv();

            stream.Write(iv, 0, iv.Length);
            var encryptor = CreateEncryptor(key, iv);

            return new DisposingCryptoStream(stream, encryptor, CryptoStreamMode.Write);
        }

        public void Probe(ProbeContext context)
        {
            var scope = context.CreateScope("aes");

            scope.Add("paddingMode", _paddingMode.ToString());

            _secureKeyProvider.Probe(scope);
        }

        ICryptoTransform CreateDecryptor(byte[] key, byte[] iv)
        {
            using var provider = CreateAes();

            return provider.CreateDecryptor(key, iv);
        }

        byte[] GenerateIv()
        {
            using var aes = CreateAes();

            aes.GenerateIV();

            return aes.IV;
        }

        ICryptoTransform CreateEncryptor(byte[] key, byte[] iv)
        {
            using var provider = CreateAes();

            return provider.CreateEncryptor(key, iv);
        }

        Aes CreateAes()
        {
            var aes = Aes.Create();
            aes.Padding = _paddingMode;
            return aes;
        }
    }
}