CBDD/src/CBDD.Core/Compression/CompressionService.cs

using System.Buffers;
using System.Collections.Concurrent;
using System.IO.Compression;

namespace ZB.MOM.WW.CBDD.Core.Compression;

/// <summary>
/// Compression codec registry and utility service.
/// </summary>
public sealed class CompressionService
{
    private readonly ConcurrentDictionary<CompressionCodec, ICompressionCodec> _codecs = new();

    public CompressionService(IEnumerable<ICompressionCodec>? additionalCodecs = null)
    {
        RegisterCodec(new NoneCompressionCodec());
        RegisterCodec(new BrotliCompressionCodec());
        RegisterCodec(new DeflateCompressionCodec());

        if (additionalCodecs == null)
            return;

        foreach (var codec in additionalCodecs)
        {
            RegisterCodec(codec);
        }
    }

    public void RegisterCodec(ICompressionCodec codec)
    {
        ArgumentNullException.ThrowIfNull(codec);
        _codecs[codec.Codec] = codec;
    }

    public bool TryGetCodec(CompressionCodec codec, out ICompressionCodec compressionCodec)
    {
        return _codecs.TryGetValue(codec, out compressionCodec!);
    }

    public ICompressionCodec GetCodec(CompressionCodec codec)
    {
        if (_codecs.TryGetValue(codec, out var compressionCodec))
            return compressionCodec;

        throw new InvalidOperationException($"Compression codec '{codec}' is not registered.");
    }

    public byte[] Compress(ReadOnlySpan<byte> input, CompressionCodec codec, CompressionLevel level)
    {
        return GetCodec(codec).Compress(input, level);
    }

    public byte[] Decompress(ReadOnlySpan<byte> input, CompressionCodec codec, int expectedLength, int maxDecompressedSizeBytes)
    {
        return GetCodec(codec).Decompress(input, expectedLength, maxDecompressedSizeBytes);
    }

    public byte[] Roundtrip(ReadOnlySpan<byte> input, CompressionCodec codec, CompressionLevel level, int maxDecompressedSizeBytes)
    {
        var compressed = Compress(input, codec, level);
        return Decompress(compressed, codec, input.Length, maxDecompressedSizeBytes);
    }

    private sealed class NoneCompressionCodec : ICompressionCodec
    {
        public CompressionCodec Codec => CompressionCodec.None;

        public byte[] Compress(ReadOnlySpan<byte> input, CompressionLevel level) => input.ToArray();

        public byte[] Decompress(ReadOnlySpan<byte> input, int expectedLength, int maxDecompressedSizeBytes)
        {
            if (input.Length > maxDecompressedSizeBytes)
                throw new InvalidDataException($"Decompressed payload exceeds max allowed size ({maxDecompressedSizeBytes} bytes).");

            if (expectedLength >= 0 && expectedLength != input.Length)
                throw new InvalidDataException($"Expected decompressed length {expectedLength}, actual {input.Length}.");

            return input.ToArray();
        }
    }

    private sealed class BrotliCompressionCodec : ICompressionCodec
    {
        public CompressionCodec Codec => CompressionCodec.Brotli;

        public byte[] Compress(ReadOnlySpan<byte> input, CompressionLevel level)
        {
            return CompressWithCodecStream(input, stream => new BrotliStream(stream, level, leaveOpen: true));
        }

        public byte[] Decompress(ReadOnlySpan<byte> input, int expectedLength, int maxDecompressedSizeBytes)
        {
            return DecompressWithCodecStream(input, stream => new BrotliStream(stream, CompressionMode.Decompress, leaveOpen: true), expectedLength, maxDecompressedSizeBytes);
        }
    }

    private sealed class DeflateCompressionCodec : ICompressionCodec
    {
        public CompressionCodec Codec => CompressionCodec.Deflate;

        public byte[] Compress(ReadOnlySpan<byte> input, CompressionLevel level)
        {
            return CompressWithCodecStream(input, stream => new DeflateStream(stream, level, leaveOpen: true));
        }

        public byte[] Decompress(ReadOnlySpan<byte> input, int expectedLength, int maxDecompressedSizeBytes)
        {
            return DecompressWithCodecStream(input, stream => new DeflateStream(stream, CompressionMode.Decompress, leaveOpen: true), expectedLength, maxDecompressedSizeBytes);
        }
    }

    private static byte[] CompressWithCodecStream(ReadOnlySpan<byte> input, Func<Stream, Stream> streamFactory)
    {
        using var output = new MemoryStream(capacity: input.Length);
        using (var codecStream = streamFactory(output))
        {
            codecStream.Write(input);
            codecStream.Flush();
        }

        return output.ToArray();
    }

    private static byte[] DecompressWithCodecStream(
        ReadOnlySpan<byte> input,
        Func<Stream, Stream> streamFactory,
        int expectedLength,
        int maxDecompressedSizeBytes)
    {
        if (maxDecompressedSizeBytes <= 0)
            throw new ArgumentOutOfRangeException(nameof(maxDecompressedSizeBytes));

        using var compressed = new MemoryStream(input.ToArray(), writable: false);
        using var codecStream = streamFactory(compressed);
        using var output = expectedLength > 0
            ? new MemoryStream(capacity: expectedLength)
            : new MemoryStream();

        var buffer = ArrayPool<byte>.Shared.Rent(8192);
        try
        {
            int totalWritten = 0;
            while (true)
            {
                var bytesRead = codecStream.Read(buffer, 0, buffer.Length);
                if (bytesRead <= 0)
                    break;

                totalWritten += bytesRead;
                if (totalWritten > maxDecompressedSizeBytes)
                    throw new InvalidDataException($"Decompressed payload exceeds max allowed size ({maxDecompressedSizeBytes} bytes).");

                output.Write(buffer, 0, bytesRead);
            }

            if (expectedLength >= 0 && totalWritten != expectedLength)
                throw new InvalidDataException($"Expected decompressed length {expectedLength}, actual {totalWritten}.");

            return output.ToArray();
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(buffer);
        }
    }
}