whitemagic/WhiteMagicTest/Discovery/PatternScannerCacheTests.cs

using System.Linq;
using System.Reflection;
using System.Runtime.InteropServices;
using WhiteMagic;
using WhiteMagic.Discovery;

namespace WhiteMagicTest.Discovery;

/// <summary>
/// Tests for <see cref="PatternScannerCache"/>.
/// </summary>
public class PatternScannerCacheTests
{
    private static InProcessReader CreateReader()
    {
        return new InProcessReader();
    }

    [Fact]
    public void FindCached_returns_same_result_on_second_call()
    {
        using var reader = CreateReader();
        var cache = new PatternScannerCache(reader);

        // Create a buffer with a known pattern
        byte[] buffer = new byte[256];
        buffer[30] = 0x11;
        buffer[31] = 0x22;
        buffer[32] = 0x33;
        buffer[33] = 0x44;

        GCHandle pin = GCHandle.Alloc(buffer, GCHandleType.Pinned);
        try
        {
            IntPtr addr = pin.AddrOfPinnedObject();
            IntPtr end = addr + buffer.Length;

            byte[] pattern = { 0x11, 0x22, 0x33, 0x44 };

            // First call should scan memory
            IntPtr first = cache.FindCached(pattern, null, addr, end);

            // Second call should return cached result
            IntPtr second = cache.FindCached(pattern, null, addr, end);

            Assert.Equal(addr + 30, first);
            Assert.Equal(first, second);

            // Value equality must mean the second call reused the cached entry.
            var cacheField = typeof(PatternScannerCache).GetField("_cache", BindingFlags.NonPublic | BindingFlags.Instance)!;
            var cacheDict = cacheField.GetValue(cache)!;
            int count = (int)cacheDict.GetType().GetProperty("Count")!.GetValue(cacheDict)!;
            Assert.Equal(1, count);
        }
        finally
        {
            pin.Free();
        }
    }

    [Fact]
    public void FindCached_value_equality_uses_content_not_reference()
    {
        using var reader = CreateReader();
        var cache = new PatternScannerCache(reader);

        byte[] buffer = new byte[256];
        buffer[10] = 0xAA;
        buffer[11] = 0xBB;

        GCHandle pin = GCHandle.Alloc(buffer, GCHandleType.Pinned);
        try
        {
            IntPtr addr = pin.AddrOfPinnedObject();
            IntPtr end = addr + buffer.Length;

            byte[] pattern1 = { 0xAA, 0xBB };
            byte[] pattern2 = { 0xAA, 0xBB };

            IntPtr first = cache.FindCached(pattern1, null, addr, end);
            IntPtr second = cache.FindCached(pattern2, null, addr, end);

            Assert.Equal(addr + 10, first);
            Assert.Equal(first, second);

            var cacheField = typeof(PatternScannerCache).GetField("_cache", BindingFlags.NonPublic | BindingFlags.Instance)!;
            var cacheDict = cacheField.GetValue(cache)!;
            int count = (int)cacheDict.GetType().GetProperty("Count")!.GetValue(cacheDict)!;
            Assert.Equal(1, count);
        }
        finally
        {
            pin.Free();
        }
    }

    [Fact]
    public void FindCached_different_ranges_are_cached_separately()
    {
        using var reader = CreateReader();
        var cache = new PatternScannerCache(reader);

        // Create two separate buffers
        byte[] buffer1 = new byte[128];
        buffer1[10] = 0xAA;
        buffer1[11] = 0xBB;

        byte[] buffer2 = new byte[128];
        buffer2[20] = 0xAA;
        buffer2[21] = 0xBB;

        GCHandle pin1 = GCHandle.Alloc(buffer1, GCHandleType.Pinned);
        GCHandle pin2 = GCHandle.Alloc(buffer2, GCHandleType.Pinned);
        try
        {
            IntPtr addr1 = pin1.AddrOfPinnedObject();
            IntPtr end1 = addr1 + buffer1.Length;

            IntPtr addr2 = pin2.AddrOfPinnedObject();
            IntPtr end2 = addr2 + buffer2.Length;

            byte[] pattern = { 0xAA, 0xBB };

            IntPtr found1 = cache.FindCached(pattern, null, addr1, end1);
            IntPtr found2 = cache.FindCached(pattern, null, addr2, end2);

            Assert.Equal(addr1 + 10, found1);
            Assert.Equal(addr2 + 20, found2);
            Assert.NotEqual(found1, found2);
        }
        finally
        {
            pin1.Free();
            pin2.Free();
        }
    }

    [Fact]
    public void FindCached_with_mask_caches_correctly()
    {
        using var reader = CreateReader();
        var cache = new PatternScannerCache(reader);

        byte[] buffer = new byte[256];
        buffer[40] = 0x99;
        buffer[41] = 0x88;  // This is wildcard
        buffer[42] = 0x77;

        GCHandle pin = GCHandle.Alloc(buffer, GCHandleType.Pinned);
        try
        {
            IntPtr addr = pin.AddrOfPinnedObject();
            IntPtr end = addr + buffer.Length;

            byte[] pattern = { 0x99, 0x00, 0x77 };
            string mask = "x?x";

            IntPtr first = cache.FindCached(pattern, mask, addr, end);
            IntPtr second = cache.FindCached(pattern, mask, addr, end);

            Assert.Equal(addr + 40, first);
            Assert.Equal(first, second);
        }
        finally
        {
            pin.Free();
        }
    }

    [Fact]
    public void Clear_clears_cached_results()
    {
        using var reader = CreateReader();
        var cache = new PatternScannerCache(reader);

        byte[] buffer = new byte[256];
        buffer[50] = 0xCC;
        buffer[51] = 0xDD;

        GCHandle pin = GCHandle.Alloc(buffer, GCHandleType.Pinned);
        try
        {
            IntPtr addr = pin.AddrOfPinnedObject();
            IntPtr end = addr + buffer.Length;

            byte[] pattern = { 0xCC, 0xDD };

            // Cache a result
            IntPtr first = cache.FindCached(pattern, null, addr, end);
            Assert.Equal(addr + 50, first);

            // Clear the cache
            cache.Clear();

            // This should rescan (not return cached result)
            IntPtr second = cache.FindCached(pattern, null, addr, end);
            Assert.Equal(addr + 50, second);
        }
        finally
        {
            pin.Free();
        }
    }

    [Fact]
    public void FindInModuleCached_caches_module_scans()
    {
        using var reader = CreateReader();
        var cache = new PatternScannerCache(reader);

        var currentProcess = System.Diagnostics.Process.GetCurrentProcess();
        var mainModule = currentProcess.MainModule;
        Assert.NotNull(mainModule);

        // MZ header is always at the start of the main module
        byte[] pattern = { 0x4D, 0x5A };

        IntPtr first = cache.FindInModuleCached(pattern, null, mainModule);
        IntPtr second = cache.FindInModuleCached(pattern, null, mainModule);

        Assert.Equal(mainModule.BaseAddress, first);
        Assert.Equal(first, second);
    }

    [Fact]
    public void FindInModulesCached_caches_multiple_modules()
    {
        using var reader = CreateReader();
        var cache = new PatternScannerCache(reader);

        var currentProcess = System.Diagnostics.Process.GetCurrentProcess();
        var modules = currentProcess.Modules.Cast<System.Diagnostics.ProcessModule>().ToList();

        Assert.NotEmpty(modules);

        // MZ header should be present in at least one module
        byte[] pattern = { 0x4D, 0x5A };

        IntPtr first = cache.FindInModulesCached(pattern, null, modules);
        IntPtr second = cache.FindInModulesCached(pattern, null, modules);

        Assert.NotEqual(IntPtr.Zero, first);
        Assert.Equal(first, second);
    }
}