Skip to content

Commit

Permalink
fix: 在进行 Read 时意外的阻塞
Browse files Browse the repository at this point in the history
  • Loading branch information
SALTWOOD committed Aug 13, 2024
1 parent ae18649 commit 0b1404d
Showing 1 changed file with 30 additions and 3 deletions.
33 changes: 30 additions & 3 deletions Network/RsaStream.cs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ public class RsaStream : Stream
protected object _lock = new object();
protected object _statusLock = new object();
private RsaStreamStatus _status;
private byte[] pendingBytes = Array.Empty<byte>();

public RsaStreamStatus Status
{
Expand Down Expand Up @@ -111,12 +112,14 @@ public override void Write(byte[] buffer, int offset, int count)
{
lock (_lock)
{

if (buffer == null) throw new ArgumentNullException(nameof(buffer));

// The max block size depends on the key size and padding. With OAEP and SHA-256, it's approximately: KeySize/8 - 2*HashSize - 2
int maxBlockSize = _remotePublicKey!.KeySize / 8 - 2 * 32 - 2;

long messageLength = count;
_stream.Write(BitConverter.GetBytes(messageLength));

using (var cryptoStream = new MemoryStream())
{
for (int i = offset; i < offset + count; i += maxBlockSize)
Expand All @@ -140,6 +143,12 @@ public override int Read(byte[] buffer, int offset, int count)
{
lock (_lock)
{
if (this.pendingBytes != null && this.pendingBytes.Length != 0)
{
(byte[] temp, this.pendingBytes) = (this.pendingBytes[0..count], this.pendingBytes[count..]);
Array.Copy(temp, 0, buffer, offset, count);
}

if (buffer == null) throw new ArgumentNullException(nameof(buffer));
if (offset < 0 || count < 0) throw new ArgumentOutOfRangeException("Offset and count must be non-negative.");
if (buffer.Length - offset < count) throw new ArgumentException("Invalid offset and count relative to buffer length.");
Expand All @@ -152,12 +161,20 @@ public override int Read(byte[] buffer, int offset, int count)
byte[] encryptedBuffer = new byte[encryptedBlockSize];
int bytesRead;

byte[] decryptedBlock = Array.Empty<byte>();

byte[] messageLengthByte = new byte[8];
if (_stream.Read(messageLengthByte) < 8) throw new Exception("Unable to read message length.");
long messageLength = BitConverter.ToInt64(messageLengthByte);
byte[] restBytes = new byte[Math.Max(0, messageLength - count)];
long actualSize = Math.Min(messageLength, count);

while ((bytesRead = _stream.Read(encryptedBuffer, 0, encryptedBlockSize)) > 0)
{
if (bytesRead != encryptedBlockSize)
throw new CryptographicException("The length of the data to decrypt is not valid for the size of this key.");

byte[] decryptedBlock = _rsaPrivate.Decrypt(encryptedBuffer, RSAEncryptionPadding.OaepSHA256);
decryptedBlock = _rsaPrivate.Decrypt(encryptedBuffer, RSAEncryptionPadding.OaepSHA256);

if (decryptedBlock.Length > 0)
{
Expand All @@ -166,9 +183,19 @@ public override int Read(byte[] buffer, int offset, int count)
totalBytesRead += copySize;
}

if (totalBytesRead >= count)
if (totalBytesRead >= actualSize)
break;
}
if (totalBytesRead > actualSize)
{
this.pendingBytes = new byte[totalBytesRead - actualSize];
Array.Copy(decryptedBlock[0..(int)(totalBytesRead - actualSize)], this.pendingBytes, (totalBytesRead - actualSize));
}
if (totalBytesRead < actualSize &&
restBytes.Length != 0 &&
_stream.Read(restBytes, 0, restBytes.Length) != restBytes.Length)
throw new Exception("Message length doesn't match.");
this.pendingBytes = restBytes;

return totalBytesRead;
}
Expand Down

0 comments on commit 0b1404d

Please sign in to comment.