StreamExtensions: Add "ReadExactlyAsync" method and use it in TorSocks5Client.
MODIFY
WalletWasabi/Extensions/StreamExtensions.cs
WalletWasabi/Extensions/StreamExtensions.cs
+using System.Buffers;
using System.Threading;
using System.Threading.Tasks;
-using System;
namespace System.IO
{
@@ -8,25 +8,59 @@ namespace System.IO
{
public static async Task<int> ReadByteAsync(this Stream stream, CancellationToken ctsToken = default)
{
- var buf = new byte[1];
- var len = await stream.ReadAsync(buf.AsMemory(0, 1), ctsToken).ConfigureAwait(false);
+ ArrayPool<byte> pool = ArrayPool<byte>.Shared;
+ byte[] buffer = pool.Rent(1);
+ try
+ {
+ int len = await stream.ReadAsync(buffer.AsMemory(0, 1), ctsToken).ConfigureAwait(false);
+
+ // End of stream.
if (len == 0)
{
return -1;
}
- return buf[0];
+ return buffer[0];
+ }
+ finally
+ {
+ pool.Return(buffer);
+ }
}
public static async Task<int> ReadBlockAsync(this Stream stream, byte[] buffer, int count, CancellationToken ctsToken = default)
{
- var left = count;
+ int left = count;
while (left != 0)
{
- var read = await stream.ReadAsync(buffer.AsMemory(count - left, left), ctsToken).ConfigureAwait(false);
+ int read = await stream.ReadAsync(buffer.AsMemory(count - left, left), ctsToken).ConfigureAwait(false);
left -= read;
}
return count - left;
}
+
+ /// <summary>
+ /// Reads exactly <paramref name="count"/> bytes from <paramref name="stream"/>.
+ /// </summary>
+ /// <param name="stream">Stream to read from.</param>
+ /// <param name="buffer">Buffer whose length must be at least <paramref name="count"/> elements.</param>
+ /// <param name="count">Number of bytes to read.</param>
+ /// <param name="cancellationToken">Cancellation token to cancel the asynchronous operation.</param>
+ /// <returns><c>true</c> if we could read exactly <paramref name="count"/> bytes from stream (stream may contain more bytes though).</returns>
+ public static async Task<bool> ReadExactlyAsync(this Stream stream, byte[] buffer, int count, CancellationToken cancellationToken = default)
+ {
+ int offset = 0;
+ while (offset < count)
+ {
+ int read = await stream.ReadAsync(buffer.AsMemory(offset, count - offset), cancellationToken).ConfigureAwait(false);
+ if (read == 0)
+ {
+ return false;
+ }
+ offset += read;
+ }
+
+ return true;
+ }
}
}

MODIFY
WalletWasabi/Tor/Socks5/TorSocks5Client.cs
WalletWasabi/Tor/Socks5/TorSocks5Client.cs
@@ -127,7 +127,7 @@ namespace WalletWasabi.Tor.Socks5
var methods = new MethodsField(isolateStream ? MethodField.UsernamePassword : MethodField.NoAuthenticationRequired);
byte[] sendBuffer = new VersionMethodRequest(methods).ToBytes();
- byte[] receiveBuffer = await SendAsync(sendBuffer, receiveBufferSize: 2, cancellationToken).ConfigureAwait(false);
+ byte[] receiveBuffer = await SendAndReceiveAsync(sendBuffer, receiveBufferSize: 2, cancellationToken).ConfigureAwait(false);
var methodSelection = new MethodSelectionResponse(receiveBuffer);
@@ -157,7 +157,7 @@ namespace WalletWasabi.Tor.Socks5
sendBuffer = usernamePasswordRequest.ToBytes();
Array.Clear(receiveBuffer, 0, receiveBuffer.Length);
- receiveBuffer = await SendAsync(sendBuffer, receiveBufferSize: 2, cancellationToken).ConfigureAwait(false);
+ receiveBuffer = await SendAndReceiveAsync(sendBuffer, receiveBufferSize: 2, cancellationToken).ConfigureAwait(false);
var userNamePasswordResponse = new UsernamePasswordResponse(receiveBuffer);
@@ -225,7 +225,7 @@ namespace WalletWasabi.Tor.Socks5
var connectionRequest = new TorSocks5Request(cmd: CmdField.Connect, new AddrField(host), new PortField(port));
var sendBuffer = connectionRequest.ToBytes();
- var receiveBuffer = await SendAsync(sendBuffer, receiveBufferSize: null, cancellationToken).ConfigureAwait(false);
+ var receiveBuffer = await SendAndReceiveAsync(sendBuffer, receiveBufferSize: null, cancellationToken).ConfigureAwait(false);
var connectionResponse = new TorSocks5Response(receiveBuffer);
@@ -331,13 +331,14 @@ namespace WalletWasabi.Tor.Socks5
}
/// <summary>
- /// Sends bytes to the Tor Socks5 connection
+ /// Sends a command to the Tor Socks5 connection and reads a response.
/// </summary>
/// <param name="sendBuffer">Sent data</param>
- /// <param name="receiveBufferSize">Maximum number of bytes expected to be received in the reply</param>
+ /// <param name="receiveBufferSize">Optionally, number of bytes expected to be received in the reply.</param>
/// <param name="cancellationToken">Cancellation token to cancel sending.</param>
/// <returns>Reply</returns>
- private async Task<byte[]> SendAsync(byte[] sendBuffer, int? receiveBufferSize = null, CancellationToken cancellationToken = default)
+ /// <exception cref="TorConnectionException">When we receive no response from Tor or the response is invalid.</exception>
+ private async Task<byte[]> SendAndReceiveAsync(byte[] sendBuffer, int? receiveBufferSize = null, CancellationToken cancellationToken = default)
{
Guard.NotNullOrEmpty(nameof(sendBuffer), sendBuffer);
@@ -363,6 +364,17 @@ namespace WalletWasabi.Tor.Socks5
// Receive the response
var receiveBuffer = new byte[actualReceiveBufferSize];
+ // Read exactly "receiveBufferSize" bytes.
+ if (receiveBufferSize != null)
+ {
+ if (await stream.ReadExactlyAsync(receiveBuffer, receiveBufferSize.Value, cancellationToken).ConfigureAwait(false))
+ {
+ return receiveBuffer;
+ }
+
+ throw new TorConnectionException($"Failed to read {receiveBufferSize.Value} bytes as expected from Tor SOCKS5.");
+ }
+
int receiveCount = await stream.ReadAsync(receiveBuffer.AsMemory(0, actualReceiveBufferSize), cancellationToken).ConfigureAwait(false);
if (receiveCount <= 0)

StreamExtensions: Add "ReadExactlyAsync" method and use it in TorSocks5Client.