DodoSSH/src/DodoSSH.Client.Shell/ViewModels/PortForwardViewModel.cs

using System.Globalization;
using CommunityToolkit.Mvvm.ComponentModel;
using CommunityToolkit.Mvvm.Input;
using DodoSSH.Client.Domain;
using DodoSSH.Client.Ssh;

namespace DodoSSH.Client.Shell.ViewModels;

/// <summary>Where one of a tab's forwards has got to.</summary>
internal enum PortForwardState
{
    /// <summary>Not listening: switched off, or its session is over.</summary>
    Stopped = 0,

    /// <summary>Asked for and not answered yet. A remote forward waits on the server.</summary>
    Starting = 1,

    /// <summary>Listening.</summary>
    Listening = 2,

    /// <summary>Could not be started; <see cref="PortForwardViewModel.Detail"/> says why.</summary>
    Failed = 3,
}

/// <summary>
/// One of a host's saved forwards, as it runs on one terminal tab.
/// </summary>
/// <remarks>
/// <para>
/// Built fresh for each session the tab holds, and started with it: the rule is the host's, the listener is
/// this session's. A reconnect is a new session, so it gets new rows — see
/// <see cref="TerminalTabViewModel.ReplaceForwards"/> — rather than trying to carry a listener across a
/// connection that no longer exists.
/// </para>
/// <para>
/// <b>One failing does not stop the others, or the terminal.</b> The commonest failure is a port already in
/// use — very often by a second tab to the same host, which is running the same forwards — and a shell that
/// would not open because of it would be the forward taking the session hostage. The row says why instead.
/// </para>
/// </remarks>
internal sealed partial class PortForwardViewModel : ObservableObject
{
    private readonly Func<SshPortForwardSpec, CancellationToken, Task<ISshPortForward>> start;
    private readonly Action<Action> post;

    private ISshPortForward? forward;
    private bool sessionOver;

    /// <summary>The detail line: why it failed, or what last went wrong through it.</summary>
    [ObservableProperty]
    private string? detail;

    [ObservableProperty]
    [NotifyPropertyChangedFor(nameof(IsListening))]
    [NotifyPropertyChangedFor(nameof(IsFailed))]
    [NotifyPropertyChangedFor(nameof(CanToggle))]
    [NotifyPropertyChangedFor(nameof(StateText))]
    [NotifyCanExecuteChangedFor(nameof(ToggleCommand))]
    private PortForwardState state;

    /// <param name="rule">The host's saved forward.</param>
    /// <param name="start">Starts a forward on this tab's session; see <c>TerminalWorkspace.StartForwardAsync</c>.</param>
    /// <param name="post">Runs an action on the UI thread, for reports SSH.NET raises on its own.</param>
    internal PortForwardViewModel(
        PortForwardRule rule,
        Func<SshPortForwardSpec, CancellationToken, Task<ISshPortForward>> start,
        Action<Action> post)
    {
        Rule = rule;
        this.start = start;
        this.post = post;
    }

    /// <summary>The saved forward this row runs.</summary>
    internal PortForwardRule Rule { get; }

    /// <summary>The forward in OpenSSH's notation, e.g. <c>-L 5432:db.internal:5432</c>.</summary>
    internal string Notation => Rule.Notation;

    /// <summary>What the forward connects to what, in words rather than flags.</summary>
    /// <remarks>
    /// Which end each port is on is the thing people get backwards about <c>-L</c> and <c>-R</c>, so it is
    /// written out: "this machine" and "the remote" rather than an arrow whose direction has to be learned.
    /// </remarks>
    internal string Route => Rule.Kind switch
    {
        PortForwardKind.Local =>
            $"localhost:{Port(Rule.BindPort)} here → {Rule.TargetHost}:{Port(Rule.TargetPort)} from the remote",
        PortForwardKind.Remote =>
            $"localhost:{Port(Rule.BindPort)} on the remote → {Rule.TargetHost}:{Port(Rule.TargetPort)} from here",
        _ => $"SOCKS proxy on localhost:{Port(Rule.BindPort)}, out through the remote",
    };

    /// <summary>Whether it is listening now.</summary>
    internal bool IsListening => State is PortForwardState.Listening;

    /// <summary>Whether it could not be started.</summary>
    internal bool IsFailed => State is PortForwardState.Failed;

    /// <summary>
    /// Whether the switch does anything: not while a start is in flight, and never once the session is over.
    /// </summary>
    internal bool CanToggle => !sessionOver && State is not PortForwardState.Starting;

    /// <summary>The state in the status bar's own lettering.</summary>
    internal string StateText => State switch
    {
        PortForwardState.Starting => "STARTING",
        PortForwardState.Listening => "LISTENING",
        PortForwardState.Failed => "FAILED",
        _ => "OFF",
    };

    /// <summary>Starts the forward, unless it is already running or its session has ended.</summary>
    /// <remarks>
    /// Never throws. It is started without anyone waiting on it — every forward on a tab at once, as the
    /// tab opens — so a failure has nowhere to go but this row.
    /// </remarks>
    internal async Task StartAsync()
    {
        if (sessionOver || State is PortForwardState.Starting or PortForwardState.Listening)
        {
            return;
        }

        State = PortForwardState.Starting;
        Detail = null;

        ISshPortForward started;

        try
        {
            started = await start(ToSpec(Rule), CancellationToken.None).ConfigureAwait(true);
        }
        catch (Exception exception) when (exception is not OutOfMemoryException)
        {
            // Unless the session ended while this was in flight, in which case that is the true story and
            // SessionEnded has already told it.
            if (!sessionOver)
            {
                Detail = exception.Message;
                State = PortForwardState.Failed;
            }

            return;
        }

        if (sessionOver)
        {
            // The session ended while the server was answering. Its connection has gone and taken the
            // listener with it; this only lets go of the handle.
            await started.DisposeAsync().ConfigureAwait(true);
            return;
        }

        forward = started;
        started.ErrorOccurred += OnForwardError;
        State = PortForwardState.Listening;
    }

    /// <summary>Stops the forward and leaves the row switched off.</summary>
    internal async Task StopAsync()
    {
        await ReleaseAsync().ConfigureAwait(true);

        if (!sessionOver)
        {
            State = PortForwardState.Stopped;
            Detail = null;
        }
    }

    /// <summary>
    /// Records that the session this forward ran over has ended, which stopped the forward with it.
    /// </summary>
    /// <remarks>
    /// A failed row keeps its reason: that the session later ended does not make it any less true that the
    /// port was taken. Anything else is off, and says it went with the session rather than being switched
    /// off, which is the difference between "I did that" and "that happened".
    /// </remarks>
    internal void SessionEnded()
    {
        if (sessionOver)
        {
            return;
        }

        sessionOver = true;
        _ = ReleaseAsync();

        if (State is not PortForwardState.Failed)
        {
            State = PortForwardState.Stopped;
            Detail = "Stopped with the session.";
        }

        OnPropertyChanged(nameof(CanToggle));
        ToggleCommand.NotifyCanExecuteChanged();
    }

    /// <summary>The SSH layer's spec for a saved rule; the two enums share their numbering on purpose.</summary>
    internal static SshPortForwardSpec ToSpec(PortForwardRule rule) => rule.Kind switch
    {
        PortForwardKind.Local => new SshPortForwardSpec(
            SshPortForwardKind.Local, rule.BindPort, rule.TargetHost, rule.TargetPort ?? 0),
        PortForwardKind.Remote => new SshPortForwardSpec(
            SshPortForwardKind.Remote, rule.BindPort, rule.TargetHost, rule.TargetPort ?? 0),
        _ => new SshPortForwardSpec(SshPortForwardKind.Dynamic, rule.BindPort),
    };

    /// <remarks>A failed row's switch is a retry: the port may have been freed since.</remarks>
    [RelayCommand(CanExecute = nameof(CanToggle))]
    private Task ToggleAsync() => IsListening ? StopAsync() : StartAsync();

    private async Task ReleaseAsync()
    {
        if (forward is not { } running)
        {
            return;
        }

        forward = null;
        running.ErrorOccurred -= OnForwardError;

        try
        {
            await running.DisposeAsync().ConfigureAwait(true);
        }
        catch (Exception exception) when (exception is not OutOfMemoryException)
        {
            // Stopping a listener has no failure anybody could act on; the connection closing would have
            // stopped it regardless.
        }
    }

    /// <remarks>
    /// A connection through the forward failed — usually the target refusing it. The listener is still up,
    /// so the row stays LISTENING and says what happened underneath.
    /// </remarks>
    private void OnForwardError(object? sender, SshPortForwardErrorEventArgs e) => post(() =>
    {
        if (State is PortForwardState.Listening)
        {
            Detail = e.Message;
        }
    });

    private static string Port(int? port) => port?.ToString(CultureInfo.InvariantCulture) ?? "?";
}