From 0b1404d8ccace7a89c7b1756cb70a577e5231c94 Mon Sep 17 00:00:00 2001
From: SALTWOOD <105980161+SALTWOOD@users.noreply.github.com>
Date: Tue, 13 Aug 2024 15:48:30 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E5=9C=A8=E8=BF=9B=E8=A1=8C=20Read=20?=
 =?UTF-8?q?=E6=97=B6=E6=84=8F=E5=A4=96=E7=9A=84=E9=98=BB=E5=A1=9E?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 Network/RsaStream.cs | 33 ++++++++++++++++++++++++++++++---
 1 file changed, 30 insertions(+), 3 deletions(-)

diff --git a/Network/RsaStream.cs b/Network/RsaStream.cs
index 66856c5..0591632 100644
--- a/Network/RsaStream.cs
+++ b/Network/RsaStream.cs
@@ -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
     {
@@ -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)
@@ -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.");
@@ -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)
                     {
@@ -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;
             }