lmxopcua/src/Core/ZB.MOM.WW.OtOpcUa.Core.Scripting.Protocol/ScriptPeCacheStore.cs

using System.Security.Cryptography;
using System.Text;

namespace ZB.MOM.WW.OtOpcUa.Core.Scripting.Protocol;

/// <summary>What a <see cref="ScriptPeCacheStore.TryRead"/> found.</summary>
public enum PeCacheLookup
{
    /// <summary>No entry for the id.</summary>
    Miss = 0,

    /// <summary>An entry whose HMAC verified; its PE bytes are returned.</summary>
    Hit = 1,

    /// <summary>
    ///     An entry that exists but did not verify — truncated, edited, written under another key, or renamed from
    ///     another id. It has been deleted; the caller recompiles.
    /// </summary>
    Rejected = 2,
}

/// <summary>
///     The node-local compiled-PE cache (lmxopcua#780 design §7, owner decision §14 #8): content-addressed files under
///     <c>Scripting:Worker:PeCacheDir</c>, each authenticated with HMAC-SHA256 under a per-node random key, so a worker
///     restart loads its scripts' PE images instead of re-entering Roslyn — and never loads bytes it did not write.
/// </summary>
/// <remarks>
///     <list type="bullet">
///         <item><b>Layout.</b> <c>&lt;dir&gt;/pe-cache.key</c> (32 random bytes, created once by the host) and
///         <c>&lt;dir&gt;/pe/&lt;id&gt;.pe</c>, one file per entry: an 8-byte magic, the 32-byte tag, then the PE.</item>
///         <item><b>Identity.</b> <see cref="EntryId"/> is SHA-256 over the compile fingerprint (Roslyn version, the
///         sandbox and analyzer build, the reference set, the runtime), the context kind and the source. A compiler or
///         analyzer change therefore misses every old entry; it can never be served an image a newer sandbox rule
///         would have refused.</item>
///         <item><b>Authentication.</b> The tag covers the magic, the id and the PE, so an entry copied to another id
///         is rejected as well as an edited or truncated one. Loading unauthenticated PE bytes would be code execution
///         outside the analyzer; a rejected entry is deleted and the source recompiled.</item>
///         <item><b>Key.</b> Readable only by the host account: on Linux/macOS the directory is 0700 and the key file
///         0600. On Windows the file inherits the directory's ACL, which must be the installer's (lmxopcua#780 installer
///         work). The worker receives the key in <c>Hello</c>, never on disk or in its environment.</item>
///         <item><b>Writes</b> go to a unique temp file and are renamed into place, so a concurrent reader (another
///         shard) sees the old entry, the new one, or none — never a torn file.</item>
///     </list>
/// </remarks>
public sealed class ScriptPeCacheStore
{
    /// <summary>The key file's name under the cache directory.</summary>
    public const string KeyFileName = "pe-cache.key";

    /// <summary>The entries' subdirectory.</summary>
    public const string EntriesDirectoryName = "pe";

    /// <summary>The HMAC key length.</summary>
    public const int KeyLength = 32;

    /// <summary>An entry above this size is not read (a 4 MiB source compiles to far less).</summary>
    public const int MaxEntryBytes = 32 * 1024 * 1024;

    private const string EntryExtension = ".pe";
    private static readonly byte[] Magic = "OTPEC\u0001\0\0"u8.ToArray();
    private const int TagLength = 32;

    private readonly byte[] _key;

    /// <summary>Opens the store over <paramref name="directory"/> with <paramref name="key"/>.</summary>
    /// <param name="directory">The cache directory (absolute).</param>
    /// <param name="key">The node's HMAC key (<see cref="KeyLength"/> bytes).</param>
    public ScriptPeCacheStore(string directory, byte[] key)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(directory);
        ArgumentNullException.ThrowIfNull(key);
        if (key.Length != KeyLength) throw new ArgumentException($"the PE cache key must be {KeyLength} bytes", nameof(key));
        Directory = directory;
        EntriesDirectory = Path.Combine(directory, EntriesDirectoryName);
        _key = (byte[])key.Clone();
    }

    /// <summary>The cache directory.</summary>
    public string Directory { get; }

    /// <summary>The entries' directory.</summary>
    public string EntriesDirectory { get; }

    /// <summary>
    ///     The host half: returns the node's key, creating the directory and the key on first use. An existing key of
    ///     the wrong length is not replaced (it may be a sign of tampering); the call throws and the caller runs without
    ///     a cache.
    /// </summary>
    /// <param name="directory">The cache directory (absolute).</param>
    /// <returns>The key.</returns>
    /// <exception cref="IOException">The directory or key cannot be created or read, or the key is malformed.</exception>
    public static byte[] LoadOrCreateKey(string directory)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(directory);
        CreatePrivateDirectory(directory);
        CreatePrivateDirectory(Path.Combine(directory, EntriesDirectoryName));
        var keyPath = Path.Combine(directory, KeyFileName);
        if (!File.Exists(keyPath))
        {
            // Written to a temp name and moved without overwrite: two hosts racing on first start agree on one key.
            var temp = keyPath + "." + Convert.ToHexString(RandomNumberGenerator.GetBytes(8)) + ".tmp";
            var options = new FileStreamOptions { Mode = FileMode.CreateNew, Access = FileAccess.Write, Share = FileShare.None };
            if (!OperatingSystem.IsWindows()) options.UnixCreateMode = UnixFileMode.UserRead | UnixFileMode.UserWrite;
            using (var stream = new FileStream(temp, options))
            {
                stream.Write(RandomNumberGenerator.GetBytes(KeyLength));
            }
            try
            {
                File.Move(temp, keyPath, overwrite: false);
            }
            catch (IOException) when (File.Exists(keyPath))
            {
                File.Delete(temp);
            }
        }
        var key = File.ReadAllBytes(keyPath);
        if (key.Length != KeyLength)
            throw new IOException($"the PE cache key '{keyPath}' is {key.Length} bytes, not {KeyLength}; delete it to start a fresh cache");
        return key;
    }

    /// <summary>The entry id for one compile: hex SHA-256 over the fingerprint, the kind and the source.</summary>
    /// <param name="fingerprint">The compile fingerprint (<c>ScriptEvaluator.CacheFingerprint</c>).</param>
    /// <param name="kind">The context kind (the wire enum value).</param>
    /// <param name="source">The script source.</param>
    /// <returns>A 64-character lowercase hex id.</returns>
    public static string EntryId(ReadOnlySpan<byte> fingerprint, int kind, string source)
    {
        ArgumentNullException.ThrowIfNull(source);
        using var hash = IncrementalHash.CreateHash(HashAlgorithmName.SHA256);
        hash.AppendData(fingerprint);
        hash.AppendData(BitConverter.GetBytes(kind));
        hash.AppendData(Encoding.UTF8.GetBytes(source));
        return Convert.ToHexString(hash.GetHashAndReset()).ToLowerInvariant();
    }

    /// <summary>The path of an entry.</summary>
    /// <param name="entryId">The id.</param>
    /// <returns>The file path.</returns>
    public string EntryPath(string entryId)
    {
        if (!IsValidId(entryId)) throw new ArgumentException("not a PE cache entry id", nameof(entryId));
        return Path.Combine(EntriesDirectory, entryId + EntryExtension);
    }

    /// <summary>Reads and authenticates an entry. A rejected entry is deleted.</summary>
    /// <param name="entryId">The id.</param>
    /// <param name="peImage">The PE bytes on a hit.</param>
    /// <returns>Hit, miss, or rejected.</returns>
    public PeCacheLookup TryRead(string entryId, out byte[]? peImage)
    {
        peImage = null;
        var path = EntryPath(entryId);
        byte[] bytes;
        try
        {
            var info = new FileInfo(path);
            if (!info.Exists) return PeCacheLookup.Miss;
            if (info.Length > MaxEntryBytes)
            {
                TryDelete(path);
                return PeCacheLookup.Rejected;
            }
            bytes = File.ReadAllBytes(path);
        }
        catch (FileNotFoundException)
        {
            return PeCacheLookup.Miss;
        }
        catch (DirectoryNotFoundException)
        {
            return PeCacheLookup.Miss;
        }

        if (bytes.Length <= Magic.Length + TagLength || !bytes.AsSpan(0, Magic.Length).SequenceEqual(Magic))
        {
            TryDelete(path);
            return PeCacheLookup.Rejected;
        }
        var tag = bytes.AsSpan(Magic.Length, TagLength);
        var pe = bytes.AsSpan(Magic.Length + TagLength);
        if (!CryptographicOperations.FixedTimeEquals(tag, Tag(entryId, pe)))
        {
            TryDelete(path);
            return PeCacheLookup.Rejected;
        }

        peImage = pe.ToArray();
        // Recency for Prune: a hit counts as a use.
        try { File.SetLastWriteTimeUtc(path, DateTime.UtcNow); } catch (IOException) { } catch (UnauthorizedAccessException) { }
        return PeCacheLookup.Hit;
    }

    /// <summary>Writes an entry (atomically replacing any existing one).</summary>
    /// <param name="entryId">The id.</param>
    /// <param name="peImage">The PE bytes.</param>
    /// <returns>False when the write failed (the cache is best effort: a failed write costs a recompile later).</returns>
    public bool TryWrite(string entryId, ReadOnlySpan<byte> peImage)
    {
        var path = EntryPath(entryId);
        var temp = Path.Combine(EntriesDirectory, $".{entryId}.{Convert.ToHexString(RandomNumberGenerator.GetBytes(6))}.tmp");
        try
        {
            System.IO.Directory.CreateDirectory(EntriesDirectory);
            var options = new FileStreamOptions { Mode = FileMode.CreateNew, Access = FileAccess.Write, Share = FileShare.None };
            if (!OperatingSystem.IsWindows()) options.UnixCreateMode = UnixFileMode.UserRead | UnixFileMode.UserWrite;
            using (var stream = new FileStream(temp, options))
            {
                stream.Write(Magic);
                stream.Write(Tag(entryId, peImage));
                stream.Write(peImage);
            }
            File.Move(temp, path, overwrite: true);
            return true;
        }
        catch (Exception ex) when (ex is IOException or UnauthorizedAccessException)
        {
            TryDelete(temp);
            return false;
        }
    }

    /// <summary>
    ///     Deletes entries not used for <paramref name="maxAge"/>, then the least recently used beyond
    ///     <paramref name="maxEntries"/>, plus stray temp files. The design prunes "entries not in the current
    ///     generation"; this runner has no generation object, so recency stands in for it.
    /// </summary>
    /// <param name="maxAge">Entries older than this go.</param>
    /// <param name="maxEntries">At most this many entries stay.</param>
    /// <returns>How many files were deleted.</returns>
    public int Prune(TimeSpan maxAge, int maxEntries)
    {
        if (!System.IO.Directory.Exists(EntriesDirectory)) return 0;
        var deleted = 0;
        var now = DateTime.UtcNow;
        var entries = new List<FileInfo>();
        foreach (var file in new DirectoryInfo(EntriesDirectory).EnumerateFiles())
        {
            var stale = now - file.LastWriteTimeUtc;
            if (file.Name.EndsWith(".tmp", StringComparison.Ordinal))
            {
                if (stale > TimeSpan.FromHours(1) && TryDelete(file.FullName)) deleted++;
                continue;
            }
            if (!file.Name.EndsWith(EntryExtension, StringComparison.Ordinal)) continue;
            if (stale > maxAge)
            {
                if (TryDelete(file.FullName)) deleted++;
                continue;
            }
            entries.Add(file);
        }
        if (entries.Count > maxEntries)
        {
            foreach (var file in entries.OrderBy(f => f.LastWriteTimeUtc).Take(entries.Count - maxEntries))
                if (TryDelete(file.FullName)) deleted++;
        }
        return deleted;
    }

    private byte[] Tag(string entryId, ReadOnlySpan<byte> pe)
    {
        using var hmac = IncrementalHash.CreateHMAC(HashAlgorithmName.SHA256, _key);
        hmac.AppendData(Magic);
        hmac.AppendData(Encoding.ASCII.GetBytes(entryId));
        hmac.AppendData(pe);
        return hmac.GetHashAndReset();
    }

    private static bool IsValidId(string id)
        => id is { Length: 64 } && id.All(static c => c is >= '0' and <= '9' or >= 'a' and <= 'f');

    private static void CreatePrivateDirectory(string path)
    {
        if (OperatingSystem.IsWindows()) System.IO.Directory.CreateDirectory(path);
        else System.IO.Directory.CreateDirectory(path, UnixFileMode.UserRead | UnixFileMode.UserWrite | UnixFileMode.UserExecute);
    }

    private static bool TryDelete(string path)
    {
        try
        {
            File.Delete(path);
            return true;
        }
        catch (Exception ex) when (ex is IOException or UnauthorizedAccessException)
        {
            return false;
        }
    }
}