From a1bcb9c7c249e5c9799718efd05a97545040f002 Mon Sep 17 00:00:00 2001 From: Robert Hague Date: Sun, 16 Nov 2025 20:20:22 +0100 Subject: [PATCH] Write into the underlying buffer in SshDataStream Several commonly used Write methods on SshDataStream end up calling Write(ReadOnlySpan) on the base MemoryStream. But since SshDataStream is a derived type, that method just rents a buffer and hands it to Write(byte[], int, int), which defeats any stackalloc'ing or renting that SshDataStream does itself. Instead, with a bit extra accounting we can just write directly into the underlying buffer. --- src/Renci.SshNet/Common/Extensions.cs | 8 +++ src/Renci.SshNet/Common/SshDataStream.cs | 92 ++++++++++++++++++------ 2 files changed, 79 insertions(+), 21 deletions(-) diff --git a/src/Renci.SshNet/Common/Extensions.cs b/src/Renci.SshNet/Common/Extensions.cs index 75d606fb7..11b8ec9a2 100644 --- a/src/Renci.SshNet/Common/Extensions.cs +++ b/src/Renci.SshNet/Common/Extensions.cs @@ -417,6 +417,14 @@ async Task WaitCore() return await completedTask.ConfigureAwait(false); } } + + extension(Array) + { + internal static int MaxLength + { + get { return 0X7FFFFFC7; } + } + } #endif } } diff --git a/src/Renci.SshNet/Common/SshDataStream.cs b/src/Renci.SshNet/Common/SshDataStream.cs index de0530bd0..291839c61 100644 --- a/src/Renci.SshNet/Common/SshDataStream.cs +++ b/src/Renci.SshNet/Common/SshDataStream.cs @@ -59,18 +59,49 @@ public bool IsEndOfData } } -#if !NET - private void Write(ReadOnlySpan buffer) + // Because this type derives from MemoryStream, the base Write(ReadOnlySpan) chooses + // to rent an array, copy the data in and delegate to Write(byte[], int, int) for + // backwards compatibility. + // With a bit of extra ceremony, we can instead allow the various Write methods here + // to write directly into the underlying buffer without the need for any intermediate + // arrays (rented or otherwise). + +#if NET9_0_OR_GREATER + /// + public override void Write(ReadOnlySpan buffer) { - var sharedBuffer = System.Buffers.ArrayPool.Shared.Rent(buffer.Length); + Write(buffer, buffer.Length, static (span, buffer) => buffer.CopyTo(span)); + } +#endif + + private delegate void WriteAction(Span span, TArg arg) +#if NET9_0_OR_GREATER + where TArg : allows ref struct +#endif + ; - buffer.CopyTo(sharedBuffer); + private void Write(TArg arg, int numBytesToWrite, WriteAction writeAction) +#if NET9_0_OR_GREATER + where TArg : allows ref struct +#endif + { + var endPosition = Position + numBytesToWrite; - Write(sharedBuffer, 0, buffer.Length); + if (Capacity < endPosition) + { + var newCapacity = Math.Max(endPosition, Math.Min(2 * (uint)Capacity, Array.MaxLength)); + Capacity = checked((int)newCapacity); + } - System.Buffers.ArrayPool.Shared.Return(sharedBuffer); + if (endPosition > Length) + { + SetLength(endPosition); + } + + writeAction(GetRemainingBuffer().AsSpan(0, numBytesToWrite), arg); + + Position = endPosition; } -#endif /// /// Writes an to the SSH data stream. @@ -78,9 +109,7 @@ private void Write(ReadOnlySpan buffer) /// data to write. public void Write(uint value) { - Span bytes = stackalloc byte[4]; - BinaryPrimitives.WriteUInt32BigEndian(bytes, value); - Write(bytes); + Write(value, 4, static (span, value) => BinaryPrimitives.WriteUInt32BigEndian(span, value)); } /// @@ -89,9 +118,7 @@ public void Write(uint value) /// data to write. public void Write(ulong value) { - Span bytes = stackalloc byte[8]; - BinaryPrimitives.WriteUInt64BigEndian(bytes, value); - Write(bytes); + Write(value, 8, static (span, value) => BinaryPrimitives.WriteUInt64BigEndian(span, value)); } /// @@ -100,9 +127,22 @@ public void Write(ulong value) /// The to write. public void Write(BigInteger data) { +#if NET + var byteCount = data.GetByteCount(); + + Write((data, byteCount), 4 + byteCount, static (span, args) => + { + BinaryPrimitives.WriteUInt32BigEndian(span, (uint)args.byteCount); + + var success = args.data.TryWriteBytes(span.Slice(4), out var bytesWritten, isBigEndian: true); + + Debug.Assert(success && bytesWritten == span.Length - 4); + }); +#else var bytes = data.ToByteArray(isBigEndian: true); WriteBinary(bytes, 0, bytes.Length); +#endif } /// @@ -129,16 +169,26 @@ public void Write(string s, Encoding encoding) ArgumentNullException.ThrowIfNull(s); ArgumentNullException.ThrowIfNull(encoding); + var byteCount = encoding.GetByteCount(s); #if NET - ReadOnlySpan value = s; - var count = encoding.GetByteCount(value); - var bytes = count <= 256 ? stackalloc byte[count] : new byte[count]; - encoding.GetBytes(value, bytes); - Write((uint)count); - Write(bytes); + Write((s, byteCount, encoding), 4 + byteCount, static (span, args) => + { + BinaryPrimitives.WriteUInt32BigEndian(span, (uint)args.byteCount); + + var bytesWritten = args.encoding.GetBytes(args.s, span.Slice(4)); + + Debug.Assert(bytesWritten == span.Length - 4); + }); #else - var bytes = encoding.GetBytes(s); - WriteBinary(bytes, 0, bytes.Length); + var rentedBuffer = System.Buffers.ArrayPool.Shared.Rent(byteCount); + + var bytesWritten = encoding.GetBytes(s, 0, s.Length, rentedBuffer, 0); + + Debug.Assert(bytesWritten == byteCount); + + WriteBinary(rentedBuffer, 0, bytesWritten); + + System.Buffers.ArrayPool.Shared.Return(rentedBuffer); #endif }