Add a loop to SocketAbstraction.ReadAsync

In order to ensure the buffer is read completely, as in SocketAbstraction.Read
sshnet · Mar 11, 2024 · 3f6accb · 3f6accb
1 parent 75ced08
commit 3f6accb
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 24 deletions.
diff --git a/src/Renci.SshNet/Abstractions/SocketAbstraction.Async.cs b/src/Renci.SshNet/Abstractions/SocketAbstraction.Async.cs
diff --git a/src/Renci.SshNet/Abstractions/SocketAbstraction.cs b/src/Renci.SshNet/Abstractions/SocketAbstraction.cs
@@ -3,9 +3,7 @@
 using System.Net;
 using System.Net.Sockets;
 using System.Threading;
-#if NET6_0_OR_GREATER == false
 using System.Threading.Tasks;
-#endif
 
 using Renci.SshNet.Common;
 
@@ -151,11 +149,34 @@ public static int Read(Socket socket, byte[] buffer, int offset, int size, TimeS
             return totalBytesRead;
         }
 
-#if NET6_0_OR_GREATER == false
-        public static ValueTask<int> ReadAsync(Socket socket, byte[] buffer, CancellationToken cancellationToken)
+        public static async ValueTask<int> ReadAsync(Socket socket, byte[] buffer, int offset, int size, CancellationToken cancellationToken)
         {
-            return socket.ReceiveAsync(new ArraySegment<byte>(buffer, 0, buffer.Length), SocketFlags.None, cancellationToken);
+            var totalBytesRead = 0;
+            var totalBytesToRead = size;
+
+            do
+            {
+                try
+                {
+                    var bytesRead = await socket.ReceiveAsync(new ArraySegment<byte>(buffer, offset + totalBytesRead, totalBytesToRead - totalBytesRead), SocketFlags.None, cancellationToken).ConfigureAwait(false);
+                    if (bytesRead == 0)
+                    {
+                        return 0;
+                    }
+
+                    totalBytesRead += bytesRead;
+                }
+                catch (SocketException ex) when (ex.SocketErrorCode == SocketError.TimedOut)
+                {
+                    throw new SshOperationTimeoutException(string.Format(CultureInfo.InvariantCulture,
+                                                            "Socket read operation has timed out after {0:F0} milliseconds.",
+                                                            socket.ReceiveTimeout),
+                                                            ex);
+                }
+            }
+            while (totalBytesRead < totalBytesToRead);
+
+            return totalBytesRead;
         }
-#endif
     }
 }
diff --git a/src/Renci.SshNet/Connection/ProtocolVersionExchange.cs b/src/Renci.SshNet/Connection/ProtocolVersionExchange.cs
@@ -187,7 +187,7 @@ private static async Task<string> SocketReadLineAsync(Socket socket, List<byte>
             // to be processed by subsequent invocations.
             while (true)
             {
-                var bytesRead = await SocketAbstraction.ReadAsync(socket, data, cancellationToken).ConfigureAwait(false);
+                var bytesRead = await SocketAbstraction.ReadAsync(socket, data, 0, data.Length, cancellationToken).ConfigureAwait(false);
                 if (bytesRead == 0)
                 {
                     throw new SshConnectionException("The connection was closed by the remote host.");