CBDD/tests/CBDD.Tests/CompressionCorruptionTests.cs

using System.Buffers.Binary;
using System.IO.Compression;
using ZB.MOM.WW.CBDD.Bson;
using ZB.MOM.WW.CBDD.Core.Compression;
using ZB.MOM.WW.CBDD.Core.Storage;
using ZB.MOM.WW.CBDD.Shared;

namespace ZB.MOM.WW.CBDD.Tests;

public class CompressionCorruptionTests
{
    /// <summary>
    /// Verifies corrupted compressed payload checksum triggers invalid data errors.
    /// </summary>
    [Fact]
    public void Read_WithBadChecksum_ShouldThrowInvalidData()
    {
        var dbPath = NewDbPath();
        var options = CompressionEnabledOptions();

        try
        {
            using var db = new TestDbContext(dbPath, options);
            var id = InsertCheckpointAndCorrupt(db, header =>
            {
                var currentChecksum = BinaryPrimitives.ReadUInt32LittleEndian(header.Slice(12, 4));
                BinaryPrimitives.WriteUInt32LittleEndian(header.Slice(12, 4), currentChecksum + 1);
            });

            var ex = Should.Throw<InvalidDataException>(() => db.Users.FindById(id));
            ex.Message.ShouldContain("checksum mismatch");
            db.GetCompressionStats().ChecksumFailureCount.ShouldBeGreaterThanOrEqualTo(1);
        }
        finally
        {
            CleanupFiles(dbPath);
        }
    }

    /// <summary>
    /// Verifies invalid original length metadata triggers invalid data errors.
    /// </summary>
    [Fact]
    public void Read_WithBadOriginalLength_ShouldThrowInvalidData()
    {
        var dbPath = NewDbPath();
        var options = CompressionEnabledOptions();

        try
        {
            using var db = new TestDbContext(dbPath, options);
            var id = InsertCheckpointAndCorrupt(db, header =>
            {
                BinaryPrimitives.WriteInt32LittleEndian(header.Slice(4, 4), -1);
            });

            var ex = Should.Throw<InvalidDataException>(() => db.Users.FindById(id));
            ex.Message.ShouldContain("decompress");
        }
        finally
        {
            CleanupFiles(dbPath);
        }
    }

    /// <summary>
    /// Verifies oversized declared decompressed length enforces safety guardrails.
    /// </summary>
    [Fact]
    public void Read_WithOversizedDeclaredLength_ShouldEnforceGuardrail()
    {
        var dbPath = NewDbPath();
        var options = CompressionEnabledOptions(maxDecompressedSizeBytes: 2048);

        try
        {
            using var db = new TestDbContext(dbPath, options);
            var id = InsertCheckpointAndCorrupt(db, header =>
            {
                BinaryPrimitives.WriteInt32LittleEndian(header.Slice(4, 4), 2049);
            });

            var ex = Should.Throw<InvalidDataException>(() => db.Users.FindById(id));
            ex.Message.ShouldContain("invalid decompressed length");
            db.GetCompressionStats().SafetyLimitRejectionCount.ShouldBeGreaterThanOrEqualTo(1);
        }
        finally
        {
            CleanupFiles(dbPath);
        }
    }

    /// <summary>
    /// Verifies invalid codec identifiers in compressed headers trigger invalid data errors.
    /// </summary>
    [Fact]
    public void Read_WithInvalidCodecId_ShouldThrowInvalidData()
    {
        var dbPath = NewDbPath();
        var options = CompressionEnabledOptions();

        try
        {
            using var db = new TestDbContext(dbPath, options);
            var id = InsertCheckpointAndCorrupt(db, header =>
            {
                header[0] = 0; // CompressionCodec.None is invalid for compressed payload header.
            });

            var ex = Should.Throw<InvalidDataException>(() => db.Users.FindById(id));
            ex.Message.ShouldContain("invalid codec");
        }
        finally
        {
            CleanupFiles(dbPath);
        }
    }

    private static ObjectId InsertCheckpointAndCorrupt(TestDbContext db, HeaderMutator mutateHeader)
    {
        var user = new User
        {
            Name = BuildPayload(16_000),
            Age = 33
        };

        var id = db.Users.Insert(user);
        db.SaveChanges();
        db.ForceCheckpoint();

        var (pageId, slot, _) = FindFirstCompressedSlot(db.Storage);
        ((slot.Flags & SlotFlags.HasOverflow) != 0).ShouldBeFalse();

        var page = new byte[db.Storage.PageSize];
        db.Storage.ReadPage(pageId, null, page);
        var headerSlice = page.AsSpan(slot.Offset, CompressedPayloadHeader.Size);
        mutateHeader(headerSlice);
        db.Storage.WritePageImmediate(pageId, page);

        return id;
    }

    private static (uint PageId, SlotEntry Slot, ushort SlotIndex) FindFirstCompressedSlot(StorageEngine storage)
    {
        var buffer = new byte[storage.PageSize];
        for (uint pageId = 1; pageId < storage.PageCount; pageId++)
        {
            storage.ReadPage(pageId, null, buffer);
            var header = SlottedPageHeader.ReadFrom(buffer);
            if (header.PageType != PageType.Data)
                continue;

            for (ushort slotIndex = 0; slotIndex < header.SlotCount; slotIndex++)
            {
                var slotOffset = SlottedPageHeader.Size + (slotIndex * SlotEntry.Size);
                var slot = SlotEntry.ReadFrom(buffer.AsSpan(slotOffset, SlotEntry.Size));
                if ((slot.Flags & SlotFlags.Deleted) != 0)
                    continue;

                if ((slot.Flags & SlotFlags.Compressed) != 0)
                    return (pageId, slot, slotIndex);
            }
        }

        throw new InvalidOperationException("No active compressed slot found for corruption test setup.");
    }

    private static CompressionOptions CompressionEnabledOptions(int maxDecompressedSizeBytes = 32 * 1024)
    {
        return new CompressionOptions
        {
            EnableCompression = true,
            MinSizeBytes = 0,
            MinSavingsPercent = 0,
            Codec = CompressionCodec.Brotli,
            Level = CompressionLevel.Fastest,
            MaxDecompressedSizeBytes = maxDecompressedSizeBytes
        };
    }

    private delegate void HeaderMutator(Span<byte> header);

    private static string BuildPayload(int approxLength)
    {
        var builder = new System.Text.StringBuilder(approxLength + 256);
        var i = 0;
        while (builder.Length < approxLength)
        {
            builder.Append("corruption-payload-");
            builder.Append(i.ToString("D8"));
            builder.Append('|');
            i++;
        }

        return builder.ToString();
    }

    private static string NewDbPath()
        => Path.Combine(Path.GetTempPath(), $"compression_corruption_{Guid.NewGuid():N}.db");

    private static void CleanupFiles(string dbPath)
    {
        var walPath = Path.ChangeExtension(dbPath, ".wal");
        var markerPath = $"{dbPath}.compact.state";
        if (File.Exists(dbPath)) File.Delete(dbPath);
        if (File.Exists(walPath)) File.Delete(walPath);
        if (File.Exists(markerPath)) File.Delete(markerPath);
    }
}