lmxopcua/tests/Server/ZB.MOM.WW.OtOpcUa.ScriptWorker.Tests/PeCacheTests.cs

using System.Diagnostics;
using System.Security.Cryptography;
using Shouldly;
using Xunit;
using ZB.MOM.WW.OtOpcUa.Core.Scripting;
using ZB.MOM.WW.OtOpcUa.Core.Scripting.Protocol;
using ZB.MOM.WW.OtOpcUa.Core.VirtualTags;

namespace ZB.MOM.WW.OtOpcUa.ScriptWorker.Tests;

/// <summary>
///     lmxopcua#780 design §7 (owner decision §14 #8): the node-local HMAC'd compiled-PE cache. The store is tested
///     directly (authentication, eviction, keying); the worker-backed cases prove a restart loads from it instead of
///     recompiling, and that the worker never runs bytes that did not verify.
/// </summary>
public sealed class PeCacheStoreTests
{
    private static readonly byte[] Fingerprint = SHA256.HashData("fp"u8);

    private static (ScriptPeCacheStore Store, Scripts.TempDir Dir) NewStore()
    {
        var dir = new Scripts.TempDir();
        var key = ScriptPeCacheStore.LoadOrCreateKey(dir.Path);
        return (new ScriptPeCacheStore(dir.Path, key), dir);
    }

    [Fact]
    public void A_written_entry_reads_back_and_an_absent_one_misses()
    {
        var (store, dir) = NewStore();
        using var cleanup = dir;
        var id = ScriptPeCacheStore.EntryId(Fingerprint, 1, "return 1;");
        store.TryRead(id, out _).ShouldBe(PeCacheLookup.Miss);
        store.TryWrite(id, [1, 2, 3, 4]).ShouldBeTrue();
        store.TryRead(id, out var pe).ShouldBe(PeCacheLookup.Hit);
        pe.ShouldBe(new byte[] { 1, 2, 3, 4 });
    }

    /// <summary>Every way an entry can be made to lie is rejected AND deleted. Positive control: the untouched entry hits.</summary>
    [Theory]
    [InlineData("flip-pe")]
    [InlineData("flip-tag")]
    [InlineData("truncate")]
    [InlineData("magic")]
    [InlineData("other-key")]
    [InlineData("renamed")]
    public void A_tampered_or_mismatched_entry_is_rejected_and_evicted(string tamper)
    {
        var (store, dir) = NewStore();
        using var cleanup = dir;
        var id = ScriptPeCacheStore.EntryId(Fingerprint, 1, "return 1;");
        var otherId = ScriptPeCacheStore.EntryId(Fingerprint, 1, "return 2;");
        store.TryWrite(id, [10, 20, 30, 40, 50]).ShouldBeTrue();
        store.TryRead(id, out _).ShouldBe(PeCacheLookup.Hit, "positive control");
        var path = store.EntryPath(id);
        var bytes = File.ReadAllBytes(path);
        var readId = id;
        switch (tamper)
        {
            case "flip-pe": bytes[^1] ^= 0xFF; File.WriteAllBytes(path, bytes); break;
            case "flip-tag": bytes[10] ^= 0x01; File.WriteAllBytes(path, bytes); break;
            case "truncate": File.WriteAllBytes(path, bytes[..(bytes.Length - 2)]); break;
            case "magic": bytes[0] = (byte)'X'; File.WriteAllBytes(path, bytes); break;
            case "other-key":
                new ScriptPeCacheStore(dir.Path, RandomNumberGenerator.GetBytes(32)).TryWrite(id, [10, 20, 30, 40, 50]).ShouldBeTrue();
                break;
            case "renamed":
                // A valid entry for one source copied under another source's id: the tag binds the id.
                File.Move(path, store.EntryPath(otherId));
                path = store.EntryPath(otherId);
                readId = otherId;
                break;
        }

        store.TryRead(readId, out var pe).ShouldBe(PeCacheLookup.Rejected);
        pe.ShouldBeNull();
        File.Exists(path).ShouldBeFalse("a rejected entry is evicted");
    }

    [Fact]
    public void The_entry_id_changes_with_the_compiler_fingerprint_the_kind_and_the_source()
    {
        var baseline = ScriptPeCacheStore.EntryId(Fingerprint, 1, "return 1;");
        ScriptPeCacheStore.EntryId(SHA256.HashData("fp2"u8), 1, "return 1;").ShouldNotBe(baseline, "a compiler/analyzer change misses");
        ScriptPeCacheStore.EntryId(Fingerprint, 2, "return 1;").ShouldNotBe(baseline);
        ScriptPeCacheStore.EntryId(Fingerprint, 1, "return 1; ").ShouldNotBe(baseline);
        ScriptPeCacheStore.EntryId(Fingerprint, 1, "return 1;").ShouldBe(baseline);
    }

    [Fact]
    public void The_compile_fingerprint_covers_the_context_and_result_types()
    {
        ScriptEvaluator<VirtualTagContext, object?>.CacheFingerprint.Length.ShouldBe(32);
        ScriptEvaluator<VirtualTagContext, object?>.CacheFingerprint
            .ShouldNotBe(ScriptEvaluator<ZB.MOM.WW.OtOpcUa.Core.ScriptedAlarms.AlarmPredicateContext, bool>.CacheFingerprint);
    }

    [Fact]
    public void The_key_is_created_once_private_and_reused()
    {
        using var dir = new Scripts.TempDir();
        var key = ScriptPeCacheStore.LoadOrCreateKey(dir.Path);
        key.Length.ShouldBe(ScriptPeCacheStore.KeyLength);
        ScriptPeCacheStore.LoadOrCreateKey(dir.Path).ShouldBe(key, "the second start reuses it");
        if (!OperatingSystem.IsWindows())
        {
            File.GetUnixFileMode(Path.Combine(dir.Path, ScriptPeCacheStore.KeyFileName))
                .ShouldBe(UnixFileMode.UserRead | UnixFileMode.UserWrite, "readable only by the host account");
        }
    }

    [Fact]
    public void A_malformed_key_is_refused_not_silently_replaced()
    {
        using var dir = new Scripts.TempDir();
        File.WriteAllBytes(Path.Combine(dir.Path, ScriptPeCacheStore.KeyFileName), [1, 2, 3]);
        Should.Throw<IOException>(() => ScriptPeCacheStore.LoadOrCreateKey(dir.Path));
        File.ReadAllBytes(Path.Combine(dir.Path, ScriptPeCacheStore.KeyFileName)).Length.ShouldBe(3);
    }

    [Fact]
    public void Prune_drops_stale_entries_then_the_least_recently_used()
    {
        var (store, dir) = NewStore();
        using var cleanup = dir;
        var ids = Enumerable.Range(0, 5).Select(i => ScriptPeCacheStore.EntryId(Fingerprint, 1, "s" + i)).ToArray();
        for (var i = 0; i < ids.Length; i++)
        {
            store.TryWrite(ids[i], [1]).ShouldBeTrue();
            File.SetLastWriteTimeUtc(store.EntryPath(ids[i]), DateTime.UtcNow.AddDays(-i * 10));
        }
        store.Prune(maxAge: TimeSpan.FromDays(25), maxEntries: 2).ShouldBe(3);
        File.Exists(store.EntryPath(ids[0])).ShouldBeTrue();
        File.Exists(store.EntryPath(ids[1])).ShouldBeTrue();
        File.Exists(store.EntryPath(ids[2])).ShouldBeFalse("LRU beyond the cap");
        File.Exists(store.EntryPath(ids[3])).ShouldBeFalse("older than the max age");
    }
}

/// <summary>The PE cache through real workers.</summary>
[Collection(ScriptWorkerCollection.Name)]
[Trait("Category", "ScriptWorker")]
public sealed class PeCacheWorkerTests
{
    private static CancellationToken Ct => TestContext.Current.CancellationToken;

    private static string Source(int i) => $"var x = (int)ctx.GetTag(\"x\").Value; var acc = 0L; for (var j = 0; j < 10; j++) acc += x * {i} + j; return acc;";

    private static long Expected(int i) => Enumerable.Range(0, 10).Sum(j => 3L * i + j);

    /// <summary>Time from a fresh runner (fresh worker process) to the first result of every one of <paramref name="count"/> scripts.</summary>
    private static async Task<(TimeSpan Elapsed, (long Hits, long Misses, long Rejected) Counts)> ColdStart(string? cacheDir, int count)
    {
        await using var runner = Scripts.Runner(Scripts.Options(peCacheDir: cacheDir));
        var sw = Stopwatch.StartNew();
        for (var i = 0; i < count; i++)
        {
            var outcome = await runner.RunAsync(Scripts.Vt("vt-" + i, Source(i), inputs: Scripts.Input("x", 3)), Ct);
            outcome.Status.ShouldBe(ScriptRunStatus.Ok, outcome.Reason);
            outcome.Value.ShouldBe(Expected(i));
        }
        return (sw.Elapsed, runner.PeCacheCounts);
    }

    /// <summary>
    ///     The measurement the slice asks for: worker restart to first result for N scripts — with no cache (every
    ///     start recompiles), with an empty cache (compiles and writes) and with a warm cache (loads only). The warm
    ///     start must serve every script from the cache and beat the cold one. Mutation check: a worker that ignores the
    ///     cache reports misses on the warm start.
    /// </summary>
    [Fact]
    public async Task A_warm_restart_loads_every_script_from_the_cache_and_is_faster_than_a_cold_one()
    {
        const int count = 40;
        using var dir = new Scripts.TempDir();
        var cacheDir = Path.Combine(dir.Path, "pe-cache");

        var uncached = await ColdStart(null, count);
        var cold = await ColdStart(cacheDir, count);
        var warm = await ColdStart(cacheDir, count);

        TestContext.Current.TestOutputHelper?.WriteLine(
            $"{count} scripts, restart to all first results: no cache {uncached.Elapsed.TotalMilliseconds:F0} ms; " +
            $"empty cache {cold.Elapsed.TotalMilliseconds:F0} ms; warm cache {warm.Elapsed.TotalMilliseconds:F0} ms");

        uncached.Counts.ShouldBe((0L, 0L, 0L), "no cache configured");
        cold.Counts.ShouldBe((0L, (long)count, 0L), "the first start compiles and writes every entry");
        warm.Counts.ShouldBe(((long)count, 0L, 0L), "the restart loads every script from the cache");
        warm.Elapsed.ShouldBeLessThan(cold.Elapsed, "loading a PE is cheaper than entering Roslyn");
        Directory.EnumerateFiles(Path.Combine(cacheDir, ScriptPeCacheStore.EntriesDirectoryName), "*.pe").Count().ShouldBe(count);
    }

    /// <summary>
    ///     The security property: a cache entry replaced with a DIFFERENT, validly framed PE (one that returns 666) under
    ///     a forged tag is never run — the worker rejects it, deletes it, recompiles the real source and returns 1.
    ///     Positive control: the forged image itself runs and returns 666 when loaded directly.
    /// </summary>
    [Fact]
    public async Task A_forged_cache_entry_is_rejected_and_the_real_source_recompiled()
    {
        using var dir = new Scripts.TempDir();
        var cacheDir = Path.Combine(dir.Path, "pe-cache");
        await using (var first = Scripts.Runner(Scripts.Options(peCacheDir: cacheDir)))
            (await first.RunAsync(Scripts.Vt("vt", "return 1;"), Ct)).Value.ShouldBe(1);
        var entry = Directory.EnumerateFiles(Path.Combine(cacheDir, ScriptPeCacheStore.EntriesDirectoryName), "*.pe").ShouldHaveSingleItem();

        var evil = ScriptEvaluator<VirtualTagContext, object?>.CompileToImage("return 666;");
        using (var direct = ScriptEvaluator<VirtualTagContext, object?>.FromImage(evil))
        {
            (await direct.RunAsync(new VirtualTagContext(Scripts.NoInputs, static (_, _) => { }, Serilog.Core.Logger.None, () => DateTime.UtcNow), Ct))
                .ShouldBe(666, "positive control: the forged image is a working script");
        }
        var original = File.ReadAllBytes(entry);
        var forged = new byte[8 + 32 + evil.Length];
        original.AsSpan(0, 8).CopyTo(forged);                    // the real magic
        RandomNumberGenerator.GetBytes(32).CopyTo(forged, 8);    // a tag the attacker cannot compute
        evil.CopyTo(forged, 40);
        File.WriteAllBytes(entry, forged);

        await using var second = Scripts.Runner(Scripts.Options(peCacheDir: cacheDir));
        var outcome = await second.RunAsync(Scripts.Vt("vt", "return 1;"), Ct);
        outcome.Value.ShouldBe(1, "the forged bytes never ran");
        second.PeCacheCounts.Rejected.ShouldBe(1);
        var key = File.ReadAllBytes(Path.Combine(cacheDir, ScriptPeCacheStore.KeyFileName));
        new ScriptPeCacheStore(cacheDir, key).TryRead(Path.GetFileNameWithoutExtension(entry), out _)
            .ShouldBe(PeCacheLookup.Hit, "the entry was rewritten from the real compile");
    }

    /// <summary>A compile error is never cached (the host's negative cache covers it) and does not poison the directory.</summary>
    [Fact]
    public async Task A_compile_error_writes_no_entry()
    {
        using var dir = new Scripts.TempDir();
        var cacheDir = Path.Combine(dir.Path, "pe-cache");
        await using var runner = Scripts.Runner(Scripts.Options(peCacheDir: cacheDir));
        (await runner.RunAsync(Scripts.Vt("bad", "return nope;"), Ct)).Reason!.ShouldStartWith("compile error:");
        Directory.EnumerateFiles(Path.Combine(cacheDir, ScriptPeCacheStore.EntriesDirectoryName), "*.pe").ShouldBeEmpty();
    }

    /// <summary>An unusable cache directory is a Warning and a cache-less runner, never a failure.</summary>
    [Fact]
    public async Task An_unusable_cache_directory_degrades_to_no_cache()
    {
        using var dir = new Scripts.TempDir();
        var blocker = Path.Combine(dir.Path, "a-file");
        File.WriteAllText(blocker, "not a directory");
        var sink = new CapturingSink();
        await using var runner = Scripts.Runner(Scripts.Options(peCacheDir: Path.Combine(blocker, "pe-cache")), sink.CreateLogger());
        (await runner.RunAsync(Scripts.Vt("vt", "return 5;"), Ct)).Value.ShouldBe(5);
        runner.Status().PeCacheEnabled.ShouldBeFalse();
        sink.Rendered.ShouldContain(m => m.Contains("is unusable"));
    }
}