Genarrative/plugins/agc-unity-editor/dotnet/src/ProtocolServer.cs

using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;

namespace Agc.Unity.Attach;

public static class ProtocolServer
{
    private static readonly UTF8Encoding StrictUtf8 = new(false, true);

    public static async Task RunAsync(Stream input, Stream output, HelperHost host, CancellationToken cancellationToken = default)
    {
        using var writeGate = new SemaphoreSlim(1, 1);
        var pending = new HashSet<Task>();
        var reader = new BoundedLineReader(input, HelperHost.MaximumMessageBytes);
        while (true)
        {
            byte[]? line;
            try { line = await reader.ReadAsync(cancellationToken); }
            catch (HostException error)
            {
                await WriteAsync(HelperHost.ErrorReply(null, error));
                break;
            }
            if (line == null) break;
            JsonObject request;
            try
            {
                request = JsonNode.Parse(StrictUtf8.GetString(line), documentOptions: new JsonDocumentOptions { MaxDepth = 32 }) as JsonObject
                    ?? throw new HostException("UnityProtocolInvalid", "A request must be a JSON object.");
            }
            catch (Exception error) when (error is JsonException or DecoderFallbackException or HostException)
            {
                await WriteAsync(HelperHost.ErrorReply(null, new HostException("UnityProtocolInvalid", "Invalid UTF-8 JSON request.")));
                continue;
            }
            var response = host.HandleAsync(request);
            var delivery = DeliverAsync(response);
            pending.RemoveWhere(task => task.IsCompleted);
            pending.Add(delivery);
            if (pending.Count >= 32) await Task.WhenAny(pending);
        }
        await Task.WhenAll(pending);

        async Task DeliverAsync(Task<JsonObject> response) => await WriteAsync(await response);
        async Task WriteAsync(JsonObject response)
        {
            var bytes = StrictUtf8.GetBytes(response.ToJsonString());
            if (bytes.Length > HelperHost.MaximumMessageBytes)
            {
                host.MarkOutputUncertain();
                var id = response["id"]?.GetValue<long>();
                bytes = StrictUtf8.GetBytes(HelperHost.ErrorReply(id, new HostException("UnityResponseTooLarge", "Unity response exceeds 2 MiB; the execution requires reconciliation.")).ToJsonString());
            }
            await writeGate.WaitAsync(cancellationToken);
            try
            {
                await output.WriteAsync(bytes, cancellationToken);
                await output.WriteAsync(new byte[] { (byte)'\n' }, cancellationToken);
                await output.FlushAsync(cancellationToken);
            }
            finally { writeGate.Release(); }
        }
    }
}

public sealed class BoundedLineReader(Stream input, int maximumBytes)
{
    private readonly byte[] buffer = new byte[8192];
    private int offset;
    private int available;

    public async Task<byte[]?> ReadAsync(CancellationToken cancellationToken = default)
    {
        using var line = new MemoryStream();
        while (true)
        {
            if (offset == available)
            {
                available = await input.ReadAsync(buffer, cancellationToken);
                offset = 0;
                if (available == 0)
                {
                    if (line.Length == 0) return null;
                    throw new HostException("UnityMessageTruncated", "A JSON message must end with a newline.");
                }
            }
            var newline = Array.IndexOf(buffer, (byte)'\n', offset, available - offset);
            var end = newline < 0 ? available : newline;
            var count = end - offset;
            if (line.Length + count > maximumBytes)
                throw new HostException("UnityMessageTooLarge", "A JSON message exceeds 2 MiB.");
            line.Write(buffer, offset, count);
            offset = end;
            if (newline < 0) continue;
            offset++;
            var bytes = line.ToArray();
            return bytes.Length > 0 && bytes[^1] == '\r' ? bytes[..^1] : bytes;
        }
    }
}