ScadaBridge/src/ZB.MOM.WW.ScadaBridge.CentralUI/ScriptAnalysis/SandboxWorkerProtocol.cs

using System.Buffers.Binary;
using System.Data.Common;
using System.Diagnostics;
using System.Text.Json;
using Microsoft.Extensions.Caching.Memory;
using Microsoft.Extensions.DependencyInjection;
using ZB.MOM.WW.ScadaBridge.Commons.Interfaces.Services;
using ZB.MOM.WW.ScadaBridge.Commons.Types;
using ZB.MOM.WW.ScadaBridge.Communication;

namespace ZB.MOM.WW.ScadaBridge.CentralUI.ScriptAnalysis;

/// <summary>Bounded private IPC for disposable, killable Test Run processes.</summary>
public static class SandboxWorkerProtocol
{
    private const int FrameLimit = 1024 * 1024;
    private static readonly SemaphoreSlim Admission = new(4, 4);
    private static readonly JsonSerializerOptions Json = new() { IgnoreReadOnlyProperties = true };
    private sealed record Frame(string Kind, string? Method, JsonElement Data, string? Error = null);
    private sealed record Start(SandboxRunRequest Request, bool External, bool Database, bool Bound);
    private sealed record ConnectionInfo(string Provider, string ConnectionString, bool Open);

    /// <summary>Runs one child under a deadline; a slot remains charged until confirmed process exit.</summary>
    public static async Task<SandboxRunResult> RunAsync(SandboxRunRequest request, IServiceProvider services,
        ISharedScriptCatalog catalog, CancellationToken caller)
        => await RunCoreAsync(request, services, catalog, caller);

    internal static async Task<SandboxRunResult> RunCoreAsync(SandboxRunRequest request, IServiceProvider services,
        ISharedScriptCatalog catalog, CancellationToken caller, Action<int>? started = null)
    {
        var seconds = Math.Clamp(request.TimeoutSeconds ?? 5, 1, 10);
        using var deadline = CancellationTokenSource.CreateLinkedTokenSource(caller);
        deadline.CancelAfter(TimeSpan.FromSeconds(seconds));
        var clock = Stopwatch.StartNew();
        var acquired = false;
        Process? process = null;
        Task? stderr = null;
        CancellationTokenRegistration terminateOnDeadline = default;
        try
        {
            await Admission.WaitAsync(deadline.Token); acquired = true;
            ISandboxInstanceGateway? gateway = null;
            if (request.Kind != ScriptKind.InboundApi && !string.IsNullOrWhiteSpace(request.BindInstanceUniqueName))
            {
                var locator = services.GetService<IInstanceLocator>();
                var comms = services.GetService<CommunicationService>();
                var site = locator is null ? null : await locator.GetSiteIdForInstanceAsync(request.BindInstanceUniqueName, deadline.Token);
                if (site is null || comms is null) throw new ScriptSandboxException("Cannot bind instance: cross-site communication or deployment is unavailable.");
                gateway = new SandboxInstanceGateway(comms, site, request.BindInstanceUniqueName, deadline.Token);
            }
            var worker = Path.Combine(AppContext.BaseDirectory, "ZB.MOM.WW.ScadaBridge.CentralUI.SandboxWorker.dll");
            if (!File.Exists(worker)) throw new InvalidOperationException("Sandbox worker is not installed.");
            var start = new ProcessStartInfo("dotnet") { UseShellExecute = false, RedirectStandardInput = true,
                RedirectStandardOutput = true, RedirectStandardError = true, CreateNoWindow = true };
            start.ArgumentList.Add(worker);
            // A finite heap bounds accidental allocation in return serialization and script code.
            start.Environment["DOTNET_GCHeapHardLimit"] = "20000000"; // 512 MiB, hexadecimal runtime setting.
            process = Process.Start(start) ?? throw new InvalidOperationException("Could not start sandbox worker.");
            terminateOnDeadline = deadline.Token.Register(static state =>
            {
                var child = (Process)state!;
                try { if (!child.HasExited) child.Kill(entireProcessTree: true); }
                catch (InvalidOperationException) { /* Concurrent normal process exit. */ }
            }, process);
            started?.Invoke(process.Id);
            stderr = DrainAsync(process.StandardError.BaseStream);
            await WriteAsync(process.StandardInput.BaseStream, new Frame("start", null,
                Encode(new Start(request, services.GetService<IExternalSystemClient>() is not null,
                    services.GetService<IDatabaseGateway>() is not null, gateway is not null))), deadline.Token);
            while (true)
            {
                var frame = await ReadAsync(process.StandardOutput.BaseStream, deadline.Token);
                if (frame.Kind == "result")
                    return frame.Data.Deserialize<SandboxRunResult>(Json) ?? throw new InvalidDataException("Missing sandbox result.");
                if (frame.Kind != "call") throw new InvalidDataException("Unexpected sandbox message.");
                deadline.Token.ThrowIfCancellationRequested();
                Frame reply;
                try
                {
                    var value = await DispatchAsync(frame.Method!, frame.Data, services, catalog, gateway, deadline.Token);
                    deadline.Token.ThrowIfCancellationRequested();
                    reply = new Frame("reply", null, Encode(value));
                }
                catch (Exception ex) when (ex is not OperationCanceledException)
                { reply = new Frame("reply", null, Encode(null), ex is ScriptSandboxException ? ex.Message : "Sandbox gateway request failed."); }
                await WriteAsync(process.StandardInput.BaseStream, reply, deadline.Token);
            }
        }
        catch (OperationCanceledException)
        { return new(false, null, null, "", $"Script execution exceeded the {seconds}-second sandbox timeout or was canceled.", SandboxErrorKind.Timeout, clock.ElapsedMilliseconds, null); }
        catch (ScriptSandboxException ex)
        { return new(false, null, null, "", ex.Message, SandboxErrorKind.SandboxLimitation, clock.ElapsedMilliseconds, null); }
        catch (Exception)
        { return new(false, null, null, "", "Sandbox worker exited or exceeded its resource budget.", SandboxErrorKind.RuntimeError, clock.ElapsedMilliseconds, null); }
        finally
        {
            deadline.Cancel(); // Stop accepting/brokering new operations before termination.
            terminateOnDeadline.Dispose();
            if (process is not null)
            {
                try { if (!process.HasExited) process.Kill(entireProcessTree: true); }
                catch (InvalidOperationException) when (process.HasExited) { }
                await process.WaitForExitAsync(CancellationToken.None);
                if (stderr is not null) await stderr;
                process.Dispose();
            }
            if (acquired) Admission.Release();
        }
    }

    private static async Task DrainAsync(Stream stream)
    { var buffer = new byte[4096]; while (await stream.ReadAsync(buffer) != 0) { } }

    private static async Task<object?> DispatchAsync(string method, JsonElement args, IServiceProvider services,
        ISharedScriptCatalog catalog, ISandboxInstanceGateway? gateway, CancellationToken ct)
    {
        string Text(int i) => args[i].GetString()!;
        IReadOnlyDictionary<string, object?>? Parameters(int i) => args[i].ValueKind == JsonValueKind.Null ? null
            : args[i].EnumerateObject().ToDictionary(p => p.Name, p => ScriptAnalysisService.JsonElementToObject(p.Value));
        switch (method)
        {
            case "shared": return await catalog.GetByNameAsync(Text(0), ct);
            case "external": return await services.GetRequiredService<IExternalSystemClient>().CallAsync(Text(0), Text(1), Parameters(2), ct);
            case "cachedExternal": return await services.GetRequiredService<IExternalSystemClient>().CachedCallAsync(Text(0), Text(1), Parameters(2), Text(3), ct);
            case "cachedDb": return await services.GetRequiredService<IDatabaseGateway>().CachedWriteAsync(Text(0), Text(1), Parameters(2), Text(3), ct);
            case "connection":
                var source = services.GetRequiredService<IDatabaseGateway>() as ISandboxDatabaseConnectionSource
                    ?? throw new ScriptSandboxException("Database worker handoff is unavailable for this provider.");
                return new ConnectionInfo("Microsoft.Data.SqlClient.SqlConnection",
                    await source.GetSandboxConnectionStringAsync(Text(0), ct), true);
            case "get": return await (gateway ?? throw new ScriptSandboxException("Instance not bound.")).GetAttributeAsync(Text(0), ct);
            case "set": await gateway!.SetAttributeAsync(Text(0), Text(1), ct); return null;
            case "script": return await gateway!.CallScriptAsync(Text(0), Parameters(1), ct);
            case "wait": return await gateway!.WaitForAttributeAsync(Text(0), args[1].GetString(), TimeSpan.FromTicks(args[2].GetInt64()), args[3].GetBoolean(), ct);
            default: throw new ScriptSandboxException("Unknown sandbox operation.");
        }
    }

    private static JsonElement Encode(object? value)
    {
        using var buffer = new BoundedSandboxStream(FrameLimit - 4096);
        JsonSerializer.Serialize(buffer, value, Json);
        using var document = JsonDocument.Parse(buffer.GetBuffer().AsMemory(0, (int)buffer.Length));
        return document.RootElement.Clone();
    }

    private static async Task WriteAsync(Stream stream, Frame frame, CancellationToken ct)
    {
        using var bytes = new BoundedSandboxStream(FrameLimit);
        JsonSerializer.Serialize(bytes, frame, Json);
        var prefix = new byte[4]; BinaryPrimitives.WriteInt32LittleEndian(prefix, (int)bytes.Length);
        await stream.WriteAsync(prefix, ct);
        await stream.WriteAsync(bytes.GetBuffer().AsMemory(0, (int)bytes.Length), ct);
        await stream.FlushAsync(ct);
    }
    private static async Task<Frame> ReadAsync(Stream stream, CancellationToken ct)
    {
        var prefix = new byte[4]; await stream.ReadExactlyAsync(prefix, ct);
        var length = BinaryPrimitives.ReadInt32LittleEndian(prefix);
        if (length is < 1 or > FrameLimit) throw new InvalidDataException("Sandbox frame budget exceeded.");
        var data = new byte[length]; await stream.ReadExactlyAsync(data, ct);
        return JsonSerializer.Deserialize<Frame>(data, Json) ?? throw new InvalidDataException("Invalid sandbox frame.");
    }

    /// <summary>Worker executable entry. Protocol streams are captured before script console routing.</summary>
    public static async Task RunWorkerAsync()
    {
        var input = Console.OpenStandardInput(); var output = Console.OpenStandardOutput();
        var first = await ReadAsync(input, CancellationToken.None);
        var start = first.Data.Deserialize<Start>(Json) ?? throw new InvalidDataException("Missing worker request.");
        var rpc = new Rpc(input, output);
        var services = new ServiceCollection();
        if (start.External) services.AddSingleton<IExternalSystemClient>(rpc);
        if (start.Database) services.AddSingleton<IDatabaseGateway>(rpc);
        if (start.Bound) services.AddSingleton<ISandboxInstanceGateway>(rpc);
        using var provider = services.BuildServiceProvider();
        using var cache = new MemoryCache(new MemoryCacheOptions());
        var runner = new ScriptAnalysisService(rpc, cache, provider);
        var result = await runner.RunInWorkerAsync(start.Request, CancellationToken.None);
        await WriteAsync(output, new Frame("result", null, Encode(result)), CancellationToken.None);
    }

    private sealed class Rpc(Stream input, Stream output) : ISharedScriptCatalog, IExternalSystemClient, IDatabaseGateway, ISandboxInstanceGateway
    {
        private readonly SemaphoreSlim _gate = new(1);
        private async Task<JsonElement> Call(string method, object?[] args, CancellationToken ct)
        {
            await _gate.WaitAsync(ct);
            try
            {
                await WriteAsync(output, new Frame("call", method, Encode(args)), ct);
                var reply = await ReadAsync(input, ct);
                if (reply.Error is not null) throw new ScriptSandboxException(reply.Error);
                return reply.Data;
            }
            finally { _gate.Release(); }
        }
        public Task<IReadOnlyList<ScriptShape>> GetShapesAsync() => Task.FromResult<IReadOnlyList<ScriptShape>>([]);
        public async Task<SharedScriptSource?> GetByNameAsync(string name, CancellationToken cancellationToken = default)
            => (await Call("shared", [name], cancellationToken)).Deserialize<SharedScriptSource>(Json);
        public async Task<ExternalCallResult> CallAsync(string systemName, string methodName, IReadOnlyDictionary<string, object?>? parameters = null, CancellationToken cancellationToken = default)
            => (await Call("external", [systemName, methodName, parameters], cancellationToken)).Deserialize<ExternalCallResult>(Json)!;
        public async Task<ExternalCallResult> CachedCallAsync(string systemName, string methodName, IReadOnlyDictionary<string, object?>? parameters = null, string? originInstanceName = null, CancellationToken cancellationToken = default, TrackedOperationId? trackedOperationId = null, Guid? executionId = null, string? sourceScript = null, Guid? parentExecutionId = null)
            => (await Call("cachedExternal", [systemName, methodName, parameters, originInstanceName], cancellationToken)).Deserialize<ExternalCallResult>(Json)!;
        public async Task<ExternalCallResult> CachedWriteAsync(string connectionName, string sql, IReadOnlyDictionary<string, object?>? parameters = null, string? originInstanceName = null, CancellationToken cancellationToken = default, TrackedOperationId? trackedOperationId = null, Guid? executionId = null, string? sourceScript = null, Guid? parentExecutionId = null)
            => (await Call("cachedDb", [connectionName, sql, parameters, originInstanceName], cancellationToken)).Deserialize<ExternalCallResult>(Json)!;
        public async Task<DbConnection> GetConnectionAsync(string connectionName, CancellationToken cancellationToken = default)
        {
            var info = (await Call("connection", [connectionName], cancellationToken)).Deserialize<ConnectionInfo>(Json)!;
            DbConnection connection = info.Provider switch {
                "Microsoft.Data.SqlClient.SqlConnection" => new Microsoft.Data.SqlClient.SqlConnection(info.ConnectionString),
                _ => throw new ScriptSandboxException("Unsupported database provider.") };
            if (info.Open) await connection.OpenAsync(cancellationToken);
            return connection;
        }
        public async Task<object?> GetAttributeAsync(string canonicalName, CancellationToken ct)
            => ScriptAnalysisService.JsonElementToObject(await Call("get", [canonicalName], ct));
        public async Task SetAttributeAsync(string canonicalName, string value, CancellationToken ct) => await Call("set", [canonicalName, value], ct);
        public async Task<object?> CallScriptAsync(string canonicalScriptName, IReadOnlyDictionary<string, object?>? parameters, CancellationToken ct)
            => ScriptAnalysisService.JsonElementToObject(await Call("script", [canonicalScriptName, parameters], ct));
        public async Task<WaitResult> WaitForAttributeAsync(string canonicalName, string? targetValueEncoded, TimeSpan timeout, bool requireGoodQuality, CancellationToken ct)
            => (await Call("wait", [canonicalName, targetValueEncoded, timeout.Ticks, requireGoodQuality], ct)).Deserialize<WaitResult>(Json)!;
    }
}