Fix SOCKS reply ordering and rejected SOCKS4 responses

src/Renci.SshNet/Channels/ChannelDirectTcpip.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Net;
 using System.Net.Sockets;
 using System.Threading;
@@ -23,6 +24,34 @@ internal sealed class ChannelDirectTcpip : ClientChannel, IChannelDirectTcpip
         private IForwardedPort _forwardedPort;
         private Socket _socket;
 
+        /// <summary>
+        /// Holds a value indicating whether data received on the channel may be written to
+        /// <see cref="_socket"/>.
+        /// </summary>
+        /// <value>
+        /// <see langword="false"/> until <see cref="Bind()"/> releases the relay; see
+        /// <see cref="StartRelay"/> for why it is held back.
+        /// </value>
+        private bool _relaying;
+
+        /// <summary>
+        /// Holds the data received on the channel before the relay was released, or
+        /// <see langword="null"/> when there is none.
+        /// </summary>
+        private Queue<byte[]> _pendingData;
+
+        /// <summary>
+        /// Holds the shutdown that was requested before the relay was released, or
+        /// <see langword="null"/> when none was.
+        /// </summary>
+        private SocketShutdown? _pendingShutdown;
+
+        /// <summary>
+        /// Holds a value indicating whether <see cref="_socket"/> must be closed as soon as the relay
+        /// is released.
+        /// </summary>
+        private bool _pendingClose;
+
         /// <summary>
         /// Initializes a new instance of the <see cref="ChannelDirectTcpip"/> class.
         /// </summary>
@@ -47,6 +76,18 @@ public override ChannelTypes ChannelType
             get { return ChannelTypes.DirectTcpip; }
         }
 
+        /// <summary>
+        /// Opens the channel for the specified remote host and port.
+        /// </summary>
+        /// <param name="remoteHost">The name of the remote host to forward to.</param>
+        /// <param name="port">The port of the remote host to forward to.</param>
+        /// <param name="forwardedPort">The forwarded port for which the channel is opened.</param>
+        /// <param name="socket">The socket to receive requests from, and send responses from the remote host to.</param>
+        /// <remarks>
+        /// The channel takes ownership of <paramref name="socket"/> here, but does not yet write to it:
+        /// the caller may still have its own reply to send first. Data received before
+        /// <see cref="Bind()"/> is buffered rather than written. See <see cref="StartRelay"/>.
+        /// </remarks>
         public void Open(string remoteHost, uint port, IForwardedPort forwardedPort, Socket socket)
         {
             if (IsOpen)
@@ -80,29 +121,46 @@ public void Open(string remoteHost, uint port, IForwardedPort forwardedPort, Soc
         /// </summary>
         private void ForwardedPort_Closing(object sender, EventArgs eventArgs)
         {
-            // signal to the client that we will not send anything anymore; this should also interrupt the
-            // blocking receive in Bind if the client sends FIN/ACK in time
-            ShutdownSocket(SocketShutdown.Send);
+            // The port is going away, so this is an abort rather than an orderly close: act on the
+            // socket right away even when the relay has not been released, discarding anything that
+            // is still buffered. There is nobody left to relay it for, and the whole point of this
+            // handler is to interrupt a blocking receive.
+            lock (_socketLock)
+            {
+                // signal to the client that we will not send anything anymore; this should also interrupt the
+                // blocking receive in Bind if the client sends FIN/ACK in time
+                ShutdownSocketCore(SocketShutdown.Send);
 
-            // if the FIN/ACK is not sent in time by the remote client, then interrupt the blocking receive
-            // by closing the socket
-            CloseSocket();
+                // if the FIN/ACK is not sent in time by the remote client, then interrupt the blocking receive
+                // by closing the socket
+                CloseSocketCore();
+            }
         }
 
         /// <summary>
         /// Binds channel to remote host.
         /// </summary>
         public void Bind()
         {
+            // Release the data that arrived while the forwarded port was still writing its own reply
+            // to the client, and only then start pumping in the other direction.
+            StartRelay();
+
             // Cannot bind if channel is not open
             if (!IsOpen)
             {
                 return;
             }
 
+            var socket = _socket;
+            if (socket is null)
+            {
+                return;
+            }
+
             var buffer = new byte[RemotePacketSize];
 
-            SocketAbstraction.ReadContinuous(_socket, buffer, 0, buffer.Length, SendData);
+            SocketAbstraction.ReadContinuous(socket, buffer, 0, buffer.Length, SendData);
 
             // even though the client has disconnected, we still want to properly close the
             // channel
@@ -113,55 +171,163 @@ public void Bind()
         }
 
         /// <summary>
-        /// Closes the socket, hereby interrupting the blocking receive in <see cref="Bind()"/>.
+        /// Starts writing data received on the channel to the socket, first flushing whatever arrived
+        /// while the relay was held back.
         /// </summary>
-        private void CloseSocket()
+        /// <remarks>
+        /// <para>
+        /// A forwarded port writes its own reply to the client between <see cref="Open"/> and
+        /// <see cref="Bind()"/> - for <see cref="ForwardedPortDynamic"/> that is the SOCKS reply, which
+        /// RFC 1928 s6 requires to precede any byte from the target. <see cref="OnData"/> runs on the
+        /// session's message-receive thread, so without this gate the two would write to the same
+        /// socket from two threads with no ordering between them, and a target that speaks first - SSH,
+        /// SMTP, IMAP, POP3, FTP, MySQL, PostgreSQL, Redis - would land its greeting where the reply
+        /// belonged.
+        /// </para>
+        /// <para>
+        /// A shutdown or close requested while the relay was closed is applied here too, after the
+        /// buffered data, so that a target which speaks and then immediately hangs up still gets its
+        /// bytes to the client in the right order.
+        /// </para>
+        /// </remarks>
+        private void StartRelay()
         {
-            if (_socket is null)
+            lock (_socketLock)
             {
-                return;
+                if (_relaying)
+                {
+                    return;
+                }
+
+                _relaying = true;
+
+                var pending = _pendingData;
+                _pendingData = null;
+
+                if (pending is not null)
+                {
+                    while (pending.Count > 0)
+                    {
+                        var data = pending.Dequeue();
+                        SendToSocket(data, 0, data.Length);
+                    }
+                }
+
+                if (_pendingShutdown is SocketShutdown pendingShutdown)
+                {
+                    _pendingShutdown = null;
+                    ShutdownSocketCore(pendingShutdown);
+                }
+
+                if (_pendingClose)
+                {
+                    _pendingClose = false;
+                    CloseSocketCore();
+                }
             }
+        }
+
+        /// <summary>
+        /// Writes data to the socket, if there still is one and it is connected.
+        /// </summary>
+        /// <param name="data">An array of <see cref="byte"/> containing the data to write.</param>
+        /// <param name="offset">The zero-based offset in <paramref name="data"/> at which to begin taking data from.</param>
+        /// <param name="count">The number of bytes of <paramref name="data"/> to write.</param>
+        /// <remarks>
+        /// The caller must hold <see cref="_socketLock"/>.
+        /// </remarks>
+        private void SendToSocket(byte[] data, int offset, int count)
+        {
+            var socket = _socket;
 
+            if (socket.IsConnected())
+            {
+                SocketAbstraction.Send(socket, data, offset, count);
+            }
+        }
+
+        /// <summary>
+        /// Closes the socket, hereby interrupting the blocking receive in <see cref="Bind()"/>.
+        /// </summary>
+        private void CloseSocket()
+        {
             lock (_socketLock)
             {
-                if (_socket is null)
+                if (!_relaying)
                 {
+                    // Disposing the socket now would strand both the reply the forwarded port is about
+                    // to write and anything OnData has buffered. StartRelay applies this afterwards.
+                    _pendingClose = true;
                     return;
                 }
 
-                // closing a socket actually disposes the socket, so we can safely dereference
-                // the field to avoid entering the lock again later
-                _socket.Dispose();
-                _socket = null;
+                CloseSocketCore();
             }
         }
 
+        /// <summary>
+        /// Disposes the socket.
+        /// </summary>
+        /// <remarks>
+        /// The caller must hold <see cref="_socketLock"/>.
+        /// </remarks>
+        private void CloseSocketCore()
+        {
+            // closing a socket actually disposes the socket, so we can safely dereference
+            // the field to avoid entering the lock again later
+            _socket?.Dispose();
+            _socket = null;
+
+            // there is no longer anywhere to write buffered data to
+            _pendingData = null;
+        }
+
         /// <summary>
         /// Shuts down the socket.
         /// </summary>
         /// <param name="how">One of the <see cref="SocketShutdown"/> values that specifies the operation that will no longer be allowed.</param>
         private void ShutdownSocket(SocketShutdown how)
         {
-            if (_socket is null)
-            {
-                return;
-            }
-
             lock (_socketLock)
             {
-                if (!_socket.IsConnected())
+                if (!_relaying)
                 {
+                    // Shutting the send side down now would break the reply the forwarded port is
+                    // about to write, and would discard whatever OnData has buffered. StartRelay
+                    // applies this after the buffer has been drained.
+                    _pendingShutdown = _pendingShutdown is null || _pendingShutdown == how
+                        ? how
+                        : SocketShutdown.Both;
                     return;
                 }
 
-                try
-                {
-                    _socket.Shutdown(how);
-                }
-                catch (SocketException ex)
-                {
-                    _logger.LogInformation(ex, "Failure shutting down socket");
-                }
+                ShutdownSocketCore(how);
+            }
+        }
+
+        /// <summary>
+        /// Shuts down the socket, without regard for the relay.
+        /// </summary>
+        /// <param name="how">One of the <see cref="SocketShutdown"/> values that specifies the operation that will no longer be allowed.</param>
+        /// <remarks>
+        /// The caller must hold <see cref="_socketLock"/>.
+        /// </remarks>
+        private void ShutdownSocketCore(SocketShutdown how)
+        {
+            var socket = _socket;
+
+            if (!socket.IsConnected())
+            {
+                return;
+            }
+
+            try
+            {
+                socket.Shutdown(how);
+            }
+            catch (SocketException ex)
+            {
+                _logger.LogInformation(ex, "Failure shutting down socket");
             }
         }
 
@@ -198,15 +364,28 @@ protected override void OnData(ArraySegment<byte> data)
         {
             base.OnData(data);
 
-            if (_socket != null)
+            lock (_socketLock)
             {
-                lock (_socketLock)
+                if (_socket is null)
                 {
-                    if (_socket.IsConnected())
-                    {
-                        SocketAbstraction.Send(_socket, data.Array, data.Offset, data.Count);
-                    }
+                    return;
+                }
+
+                if (!_relaying)
+                {
+                    // We are on the session's message-receive thread and the forwarded port has not
+                    // finished writing its own reply to the client yet. Hold on to the data instead of
+                    // racing that thread for the socket; StartRelay writes it out in order. The array
+                    // belongs to the message being processed, so it has to be copied.
+                    var buffered = new byte[data.Count];
+                    Buffer.BlockCopy(data.Array, data.Offset, buffered, 0, data.Count);
+
+                    _pendingData ??= new Queue<byte[]>();
+                    _pendingData.Enqueue(buffered);
+                    return;
                 }
+
+                SendToSocket(data.Array, data.Offset, data.Count);
             }
         }
 
@@ -288,17 +467,9 @@ protected override void Dispose(bool disposing)
 
             if (disposing)
             {
-                if (_socket != null)
+                lock (_socketLock)
                 {
-                    lock (_socketLock)
-                    {
-                        var socket = _socket;
-                        if (socket != null)
-                        {
-                            _socket = null;
-                            socket.Dispose();
-                        }
-                    }
+                    CloseSocketCore();
                 }
 
                 var channelOpen = _channelOpen;