XFE.SpaceEngineers.AgentBridge

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.IO.Pipes;
using System.Runtime.InteropServices;
using System.Security.AccessControl;
using System.Security.Principal;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Win32.SafeHandles;
using Newtonsoft.Json.Linq;
using XFE.SeAgent.Plugin.Protocol;

namespace XFE.SeAgent.Plugin.Transport
{
    /// <summary>Local, current-user JSON RPC. Only Pump invokes game-facing callbacks.</summary>
    public sealed class AgentServer : IDisposable
    {
        public const int MaximumConnections = 4;
        public const int MaximumQueuedRequests = 128;
        private const int ReadTimeoutMs = 10000;
        private const int WriteTimeoutMs = 5000;
        private readonly Func<string, JObject, JObject> execute;
        private readonly Action<string> log;
        private readonly object gate = new object();
        private readonly Queue<PendingRequest> requests = new Queue<PendingRequest>();
        private readonly Queue<string> logs = new Queue<string>();
        private readonly HashSet<NamedPipeServerStream> pipes = new HashSet<NamedPipeServerStream>();
        private readonly CancellationTokenSource shutdown = new CancellationTokenSource();
        private int disposed;
        private int pumpThread;
        private int pumping;
        private bool started;

        public string PipeName { get; }

        public AgentServer(Func<string, JObject, JObject> execute, Action<string> log)
        {
            this.execute = execute ?? throw new ArgumentNullException(nameof(execute));
            this.log = log ?? throw new ArgumentNullException(nameof(log));
            PipeName = "XFE.SE.Agent." + Process.GetCurrentProcess().Id;
        }

        public void Start()
        {
            lock (gate)
            {
                ThrowIfDisposed();
                if (started) return;
                var initial = new List<NamedPipeServerStream>();
                try
                {
                    for (int i = 0; i < MaximumConnections; i++)
                    {
                        var pipe = CreatePipe();
                        initial.Add(pipe);
                        pipes.Add(pipe);
                    }
                    started = true;
                    foreach (var pipe in initial) _ = Task.Run(() => ListenAsync(pipe));
                    EnqueueLog("Agent pipe ready: " + PipeName);
                }
                catch
                {
                    foreach (var pipe in initial) { pipes.Remove(pipe); pipe.Dispose(); }
                    throw;
                }
            }
        }

        private NamedPipeServerStream CreatePipe()
        {
            using (var identity = WindowsIdentity.GetCurrent())
            {
                var user = identity.User ?? throw new InvalidOperationException("Current Windows user SID is unavailable.");
                var security = new PipeSecurity();
                security.SetAccessRuleProtection(true, false);
                security.SetOwner(user);
                security.AddAccessRule(new PipeAccessRule(new SecurityIdentifier(WellKnownSidType.NetworkSid, null), PipeAccessRights.FullControl, AccessControlType.Deny));
                security.AddAccessRule(new PipeAccessRule(user, PipeAccessRights.FullControl, AccessControlType.Allow));
#if NETFRAMEWORK
                return new NamedPipeServerStream(PipeName, PipeDirection.InOut, MaximumConnections,
                    PipeTransmissionMode.Byte, PipeOptions.Asynchronous, 4096, 4096, security);
#else
                return NamedPipeServerStreamAcl.Create(PipeName, PipeDirection.InOut, MaximumConnections,
                    PipeTransmissionMode.Byte, PipeOptions.Asynchronous, 4096, 4096, security);
#endif
            }
        }

        private async Task ListenAsync(NamedPipeServerStream first)
        {
            NamedPipeServerStream pipe = first;
            while (Volatile.Read(ref disposed) == 0)
            {
                try
                {
                    await pipe.WaitForConnectionAsync(shutdown.Token).ConfigureAwait(false);
                    await HandleConnectionAsync(pipe).ConfigureAwait(false);
                }
                catch (OperationCanceledException) when (Volatile.Read(ref disposed) != 0) { }
                catch (ObjectDisposedException) when (Volatile.Read(ref disposed) != 0) { }
                catch (Exception error)
                {
                    if (Volatile.Read(ref disposed) == 0) EnqueueLog("Agent connection ended: " + error.GetType().Name);
                }
                finally
                {
                    lock (gate) pipes.Remove(pipe);
                    pipe.Dispose();
                }
                lock (gate)
                {
                    if (disposed != 0) return;
                    try { pipe = CreatePipe(); pipes.Add(pipe); }
                    catch (Exception error) { EnqueueLog("Agent pipe listener stopped: " + error.GetType().Name); return; }
                }
            }
        }

        private async Task HandleConnectionAsync(NamedPipeServerStream pipe)
        {
            AgentRequest request;
            try
            {
                string line;
                using (var readTimeout = CancellationTokenSource.CreateLinkedTokenSource(shutdown.Token))
                {
                    readTimeout.CancelAfter(ReadTimeoutMs);
                    using (readTimeout.Token.Register(() => ClosePipe(pipe)))
                        line = await ReadLineAsync(pipe, readTimeout.Token).ConfigureAwait(false);
                }
                request = AgentProtocol.ParseRequest(line);
            }
            catch (AgentProtocolException error)
            {
                await WriteResponseAsync(pipe, AgentProtocol.Error(error.RequestId, error.Code, error.Message)).ConfigureAwait(false);
                return;
            }
            catch (System.Text.DecoderFallbackException)
            {
                await WriteResponseAsync(pipe, AgentProtocol.Error(null, "parse_error", "Request must be valid UTF-8.")).ConfigureAwait(false);
                return;
            }

            var pending = new PendingRequest(request, () => IsClientPresent(pipe), () => Volatile.Read(ref disposed) != 0);
            lock (gate)
            {
                if (disposed != 0) pending.Cancel("server_stopped", "Agent server has stopped.");
                else if (requests.Count >= MaximumQueuedRequests) pending.Cancel("server_busy", "Request queue is full. Retry later.");
                else requests.Enqueue(pending);
            }

            // One connection carries exactly one request. A second read detects remote close
            // while the game thread is busy, without a worker touching any game objects.
            Task disconnected = ObserveDisconnectAsync(pipe, pending);
            using (var wait = CancellationTokenSource.CreateLinkedTokenSource(shutdown.Token))
            {
                Task timeout = Task.Delay(request.TimeoutMs, wait.Token);
                Task finished = await Task.WhenAny(pending.Completion, disconnected, timeout).ConfigureAwait(false);
                if (finished == timeout)
                    pending.Cancel(Volatile.Read(ref disposed) == 0 ? "timeout" : "server_stopped", "Request expired before execution.");
                JObject response = await pending.Completion.ConfigureAwait(false);
                wait.Cancel();
                await WriteResponseAsync(pipe, response).ConfigureAwait(false);
            }
        }

        private static async Task<string> ReadLineAsync(NamedPipeServerStream pipe, CancellationToken token)
        {
            var buffer = new byte[4096];
            using (var line = new MemoryStream())
            {
                while (true)
                {
                    int count = await pipe.ReadAsync(buffer, 0, buffer.Length, token).ConfigureAwait(false);
                    if (count == 0) throw new EndOfStreamException("Client closed before sending a complete request line.");
                    int newline = Array.IndexOf(buffer, (byte)'\n', 0, count);
                    int length = newline < 0 ? count : newline;
                    if (line.Length + length > AgentProtocol.MaxRequestBytes)
                        throw new AgentProtocolException("request_too_large", "Request exceeds the 1 MiB limit.");
                    line.Write(buffer, 0, length);
                    if (newline < 0) continue;
                    if (newline != count - 1)
                        throw new AgentProtocolException("invalid_request", "Only one request line is allowed per connection.");
                    byte[] data = line.ToArray();
                    int textLength = data.Length > 0 && data[data.Length - 1] == '\r' ? data.Length - 1 : data.Length;
                    return AgentProtocol.Utf8.GetString(data, 0, textLength);
                }
            }
        }

        private static async Task ObserveDisconnectAsync(NamedPipeServerStream pipe, PendingRequest pending)
        {
            try
            {
                var extra = new byte[1];
                int count = await pipe.ReadAsync(extra, 0, 1).ConfigureAwait(false);
                pending.Cancel(count == 0 ? "disconnected" : "invalid_request",
                    count == 0 ? "Client disconnected before completion." : "Only one request line is allowed per connection.");
            }
            catch (Exception error) when (error is IOException || error is ObjectDisposedException || error is InvalidOperationException || error is OperationCanceledException)
            {
                pending.Cancel("disconnected", "Client disconnected before completion.");
            }
        }

        private async Task WriteResponseAsync(NamedPipeServerStream pipe, JObject response)
        {
            byte[] bytes;
            try { bytes = AgentProtocol.SerializeResponse(response); }
            catch (AgentRpcException error) { bytes = AgentProtocol.SerializeResponse(AgentProtocol.Error(response["id"], error.Code, error.Message)); }
            using (var timeout = CancellationTokenSource.CreateLinkedTokenSource(shutdown.Token))
            {
                timeout.CancelAfter(WriteTimeoutMs);
                using (timeout.Token.Register(() => ClosePipe(pipe)))
                {
                    await pipe.WriteAsync(bytes, 0, bytes.Length, timeout.Token).ConfigureAwait(false);
                    await pipe.FlushAsync(timeout.Token).ConfigureAwait(false);
                }
            }
        }

        public void Pump(int maxRequests = 4)
        {
            if (maxRequests < 0) throw new ArgumentOutOfRangeException(nameof(maxRequests));
            if (Volatile.Read(ref disposed) != 0) return;
            int thread = Thread.CurrentThread.ManagedThreadId;
            int owner = Interlocked.CompareExchange(ref pumpThread, thread, 0);
            if (owner != 0 && owner != thread) throw new InvalidOperationException("AgentServer.Pump must remain on the game Update thread.");
            if (Interlocked.CompareExchange(ref pumping, 1, 0) != 0) return;
            try
            {
                for (int i = 0; i < 8; i++)
                {
                    string entry;
                    lock (gate) { if (logs.Count == 0) break; entry = logs.Dequeue(); }
                    try { log(entry); } catch { /* Logging must not stop the game update loop. */ }
                }
                int executed = 0;
                for (int examined = 0; examined < MaximumQueuedRequests && executed < Math.Min(maxRequests, MaximumQueuedRequests); examined++)
                {
                    PendingRequest pending;
                    lock (gate) { if (requests.Count == 0 || disposed != 0) break; pending = requests.Dequeue(); }
                    if (!pending.TryStart()) continue;
                    executed++;
                    try { pending.Complete(AgentProtocol.Success(pending.Request.Id, execute(pending.Request.Method, pending.Request.Parameters))); }
                    catch (AgentRpcException error) { pending.Complete(AgentProtocol.Error(pending.Request.Id, error.Code, error.Message)); }
                    catch (Exception error) { pending.Complete(AgentProtocol.Error(pending.Request.Id, "execution_error", error.Message)); }
                }
            }
            finally { Volatile.Write(ref pumping, 0); }
        }

        private void EnqueueLog(string message)
        {
            lock (gate)
            {
                if (logs.Count >= MaximumQueuedRequests) logs.Dequeue();
                logs.Enqueue(message);
            }
        }

        private static bool IsClientPresent(NamedPipeServerStream pipe)
        {
            try
            {
                uint available;
                return pipe.IsConnected && PeekNamedPipe(pipe.SafePipeHandle, IntPtr.Zero, 0, IntPtr.Zero, out available, IntPtr.Zero);
            }
            catch (Exception error) when (error is IOException || error is ObjectDisposedException || error is InvalidOperationException) { return false; }
        }

        [DllImport("kernel32.dll", SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        private static extern bool PeekNamedPipe(SafePipeHandle pipe, IntPtr buffer, uint bufferSize, IntPtr bytesRead, out uint bytesAvailable, IntPtr bytesLeftThisMessage);

        private static void ClosePipe(NamedPipeServerStream pipe)
        {
            try { pipe.Dispose(); } catch (IOException) { }
        }

        private void ThrowIfDisposed()
        {
            if (Volatile.Read(ref disposed) != 0) throw new ObjectDisposedException(nameof(AgentServer));
        }

        public void Dispose()
        {
            if (Interlocked.Exchange(ref disposed, 1) != 0) return;
            NamedPipeServerStream[] open;
            lock (gate)
            {
                foreach (var pending in requests) pending.Cancel("server_stopped", "Agent server has stopped.");
                requests.Clear();
                open = new NamedPipeServerStream[pipes.Count];
                pipes.CopyTo(open);
            }
            shutdown.Cancel();
            foreach (var pipe in open) ClosePipe(pipe);
            // The cancellation source stays alive until outstanding async pipe operations finish.
        }
    }
}