PipeInputStream will read as much as it can without blocking

TrapTech · Nov 24, 2022 · b430a83 · b430a83
1 parent 796d7b0
commit b430a83
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 12 deletions.
diff --git a/src/Renci.SshNet/Common/LinkedListQueue.cs b/src/Renci.SshNet/Common/LinkedListQueue.cs
@@ -93,10 +93,17 @@ public void CompleteAdding()
         /// </summary>
         /// <returns><c>true</c>, if an item could be removed; otherwise <c>false</c>.</returns>
         /// <param name="item">The item to be removed from the queue.</param>
-        public bool TryTake(out T item)
+        /// <param name="wait">Wait for data or fail immediately if empty.</param>
+        public bool TryTake(out T item, bool wait)
         {
             lock (_lock)
             {
+                if (_first == null && !wait)
+                {
+                    item = default(T);
+                    return false;
+                }
+
                 while (_first == null && !_isAddingCompleted)
                     Monitor.Wait(_lock);
 

diff --git a/src/Renci.SshNet/Common/PipeInputStream.cs b/src/Renci.SshNet/Common/PipeInputStream.cs
@@ -40,22 +40,32 @@ public override int Read(byte[] buffer, int offset, int count)
             if (_isDisposed)
                 throw CreateObjectDisposedException();
 
-            if (_current == null || _currentPosition == _current.Length)
+            var bytesRead = 0;
+
+            while (bytesRead < count)
             {
-                if (_queue.IsCompleted || !_queue.TryTake(out _current))
-                    return 0;
+                if (_current == null || _currentPosition == _current.Length)
+                {
+                    if (!_queue.TryTake(out _current, (bytesRead == 0)))
+                    {
+                        _current = null;
+                        return bytesRead;
+                    }
 
-                _currentPosition = 0;
-            }
+                    _currentPosition = 0;
+                }
 
-            var avail = _current.Length - _currentPosition;
-            if (count > avail)
-                count = avail;
+                var toRead = _current.Length - _currentPosition;
+                if (toRead > count - bytesRead)
+                    toRead = count - bytesRead;
 
-            Buffer.BlockCopy(_current, _currentPosition, buffer, offset, count);
+                Buffer.BlockCopy(_current, _currentPosition, buffer, offset + bytesRead, toRead);
+
+                _currentPosition += toRead;
+                bytesRead += toRead;
+            }
 
-            _currentPosition += count;
-            return count;
+            return bytesRead;
         }
 
         public override void Write(byte[] buffer, int offset, int count)