path: root/src/Agent.Admin/Agent.Admin.Client/ElevatedNetworkSession.cs

using System.ComponentModel;
using System.Diagnostics;
using System.IO.Pipes;
using System.Runtime.InteropServices;
using System.Security.AccessControl;
using System.Security.Principal;
using Microsoft.Win32.SafeHandles;

namespace Agent.Admin.Client;

public sealed class AdminOperationException(string code, string message, Exception? inner = null) : Exception(message, inner)
{
    public string Code { get; } = code;
}

public interface IElevatedNetworkSession : IAsyncDisposable
{
    Task<bool> InspectLoopbackAsync(string applicationSid, CancellationToken cancellationToken);
    Task ApplyAsync(string applicationSid, string networkAccess, IReadOnlyList<AdminNetworkEndpoint> endpoints,
        bool originallyLoopbackExempt, CancellationToken cancellationToken);
    Task RecoverLoopbackAsync(string applicationSid, bool originallyLoopbackExempt, CancellationToken cancellationToken);
    Task CleanupAsync(CancellationToken cancellationToken);
}

public interface IElevatedNetworkSessionFactory
{
    Task<IElevatedNetworkSession> StartAsync(Guid taskId, CancellationToken cancellationToken);
}

public sealed partial class ElevatedNetworkSessionFactory(string helperPath) : IElevatedNetworkSessionFactory
{
    private const uint PipeAccessDuplex = 0x00000003;
    private const uint FileFlagOverlapped = 0x40000000;
    private const uint FileFlagWriteThrough = 0x80000000;
    private const uint PipeRejectRemoteClients = 0x00000008;

    public async Task<IElevatedNetworkSession> StartAsync(Guid taskId, CancellationToken cancellationToken)
    {
        string canonicalHelper = Path.GetFullPath(helperPath);
        if (!File.Exists(canonicalHelper))
            throw new AdminOperationException("sandbox.helper_incompatible", $"Elevated helper does not exist at the fixed Bundle path '{canonicalHelper}'.");

        string pipeName = $"AgentSandboxMcp-{Environment.ProcessId}-{Guid.NewGuid():N}";
        using WindowsIdentity identity = WindowsIdentity.GetCurrent();
        SecurityIdentifier logon = GetLogonSid(identity);
        var security = new PipeSecurity();
        security.SetAccessRuleProtection(isProtected: true, preserveInheritance: false);
        security.AddAccessRule(new PipeAccessRule(logon, PipeAccessRights.ReadWrite, AccessControlType.Allow));
        security.AddAccessRule(new PipeAccessRule(new SecurityIdentifier(WellKnownSidType.BuiltinAdministratorsSid, null),
            PipeAccessRights.ReadWrite, AccessControlType.Allow));
        security.AddAccessRule(new PipeAccessRule(new SecurityIdentifier(WellKnownSidType.NetworkSid, null),
            PipeAccessRights.FullControl, AccessControlType.Deny));

        NamedPipeServerStream pipe = CreateLocalPipe(pipeName, security);
        Process? helper = null;
        try
        {
            var start = new ProcessStartInfo
            {
                FileName = canonicalHelper,
                Arguments = $"--pipe {Quote(pipeName)} --parent-pid {Environment.ProcessId} --task {taskId:N} --protocol {AdminProtocol.Version}",
                UseShellExecute = true,
                Verb = "runas",
                WorkingDirectory = Path.GetDirectoryName(canonicalHelper)!
            };
            try { helper = Process.Start(start); }
            catch (Win32Exception ex) when (ex.NativeErrorCode == 1223)
            {
                throw new AdminOperationException("sandbox.elevation_cancelled", "The Windows elevation request was cancelled.", ex);
            }
            catch (Exception ex) when (ex is Win32Exception or InvalidOperationException)
            {
                throw new AdminOperationException("sandbox.elevation_failed", $"The elevated helper could not be started: {ex.Message}", ex);
            }
            if (helper is null) throw new AdminOperationException("sandbox.elevation_failed", "The elevated helper did not return a Process handle.");

            using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            Task connection = pipe.WaitForConnectionAsync(linked.Token);
            Task exit = helper.WaitForExitAsync(linked.Token);
            Task timeout = Task.Delay(TimeSpan.FromSeconds(30), linked.Token);
            Task completed = await Task.WhenAny(connection, exit, timeout).ConfigureAwait(false);
            if (completed == exit)
                throw new AdminOperationException("sandbox.elevation_failed", $"The elevated helper exited before connecting (exit code {helper.ExitCode}).");
            if (completed == timeout)
                throw new AdminOperationException("sandbox.elevation_failed", "The elevated helper did not connect within 30 seconds.");
            await connection.ConfigureAwait(false);
            linked.Cancel();

            if (!Native.GetNamedPipeClientProcessId(pipe.SafePipeHandle, out uint clientPid) || clientPid != helper.Id || helper.HasExited)
                throw new AdminOperationException("sandbox.ipc_rejected", "The Named Pipe peer does not match the elevated Process started for this Task.");
            await AdminProtocol.WriteAsync(pipe, new AdminRequest(AdminProtocol.Version, "hello", taskId), cancellationToken).ConfigureAwait(false);
            AdminResponse hello = await AdminProtocol.ReadAsync<AdminResponse>(pipe, cancellationToken).ConfigureAwait(false);
            EnsureSuccess(hello);
            return new ElevatedNetworkSession(pipe, helper, taskId);
        }
        catch
        {
            pipe.Dispose();
            helper?.Dispose();
            throw;
        }
    }

    private static string Quote(string value) => '"' + value.Replace("\"", "\\\"") + '"';

    private static SecurityIdentifier GetLogonSid(WindowsIdentity identity)
    {
        const int tokenLogonSid = 28;
        _ = Native.GetTokenInformation(identity.Token, tokenLogonSid, 0, 0, out uint length);
        if (length < (uint)(IntPtr.Size * 2))
            throw new AdminOperationException("sandbox.ipc_rejected", "The invoking logon SID size is invalid.");
        nint buffer = Marshal.AllocHGlobal(checked((int)length));
        try
        {
            if (!Native.GetTokenInformation(identity.Token, tokenLogonSid, buffer, length, out _))
                throw new AdminOperationException("sandbox.ipc_rejected",
                    $"The invoking logon SID is unavailable: {new Win32Exception(Marshal.GetLastPInvokeError()).Message}");
            // TOKEN_GROUPS contains a 32-bit count followed by pointer-aligned SID_AND_ATTRIBUTES.
            nint sid = Marshal.ReadIntPtr(buffer, IntPtr.Size);
            return new SecurityIdentifier(sid);
        }
        finally { Marshal.FreeHGlobal(buffer); }
    }

    private static NamedPipeServerStream CreateLocalPipe(string pipeName, PipeSecurity security)
    {
        byte[] descriptor = security.GetSecurityDescriptorBinaryForm();
        GCHandle pinned = GCHandle.Alloc(descriptor, GCHandleType.Pinned);
        try
        {
            var attributes = new Native.SecurityAttributes
            {
                Length = Marshal.SizeOf<Native.SecurityAttributes>(),
                SecurityDescriptor = pinned.AddrOfPinnedObject(),
                InheritHandle = 0
            };
            SafePipeHandle handle = Native.CreateNamedPipe($@"\\.\pipe\{pipeName}",
                PipeAccessDuplex | FileFlagOverlapped | FileFlagWriteThrough,
                PipeRejectRemoteClients, 1, 4096, 4096, 0, in attributes);
            if (handle.IsInvalid)
            {
                int error = Marshal.GetLastPInvokeError();
                handle.Dispose();
                throw new AdminOperationException("sandbox.ipc_rejected",
                    $"The local administrative Named Pipe could not be created: {new Win32Exception(error).Message}");
            }
            return new NamedPipeServerStream(PipeDirection.InOut, isAsync: true, isConnected: false, handle);
        }
        finally { pinned.Free(); }
    }

    internal static void EnsureSuccess(AdminResponse response)
    {
        if (!response.Success) throw new AdminOperationException(response.ErrorCode ?? "sandbox.network_apply_failed",
            response.Message ?? "The elevated network operation failed.");
    }

    private static partial class Native
    {
        [StructLayout(LayoutKind.Sequential)]
        internal struct SecurityAttributes
        {
            internal int Length;
            internal nint SecurityDescriptor;
            internal int InheritHandle;
        }

        [LibraryImport("kernel32.dll", EntryPoint = "CreateNamedPipeW", SetLastError = true,
            StringMarshalling = StringMarshalling.Utf16)]
        internal static partial SafePipeHandle CreateNamedPipe(string name, uint openMode, uint pipeMode,
            uint maxInstances, uint outBufferSize, uint inBufferSize, uint defaultTimeout,
            in SecurityAttributes securityAttributes);

        [LibraryImport("advapi32.dll", SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        internal static partial bool GetTokenInformation(nint tokenHandle, int tokenInformationClass,
            nint tokenInformation, uint tokenInformationLength, out uint returnLength);

        [LibraryImport("kernel32.dll", SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        internal static partial bool GetNamedPipeClientProcessId(SafePipeHandle pipe, out uint clientProcessId);
    }
}

internal sealed class ElevatedNetworkSession(NamedPipeServerStream pipe, Process helper, Guid taskId) : IElevatedNetworkSession
{
    private bool _cleaned;

    public async Task<bool> InspectLoopbackAsync(string applicationSid, CancellationToken cancellationToken)
    {
        AdminResponse response = await SendAsync(new(AdminProtocol.Version, "inspect-loopback", taskId, applicationSid), cancellationToken).ConfigureAwait(false);
        return response.LoopbackExempt ?? throw new AdminOperationException("sandbox.ipc_rejected", "The elevated helper omitted the loopback state.");
    }

    public async Task ApplyAsync(string applicationSid, string networkAccess, IReadOnlyList<AdminNetworkEndpoint> endpoints,
        bool originallyLoopbackExempt, CancellationToken cancellationToken) =>
        ElevatedNetworkSessionFactory.EnsureSuccess(await SendAsync(new(AdminProtocol.Version, "apply-network", taskId,
            applicationSid, networkAccess, endpoints, originallyLoopbackExempt), cancellationToken).ConfigureAwait(false));

    public async Task RecoverLoopbackAsync(string applicationSid, bool originallyLoopbackExempt, CancellationToken cancellationToken) =>
        ElevatedNetworkSessionFactory.EnsureSuccess(await SendAsync(new(AdminProtocol.Version, "recover-loopback", taskId,
            applicationSid, OriginallyLoopbackExempt: originallyLoopbackExempt), cancellationToken).ConfigureAwait(false));

    public async Task CleanupAsync(CancellationToken cancellationToken)
    {
        if (_cleaned) return;
        ElevatedNetworkSessionFactory.EnsureSuccess(await SendAsync(new(AdminProtocol.Version, "cleanup-network", taskId), cancellationToken).ConfigureAwait(false));
        _cleaned = true;
    }

    private async Task<AdminResponse> SendAsync(AdminRequest request, CancellationToken cancellationToken)
    {
        if (helper.HasExited) throw new AdminOperationException("sandbox.elevation_failed", "The elevated helper exited before Task Resource cleanup.");
        await AdminProtocol.WriteAsync(pipe, request, cancellationToken).ConfigureAwait(false);
        return await AdminProtocol.ReadAsync<AdminResponse>(pipe, cancellationToken).ConfigureAwait(false);
    }

    public async ValueTask DisposeAsync()
    {
        try
        {
            if (pipe.IsConnected)
            {
                if (!_cleaned) await CleanupAsync(CancellationToken.None).ConfigureAwait(false);
                await AdminProtocol.WriteAsync(pipe, new AdminRequest(AdminProtocol.Version, "shutdown", taskId), CancellationToken.None).ConfigureAwait(false);
            }
        }
        catch { }
        pipe.Dispose();
        try
        {
            if (!helper.HasExited) await helper.WaitForExitAsync().WaitAsync(TimeSpan.FromSeconds(5)).ConfigureAwait(false);
        }
        catch { if (!helper.HasExited) helper.Kill(entireProcessTree: true); }
        helper.Dispose();
    }
}