DodoSSH/src/DodoSSH.Client.Terminal/TerminalWorkspace.cs

using DodoSSH.Client.Ssh;

namespace DodoSSH.Client.Terminal;

/// <summary>Tuning for the workspace.</summary>
public sealed class TerminalWorkspaceOptions
{
    /// <summary>
    /// How long <see cref="TerminalWorkspace.WaitForRendererAsync"/> waits for the renderer's socket
    /// before giving up.
    /// </summary>
    /// <remarks>
    /// <para>
    /// The value has to separate two cases. Attaching is normally near-instant: WebView2 starts with the
    /// window and the page has usually attached while the user was still typing a passphrase. But a first
    /// run on a cold profile creates a user-data directory and starts a process tree of some thirty-five
    /// processes, and on a slow or loaded machine that is seconds rather than milliseconds. A renderer
    /// that will never attach — no Evergreen runtime, an install blocked by policy, an AppContainer that
    /// cannot reach loopback — will not attach however long the wait is.
    /// </para>
    /// <para>
    /// So being generous costs only how long a genuinely broken WebView2 takes to say so, while being
    /// tight costs telling someone their runtime is broken when it was merely slow. Fifteen seconds is
    /// well clear of any cold start observed here and is still an answer rather than a hang.
    /// </para>
    /// </remarks>
    public TimeSpan RendererTimeout { get; init; } = TimeSpan.FromSeconds(15);
}

/// <summary>A session whose shell has ended.</summary>
/// <param name="sessionId">The session, as the renderer and the workspace know it.</param>
public sealed class TerminalSessionEndedEventArgs(uint sessionId) : EventArgs
{
    /// <summary>The session that ended.</summary>
    public uint SessionId { get; } = sessionId;
}

/// <summary>
/// Owns the loopback data plane and every live terminal session.
/// </summary>
/// <remarks>
/// <para>
/// One data plane and one renderer page for the whole application, with a session id per terminal.
/// Not one WebView per tab: each WebView2 is a separate browser process, so twenty tabs would mean
/// twenty renderer processes and several hundred megabytes for a working set a user would call
/// ordinary. Splits and tabs are layout inside the single page.
/// </para>
/// <para>
/// <b>A session's lifetime is the application's, not the vault's.</b> This object is composed once at
/// startup and outlives every lock, deliberately: locking the vault zeroes keys, and a shell needs no
/// vault key to keep running, so a job started before the lock keeps running through it. That is a
/// policy rather than an oversight — <c>MainWindowViewModel.LockAsync</c> says why, and the shell shows
/// <see cref="LiveSessionCount"/> on the unlock screen so it is not a hidden state.
/// </para>
/// </remarks>
public sealed class TerminalWorkspace : IAsyncDisposable
{
    private readonly TerminalDataPlane dataPlane;
    private readonly ISshConnectionFactory connections;
    private readonly TimeProvider clock;
    private readonly TerminalWorkspaceOptions options;
    private readonly Dictionary<uint, LiveSession> sessions = [];

    /// <summary>
    /// Guards <see cref="sessions"/> and <see cref="nextSessionId"/>.
    /// </summary>
    /// <remarks>
    /// This dictionary is genuinely touched from more than one thread, which the tab strip made true rather
    /// than merely arguable: sessions are opened and closed from the UI thread, and
    /// <see cref="RunSessionAsync"/>'s finally now asks whether the session is still known — on whichever
    /// thread-pool thread the pump happened to unwind on. An unsynchronised <see cref="Dictionary{TKey,
    /// TValue}"/> read concurrent with a write does not merely return a stale answer; it can throw or spin.
    /// The data plane guards its own registration table the same way and for the same reason.
    /// </remarks>
    private readonly Lock sessionGate = new();
    private readonly CancellationTokenSource lifetime = new();

    private uint nextSessionId = 1;
    private Task? server;
    private int disposed;

    /// <param name="assets">Where the renderer's files come from.</param>
    /// <param name="connections">How SSH connections are made.</param>
    /// <param name="clock">Time source, so the pumps' flush interval is testable.</param>
    /// <param name="options">Tuning, or null for the defaults.</param>
    public TerminalWorkspace(
        ITerminalAssetProvider assets,
        ISshConnectionFactory connections,
        TimeProvider clock,
        TerminalWorkspaceOptions? options = null)
    {
        this.connections = connections;
        this.clock = clock;
        this.options = options ?? new TerminalWorkspaceOptions();

        dataPlane = new TerminalDataPlane(assets);
    }

    /// <summary>Where the WebView should navigate.</summary>
    public Uri PageUrl => dataPlane.PageUrl;

    /// <summary>
    /// How many terminals still have a live shell behind them.
    /// </summary>
    /// <remarks>
    /// <para>
    /// Not <c>sessions.Count</c>, which over-reports. Nothing removes an entry when the remote closes
    /// the channel on its own — <see cref="RunSessionAsync"/> only drops the renderer registration — so
    /// a session whose shell exited half an hour ago is still in the dictionary. A completed
    /// <c>Run</c> task is what "the shell is gone" actually looks like: the pump's loops have finished
    /// and it has already sent <c>SessionClosed</c> to the renderer.
    /// </para>
    /// <para>
    /// This exists because the shell shows the number on the unlock screen, and a lock screen that
    /// claims a shell is still running when it is not would be the same class of dishonesty the number
    /// is there to prevent.
    /// </para>
    /// </remarks>
    public int LiveSessionCount
    {
        get
        {
            lock (sessionGate)
            {
                return sessions.Values.Count(session => !session.Run.IsCompleted);
            }
        }
    }

    /// <summary>
    /// Whether one session's shell is still running.
    /// </summary>
    /// <remarks>
    /// The same question <see cref="LiveSessionCount"/> answers in aggregate, and answered the same way: a
    /// completed <c>Run</c> is what "the shell is gone" looks like, because nothing removes the entry when
    /// the remote closes the channel on its own. A session id this workspace never issued is not live.
    /// </remarks>
    public bool IsSessionLive(uint sessionId)
    {
        lock (sessionGate)
        {
            return sessions.TryGetValue(sessionId, out var session) && !session.Run.IsCompleted;
        }
    }

    /// <summary>
    /// Raised with the session id when a shell ends on its own.
    /// </summary>
    /// <remarks>
    /// <para>
    /// A tab has to be able to stop claiming it is connected, and polling would be the alternative: a timer
    /// asking <see cref="IsSessionLive"/> often enough to look immediate, for a thing that happens a handful
    /// of times a day. The renderer already learns about this instantly — the pump sends
    /// <c>SessionClosed</c> and the page writes the reason into the pane — so this is the same fact reaching
    /// the half of the interface Avalonia draws.
    /// </para>
    /// <para>
    /// <b>Raised on whatever thread the pump finished on</b>, which is a thread-pool thread. A handler that
    /// touches an observable collection has to marshal; this type has no toolkit to do it with, which is
    /// exactly why it does not try.
    /// </para>
    /// <para>
    /// Not raised by <see cref="CloseSessionAsync"/>. That path already has a caller who knows the session is
    /// going, and telling it what it just asked for is how a tab close turns into a second tab close.
    /// </para>
    /// </remarks>
    public event EventHandler<TerminalSessionEndedEventArgs>? SessionEnded;

    /// <summary>Starts the loopback listener.</summary>
    public void Start() => server = dataPlane.RunAsync(lifetime.Token);

    /// <summary>
    /// Waits until the renderer page has attached its socket.
    /// </summary>
    /// <remarks>
    /// <para>
    /// A session opened before the renderer attaches would have its <c>SessionOpened</c> frame dropped —
    /// the transport discards frames when nothing is connected — leaving output arriving for a terminal
    /// that was never created. The gate is the invariant and stays.
    /// </para>
    /// <para>
    /// Bounded, because whether the renderer attaches at all depends on a WebView2 runtime this process
    /// does not install. An unbounded wait turned a missing runtime into a Connect that never returned,
    /// with the caller's busy state never cleared and nothing on screen to explain it. Callers are
    /// expected to translate the timeout into something that names the runtime, because
    /// <see cref="TimeoutException"/>'s own message names nothing.
    /// </para>
    /// </remarks>
    /// <param name="cancellationToken">Abandons the wait.</param>
    /// <exception cref="TimeoutException">
    /// No renderer attached within <see cref="TerminalWorkspaceOptions.RendererTimeout"/>.
    /// </exception>
    public Task WaitForRendererAsync(CancellationToken cancellationToken) =>
        dataPlane.RendererAttached.WaitAsync(options.RendererTimeout, cancellationToken);

    /// <summary>Connects to a host and starts a terminal for it.</summary>
    /// <returns>The session id, which identifies this terminal in the renderer.</returns>
    public async Task<uint> OpenSessionAsync(
        SshConnectionRequest request,
        TerminalSize size,
        CancellationToken cancellationToken)
    {
        var connection = await connections.ConnectAsync(request, cancellationToken).ConfigureAwait(false);

        ISshShellSession shell;
        try
        {
            shell = await connection.OpenShellAsync(size, cancellationToken).ConfigureAwait(false);
        }
        catch
        {
            await connection.DisposeAsync().ConfigureAwait(false);
            throw;
        }

        uint sessionId;
        TerminalSessionPump pump;

        lock (sessionGate)
        {
            sessionId = nextSessionId++;
            pump = new TerminalSessionPump(sessionId, shell, dataPlane, clock);
        }

        dataPlane.Register(sessionId, pump);

        // Registered before running, so an acknowledgement that arrives with the very first output
        // frame has somewhere to go.
        var run = RunSessionAsync(sessionId, pump);

        lock (sessionGate)
        {
            sessions[sessionId] = new LiveSession(connection, pump, run);
        }

        return sessionId;
    }

    /// <summary>
    /// Shows one terminal's pane and hides the others.
    /// </summary>
    /// <remarks>
    /// <para>
    /// The page keeps a pane per session and displays one of them, so switching tabs is this frame and
    /// nothing else — no terminal is destroyed, no scrollback is lost, and the shell behind a hidden pane
    /// goes on running and goes on being read. That is the whole reason tabs cost so little here: the
    /// expensive object is the WebView, and there is one of those however many tabs are open.
    /// </para>
    /// <para>
    /// Sent for a session id this workspace does not know as readily as for one it does. The page ignores
    /// a pane it has never created, and refusing here would mean holding a second copy of the tab strip's
    /// idea of what exists — which is the sort of duplicated truth that ends up disagreeing.
    /// </para>
    /// </remarks>
    public ValueTask ActivateSessionAsync(uint sessionId, CancellationToken cancellationToken) =>
        dataPlane.SendAsync(
            TerminalFrame.Create((byte)TerminalServerOpcode.SessionActivated, sessionId, []),
            cancellationToken);

    /// <summary>
    /// Closes one terminal.
    /// </summary>
    /// <remarks>
    /// Closing one terminal without taking the process down is what a tab close needs, and what it is now
    /// reached for. It is also what <see cref="DisposeAsync"/> uses on the way out, which was its only
    /// caller while the interface had no tabs.
    /// </remarks>
    public async Task CloseSessionAsync(uint sessionId)
    {
        LiveSession? session;

        lock (sessionGate)
        {
            if (!sessions.Remove(sessionId, out session))
            {
                return;
            }
        }

        dataPlane.Unregister(sessionId);

        // Before the teardown, while the socket is still up. The page holds an xterm instance, its whole
        // scrollback and a WebGL context per pane and never reclaimed one on its own; a frame dropped here
        // is a leak nothing else would notice, because it is inside the WebView.
        await dataPlane
            .SendAsync(
                TerminalFrame.Create((byte)TerminalServerOpcode.SessionRemoved, sessionId, []),
                CancellationToken.None)
            .ConfigureAwait(false);

        // The connection goes even if the pump's disposal throws. The entry is already out of the dictionary,
        // so nothing will come back for it — not DisposeAsync's loop, which snapshots the keys — and an
        // undisposed ISshConnection is a live authenticated channel to a remote host that this process has
        // forgotten about. Which of the two is more important is not close.
        try
        {
            await session.Pump.DisposeAsync().ConfigureAwait(false);
        }
        finally
        {
            await session.Connection.DisposeAsync().ConfigureAwait(false);
        }

        // Draining, not doing work. This await exists so the pump's loops have finished before the caller
        // moves on; the session is already dead either way, and a fault here is nothing the caller can act
        // on — the pump's last act is a best-effort SessionClosed frame on a socket that may have gone with
        // the window. Letting that escape would fault a tab-close command over a session that closed fine.
        try
        {
            await session.Run.ConfigureAwait(false);
        }
        catch (Exception exception) when (exception is not OutOfMemoryException)
        {
            // Expected on the ordinary path: disposing the pump cancels its run.
        }
    }

    /// <inheritdoc />
    public async ValueTask DisposeAsync()
    {
        if (Interlocked.Exchange(ref disposed, 1) == 1)
        {
            return;
        }

        uint[] open;

        lock (sessionGate)
        {
            open = [.. sessions.Keys];
        }

        foreach (var sessionId in open)
        {
            await CloseSessionAsync(sessionId).ConfigureAwait(false);
        }

        await lifetime.CancelAsync().ConfigureAwait(false);
        await dataPlane.DisposeAsync().ConfigureAwait(false);

        if (server is not null)
        {
            try
            {
                await server.ConfigureAwait(false);
            }
            catch (OperationCanceledException)
            {
                // Expected: the accept loop is stopped by cancelling it.
            }
        }

        lifetime.Dispose();
    }

    private async Task RunSessionAsync(uint sessionId, TerminalSessionPump pump)
    {
        try
        {
            await pump.RunAsync(lifetime.Token).ConfigureAwait(false);
        }
        finally
        {
            dataPlane.Unregister(sessionId);

            // Only when the session is still one this workspace knows about. CloseSessionAsync removes the
            // entry before it disposes the pump, so a tab the user closed does not come back as news.
            bool announce;

            lock (sessionGate)
            {
                announce = sessions.ContainsKey(sessionId);
            }

            if (announce)
            {
                SessionEnded?.Invoke(this, new TerminalSessionEndedEventArgs(sessionId));
            }
        }
    }

    private sealed record LiveSession(ISshConnection Connection, TerminalSessionPump Pump, Task Run);
}