DodoSSH/tests/DodoSSH.Client.Ssh.Tests/PortForwardTests.cs

using System.Net;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Text;

namespace DodoSSH.Client.Ssh.Tests;

/// <summary>
/// The three kinds of forward, each carrying real bytes through the container's sshd.
/// </summary>
/// <remarks>
/// <para>
/// ◆ <b>The target of a local forward is named as the container sees it</b>, for the reason
/// <c>LoopbackProxyTests</c> gives: the remote resolves it, so <c>localhost:2222</c> is the container's own
/// sshd, and the published port this test host uses means nothing in there. What comes back is that sshd's
/// banner, which is bytes nothing but a real forward could have fetched.
/// </para>
/// <para>
/// A remote forward is proved the other way round. A listener on this test host answers with a known line,
/// the remote is asked to listen for it, and a local forward on the same connection reaches that remote
/// port from the inside — so the line has crossed the connection twice, once in each direction.
/// </para>
/// </remarks>
[Collection(SshCollection.Name)]
public sealed class PortForwardTests(SshServerFixture fixture)
{
    /// <summary>What the container calls itself. See the class remarks.</summary>
    private const string InternalHost = "localhost";

    private static CancellationToken Token => TestContext.Current.CancellationToken;

    [Fact]
    public async Task ALocalForward_CarriesAConnectionToATargetOnlyTheRemoteCanSee()
    {
        await using var connection = await ConnectAsync();

        await using var forward = await connection.StartForwardAsync(
            new SshPortForwardSpec(SshPortForwardKind.Local, 0, InternalHost, SshServerFixture.InternalPort),
            Token);

        forward.IsListening.ShouldBeTrue();
        forward.BoundPort.ShouldBeGreaterThan(0, "an ephemeral port is reported back once bound");

        using var client = new TcpClient();
        await client.ConnectAsync(IPAddress.Loopback, forward.BoundPort, Token);

        (await ReadLineAsync(client.GetStream())).ShouldStartWith("SSH-2.0-");
    }

    [Fact]
    public async Task ADynamicForward_IsASocksProxyWhoseConnectionsLeaveFromTheRemote()
    {
        await using var connection = await ConnectAsync();

        await using var forward = await connection.StartForwardAsync(
            new SshPortForwardSpec(SshPortForwardKind.Dynamic, 0), Token);

        using var client = new TcpClient();
        await client.ConnectAsync(IPAddress.Loopback, forward.BoundPort, Token);

        var stream = client.GetStream();
        await Socks5ConnectAsync(stream, InternalHost, SshServerFixture.InternalPort);

        (await ReadLineAsync(stream)).ShouldStartWith("SSH-2.0-");
    }

    [Fact]
    public async Task ARemoteForward_CarriesAConnectionOnTheRemoteBackToThisMachine()
    {
        const string Greeting = "hello from the client side";

        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();

        var served = ServeOneLineAsync(listener, Greeting);

        await using var connection = await ConnectAsync();

        var remotePort = RandomNumberGenerator.GetInt32(20000, 60000);

        await using var remote = await connection.StartForwardAsync(
            new SshPortForwardSpec(
                SshPortForwardKind.Remote,
                remotePort,
                "127.0.0.1",
                ((IPEndPoint)listener.LocalEndpoint).Port),
            Token);

        remote.BoundPort.ShouldBe(remotePort);

        // Reaches the remote's own loopback, which is where the remote forward is listening.
        await using var inward = await connection.StartForwardAsync(
            new SshPortForwardSpec(SshPortForwardKind.Local, 0, "127.0.0.1", remotePort),
            Token);

        using var client = new TcpClient();
        await client.ConnectAsync(IPAddress.Loopback, inward.BoundPort, Token);

        (await ReadLineAsync(client.GetStream())).ShouldBe(Greeting);

        await served;
    }

    [Fact]
    public async Task APortAlreadyInUse_IsRefusedWithASentenceNamingIt()
    {
        using var squatter = new TcpListener(IPAddress.Loopback, 0);
        squatter.Start();

        var taken = ((IPEndPoint)squatter.LocalEndpoint).Port;

        await using var connection = await ConnectAsync();

        var refused = await Should.ThrowAsync<SshPortForwardException>(async () =>
            await connection.StartForwardAsync(
                new SshPortForwardSpec(SshPortForwardKind.Local, taken, InternalHost, SshServerFixture.InternalPort),
                Token));

        refused.Message.ShouldBe($"Port {taken} on this machine is already in use.");

        // And the refusal left nothing half-registered on the client: the next forward still starts.
        await using var next = await connection.StartForwardAsync(
            new SshPortForwardSpec(SshPortForwardKind.Local, 0, InternalHost, SshServerFixture.InternalPort),
            Token);

        next.IsListening.ShouldBeTrue();
    }

    [Fact]
    public async Task DisposingAForward_StopsItsListener()
    {
        await using var connection = await ConnectAsync();

        var forward = await connection.StartForwardAsync(
            new SshPortForwardSpec(SshPortForwardKind.Local, 0, InternalHost, SshServerFixture.InternalPort),
            Token);

        var port = forward.BoundPort;

        await forward.DisposeAsync();

        forward.IsListening.ShouldBeFalse();
        await ShouldRefuseAsync(port);

        // Idempotent, because the connection closing will try again.
        await forward.DisposeAsync();
    }

    [Fact]
    public async Task ClosingTheConnection_StopsEveryForwardOnIt()
    {
        var connection = await ConnectAsync();

        var forward = await connection.StartForwardAsync(
            new SshPortForwardSpec(SshPortForwardKind.Dynamic, 0), Token);

        var port = forward.BoundPort;

        await connection.DisposeAsync();

        forward.IsListening.ShouldBeFalse();
        await ShouldRefuseAsync(port);

        // What the workspace does after closing a session, and it must be harmless.
        await forward.DisposeAsync();
    }

    [Fact]
    public async Task AnInvalidSpec_IsRefusedBeforeAnythingListens()
    {
        await using var connection = await ConnectAsync();

        await Should.ThrowAsync<ArgumentException>(async () =>
            await connection.StartForwardAsync(
                new SshPortForwardSpec(SshPortForwardKind.Local, 0, TargetHost: null, TargetPort: 0),
                Token));

        await Should.ThrowAsync<ArgumentException>(async () =>
            await connection.StartForwardAsync(
                new SshPortForwardSpec(SshPortForwardKind.Remote, 0, "127.0.0.1", 80),
                Token));
    }

    private static async Task ShouldRefuseAsync(int port)
    {
        using var client = new TcpClient();

        await Should.ThrowAsync<SocketException>(async () =>
            await client.ConnectAsync(IPAddress.Loopback, port, Token));
    }

    private static async Task ServeOneLineAsync(TcpListener listener, string line)
    {
        using var accepted = await listener.AcceptTcpClientAsync(Token);

        await accepted.GetStream().WriteAsync(Encoding.ASCII.GetBytes(line + "\n"), Token);
    }

    /// <summary>
    /// The smallest SOCKS5 exchange there is: no authentication, then CONNECT by name.
    /// </summary>
    private static async Task Socks5ConnectAsync(NetworkStream stream, string host, int port)
    {
        await stream.WriteAsync(new byte[] { 5, 1, 0 }, Token);

        var method = new byte[2];
        await stream.ReadExactlyAsync(method, Token);
        method.ShouldBe(new byte[] { 5, 0 });

        var name = Encoding.ASCII.GetBytes(host);
        byte[] request = [5, 1, 0, 3, (byte)name.Length, .. name, (byte)(port >> 8), (byte)port];
        await stream.WriteAsync(request, Token);

        // An IPv4 bound address, which is what SSH.NET always answers with: ten bytes.
        var reply = new byte[10];
        await stream.ReadExactlyAsync(reply, Token);
        reply[1].ShouldBe((byte)0, "the proxy should report the CONNECT as granted");
    }

    private static async Task<string> ReadLineAsync(NetworkStream stream)
    {
        var line = new List<byte>();
        var one = new byte[1];

        using var timeout = CancellationTokenSource.CreateLinkedTokenSource(Token);
        timeout.CancelAfter(TimeSpan.FromSeconds(15));

        while (await stream.ReadAsync(one, timeout.Token) == 1 && one[0] != (byte)'\n')
        {
            line.Add(one[0]);
        }

        return Encoding.ASCII.GetString([.. line]).TrimEnd('\r');
    }

    private SshConnectionRequest Request() =>
        new(fixture.Host, fixture.Port, SshServerFixture.Username, new SshPasswordCredential(SshServerFixture.Password));

    /// <summary>Connects with the container's host key already trusted, since first contact is not the subject.</summary>
    private async Task<ISshConnection> ConnectAsync()
    {
        var knownHosts = new InMemoryKnownHostStore();
        var factory = new SshNetConnectionFactory(knownHosts);

        var unknown = await Should.ThrowAsync<SshHostKeyUnknownException>(async () =>
            await factory.ConnectAsync(Request(), progress: null, Token));

        await knownHosts.TrustAsync(unknown.Presentation, Token);

        return await factory.ConnectAsync(Request(), progress: null, Token);
    }
}