natsdotnet/tests/NATS.Server.Tests/JetStream/Storage/FileStoreTombstoneTrackingTests.cs

// Reference: golang/nats-server/server/filestore.go:5267 (removeMsg)
//            golang/nats-server/server/filestore.go:5890 (eraseMsg)
//
// Tests verifying:
//   1. SequenceSet correctly tracks deleted sequences in MsgBlock
//   2. Tombstones survive MsgBlock recovery (RebuildIndex populates SequenceSet)
//   3. Secure erase (Delete with secureErase=true) overwrites payload bytes
//   4. EraseMsg at FileStore level marks the sequence as deleted
//
// Go test analogs:
//   TestFileStoreEraseMsgDoesNotLoseTombstones   (filestore_test.go:10781)
//   TestFileStoreTombstonesNoFirstSeqRollback     (filestore_test.go:10911)
//   TestFileStoreRemoveMsg                        (filestore_test.go:5267)

using System.Security.Cryptography;
using System.Text;
using NATS.Server.JetStream.Storage;

namespace NATS.Server.Tests.JetStream.Storage;

/// <summary>
/// Tests for SequenceSet-backed deletion tracking and secure erase in MsgBlock.
/// Reference: golang/nats-server/server/filestore.go eraseMsg / removeMsg.
/// </summary>
public sealed class FileStoreTombstoneTrackingTests : IDisposable
{
    private readonly string _testDir;

    public FileStoreTombstoneTrackingTests()
    {
        _testDir = Path.Combine(Path.GetTempPath(), $"nats-tombstone-tracking-{Guid.NewGuid():N}");
        Directory.CreateDirectory(_testDir);
    }

    public void Dispose()
    {
        if (Directory.Exists(_testDir))
            Directory.Delete(_testDir, recursive: true);
    }

    private string UniqueDir()
    {
        var dir = Path.Combine(_testDir, Guid.NewGuid().ToString("N"));
        Directory.CreateDirectory(dir);
        return dir;
    }

    // -------------------------------------------------------------------------
    // SequenceSet tracking in MsgBlock
    // -------------------------------------------------------------------------

    // Go: removeMsg — after Delete, IsDeleted returns true and DeletedCount == 1
    [Fact]
    public void MsgBlock_Delete_TracksDeletionInSequenceSet()
    {
        var dir = UniqueDir();
        using var block = MsgBlock.Create(0, dir, maxBytes: 1024 * 1024);

        block.Write("a", ReadOnlyMemory<byte>.Empty, "payload"u8.ToArray());
        block.Write("b", ReadOnlyMemory<byte>.Empty, "payload"u8.ToArray());
        block.Write("c", ReadOnlyMemory<byte>.Empty, "payload"u8.ToArray());

        block.Delete(2).ShouldBeTrue();

        block.IsDeleted(2).ShouldBeTrue();
        block.IsDeleted(1).ShouldBeFalse();
        block.IsDeleted(3).ShouldBeFalse();
        block.DeletedCount.ShouldBe(1UL);
        block.MessageCount.ShouldBe(2UL);
    }

    // Multiple deletes tracked correctly — SequenceSet merges contiguous ranges.
    [Fact]
    public void MsgBlock_MultipleDeletes_AllTrackedInSequenceSet()
    {
        var dir = UniqueDir();
        using var block = MsgBlock.Create(0, dir, maxBytes: 1024 * 1024);

        for (var i = 0; i < 10; i++)
            block.Write($"subj.{i}", ReadOnlyMemory<byte>.Empty, "payload"u8.ToArray());

        // Delete seqs 3, 4, 5 (contiguous — SequenceSet will merge into one range).
        block.Delete(3).ShouldBeTrue();
        block.Delete(4).ShouldBeTrue();
        block.Delete(5).ShouldBeTrue();

        block.DeletedCount.ShouldBe(3UL);
        block.MessageCount.ShouldBe(7UL);

        block.IsDeleted(3).ShouldBeTrue();
        block.IsDeleted(4).ShouldBeTrue();
        block.IsDeleted(5).ShouldBeTrue();
        block.IsDeleted(2).ShouldBeFalse();
        block.IsDeleted(6).ShouldBeFalse();
    }

    // -------------------------------------------------------------------------
    // Tombstones survive recovery (RebuildIndex populates SequenceSet)
    // -------------------------------------------------------------------------

    // Go: TestFileStoreTombstonesNoFirstSeqRollback — after restart, deleted seqs still deleted.
    // Reference: filestore.go RebuildIndex reads ebit from block file.
    [Fact]
    public void MsgBlock_Recovery_TombstonesInSequenceSet()
    {
        var dir = UniqueDir();

        // Phase 1: write messages and delete one, then close.
        using (var block = MsgBlock.Create(0, dir, maxBytes: 1024 * 1024))
        {
            block.Write("a", ReadOnlyMemory<byte>.Empty, "one"u8.ToArray());
            block.Write("b", ReadOnlyMemory<byte>.Empty, "two"u8.ToArray());
            block.Write("c", ReadOnlyMemory<byte>.Empty, "three"u8.ToArray());
            block.Delete(2); // marks seq 2 with ebit on disk
            block.Flush();
        }

        // Phase 2: recover from file — SequenceSet must be populated by RebuildIndex.
        using var recovered = MsgBlock.Recover(0, dir);
        recovered.DeletedCount.ShouldBe(1UL);
        recovered.MessageCount.ShouldBe(2UL);
        recovered.IsDeleted(1).ShouldBeFalse();
        recovered.IsDeleted(2).ShouldBeTrue();
        recovered.IsDeleted(3).ShouldBeFalse();

        // Read should return null for deleted seq.
        recovered.Read(2).ShouldBeNull();
        recovered.Read(1).ShouldNotBeNull();
        recovered.Read(3).ShouldNotBeNull();
    }

    // Multiple tombstones survive recovery.
    [Fact]
    public void MsgBlock_Recovery_MultipleDeletedSeqs_AllInSequenceSet()
    {
        var dir = UniqueDir();

        using (var block = MsgBlock.Create(0, dir, maxBytes: 1024 * 1024))
        {
            for (var i = 0; i < 10; i++)
                block.Write($"subj", ReadOnlyMemory<byte>.Empty, "payload"u8.ToArray());

            block.Delete(1);
            block.Delete(3);
            block.Delete(5);
            block.Delete(7);
            block.Delete(9);
            block.Flush();
        }

        using var recovered = MsgBlock.Recover(0, dir);
        recovered.DeletedCount.ShouldBe(5UL);
        recovered.MessageCount.ShouldBe(5UL);

        for (ulong seq = 1; seq <= 9; seq += 2)
            recovered.IsDeleted(seq).ShouldBeTrue($"seq {seq} should be deleted");
        for (ulong seq = 2; seq <= 10; seq += 2)
            recovered.IsDeleted(seq).ShouldBeFalse($"seq {seq} should NOT be deleted");
    }

    // Skip records (WriteSkip) survive recovery and appear in SequenceSet.
    [Fact]
    public void MsgBlock_Recovery_SkipRecordsInSequenceSet()
    {
        var dir = UniqueDir();

        using (var block = MsgBlock.Create(0, dir, maxBytes: 1024 * 1024, firstSequence: 1))
        {
            block.Write("a", ReadOnlyMemory<byte>.Empty, "payload"u8.ToArray()); // seq=1
            block.WriteSkip(2);  // tombstone
            block.WriteSkip(3);  // tombstone
            block.Write("b", ReadOnlyMemory<byte>.Empty, "payload"u8.ToArray()); // seq=4
            block.Flush();
        }

        using var recovered = MsgBlock.Recover(0, dir);
        // Seqs 2 and 3 are skip records → deleted.
        recovered.IsDeleted(2).ShouldBeTrue();
        recovered.IsDeleted(3).ShouldBeTrue();
        recovered.IsDeleted(1).ShouldBeFalse();
        recovered.IsDeleted(4).ShouldBeFalse();
        recovered.DeletedCount.ShouldBe(2UL);
        recovered.MessageCount.ShouldBe(2UL);
    }

    // -------------------------------------------------------------------------
    // Secure erase — payload bytes are overwritten with random data
    // -------------------------------------------------------------------------

    // Go: eraseMsg (filestore.go:5890) — payload bytes replaced with random bytes.
    [Fact]
    public void MsgBlock_SecureErase_OverwritesPayloadBytes()
    {
        var dir = UniqueDir();
        var original = Encoding.UTF8.GetBytes("this is a secret payload");

        using (var block = MsgBlock.Create(0, dir, maxBytes: 1024 * 1024))
        {
            block.Write("secret", ReadOnlyMemory<byte>.Empty, original);

            // Perform secure erase — overwrites payload bytes in-place on disk.
            block.Delete(1, secureErase: true).ShouldBeTrue();
            block.Flush();
        }

        // Read the raw block file and verify the original payload bytes are gone.
        var blockFile = Path.Combine(dir, "000000.blk");
        var rawBytes = File.ReadAllBytes(blockFile);

        // The payload "this is a secret payload" should no longer appear as a substring.
        var payloadBytes = Encoding.UTF8.GetBytes("this is a secret");
        var rawAsSpan = rawBytes.AsSpan();
        var found = false;
        for (var i = 0; i <= rawBytes.Length - payloadBytes.Length; i++)
        {
            if (rawAsSpan[i..].StartsWith(payloadBytes.AsSpan()))
            {
                found = true;
                break;
            }
        }
        found.ShouldBeFalse("Secret payload bytes should have been overwritten by secure erase");
    }

    // After secure erase, the message appears deleted (returns null on Read).
    [Fact]
    public void MsgBlock_SecureErase_MessageAppearsDeleted()
    {
        var dir = UniqueDir();
        using var block = MsgBlock.Create(0, dir, maxBytes: 1024 * 1024);

        block.Write("sensitive", ReadOnlyMemory<byte>.Empty, "secret data"u8.ToArray());
        block.Write("other", ReadOnlyMemory<byte>.Empty, "normal"u8.ToArray());

        block.Delete(1, secureErase: true).ShouldBeTrue();

        block.IsDeleted(1).ShouldBeTrue();
        block.Read(1).ShouldBeNull();
        block.Read(2).ShouldNotBeNull(); // other message unaffected
        block.DeletedCount.ShouldBe(1UL);
        block.MessageCount.ShouldBe(1UL);
    }

    // Secure erase with secureErase=false is identical to regular delete (no overwrite).
    [Fact]
    public void MsgBlock_Delete_WithSecureEraseFalse_NormalDelete()
    {
        var dir = UniqueDir();
        using var block = MsgBlock.Create(0, dir, maxBytes: 1024 * 1024);

        block.Write("x", ReadOnlyMemory<byte>.Empty, "content"u8.ToArray());
        block.Delete(1, secureErase: false).ShouldBeTrue();
        block.IsDeleted(1).ShouldBeTrue();
        block.Read(1).ShouldBeNull();
    }

    // Double secure erase returns false on second call.
    [Fact]
    public void MsgBlock_SecureErase_DoubleErase_ReturnsFalse()
    {
        var dir = UniqueDir();
        using var block = MsgBlock.Create(0, dir, maxBytes: 1024 * 1024);

        block.Write("x", ReadOnlyMemory<byte>.Empty, "content"u8.ToArray());
        block.Delete(1, secureErase: true).ShouldBeTrue();
        block.Delete(1, secureErase: true).ShouldBeFalse(); // already deleted
    }

    // -------------------------------------------------------------------------
    // DeletedSequences property returns snapshot of SequenceSet
    // -------------------------------------------------------------------------

    // DeletedSequences snapshot contains all deleted seqs (still IReadOnlySet from HashSet copy).
    [Fact]
    public void DeletedSequences_ReturnsCorrectSnapshot()
    {
        var dir = UniqueDir();
        using var block = MsgBlock.Create(0, dir, maxBytes: 1024 * 1024);

        block.Write("a", ReadOnlyMemory<byte>.Empty, "one"u8.ToArray());
        block.Write("b", ReadOnlyMemory<byte>.Empty, "two"u8.ToArray());
        block.Write("c", ReadOnlyMemory<byte>.Empty, "three"u8.ToArray());
        block.Write("d", ReadOnlyMemory<byte>.Empty, "four"u8.ToArray());

        block.Delete(2);
        block.Delete(4);

        var snapshot = block.DeletedSequences;
        snapshot.Count.ShouldBe(2);
        snapshot.ShouldContain(2UL);
        snapshot.ShouldContain(4UL);
        snapshot.ShouldNotContain(1UL);
        snapshot.ShouldNotContain(3UL);
    }

    // -------------------------------------------------------------------------
    // FileStore EraseMsg integration
    // -------------------------------------------------------------------------

    // Go: eraseMsg — after EraseMsg, message is gone and state reflects deletion.
    [Fact]
    public void FileStore_EraseMsg_MessageGoneAfterErase()
    {
        var dir = UniqueDir();
        var opts = new FileStoreOptions { Directory = dir };
        using var store = new FileStore(opts);

        store.StoreMsg("foo", null, "secret"u8.ToArray(), 0);
        store.StoreMsg("foo", null, "normal"u8.ToArray(), 0);

        var state1 = store.State();
        state1.Msgs.ShouldBe(2UL);

        store.EraseMsg(1).ShouldBeTrue();

        var state2 = store.State();
        state2.Msgs.ShouldBe(1UL);

        // Erasing same seq twice returns false.
        store.EraseMsg(1).ShouldBeFalse();
    }

    // Go: TestFileStoreEraseMsgDoesNotLoseTombstones — erase does not disturb other tombstones.
    // Reference: filestore_test.go:10781
    [Fact]
    public void FileStore_EraseMsg_DoesNotLoseTombstones()
    {
        var dir = UniqueDir();
        var opts = new FileStoreOptions { Directory = dir };
        using var store = new FileStore(opts);

        store.StoreMsg("foo", null, [], 0);                                                         // seq=1
        store.StoreMsg("foo", null, [], 0);                                                         // seq=2 (tombstone)
        store.StoreMsg("foo", null, "secret"u8.ToArray(), 0);                                      // seq=3 (erased)

        store.RemoveMsg(2); // tombstone seq=2
        store.StoreMsg("foo", null, [], 0); // seq=4

        store.EraseMsg(3); // erase seq=3

        var state = store.State();
        state.Msgs.ShouldBe(2UL);        // msgs 1 and 4 remain
        state.NumDeleted.ShouldBe(2);    // seqs 2 and 3 deleted
        state.Deleted.ShouldNotBeNull();
        state.Deleted!.ShouldContain(2UL);
        state.Deleted.ShouldContain(3UL);

        // Restart — state should be identical.
        store.Dispose();
        using var store2 = new FileStore(opts);
        var after = store2.State();
        after.Msgs.ShouldBe(2UL);
        after.NumDeleted.ShouldBe(2);
        after.Deleted.ShouldNotBeNull();
        after.Deleted!.ShouldContain(2UL);
        after.Deleted.ShouldContain(3UL);
    }
}