From 89dcb7f6c989907505f72f47f65528774e3ec492 Mon Sep 17 00:00:00 2001
From: Sanae <atmatm6000@outlook.com>
Date: Fri, 11 Mar 2022 19:08:50 -0600
Subject: [PATCH] More changes to reading varsize packets

---
 Server/Server.cs | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/Server/Server.cs b/Server/Server.cs
index eaff525..9b82415 100644
--- a/Server/Server.cs
+++ b/Server/Server.cs
@@ -120,7 +120,7 @@ public class Server {
             while (true) {
                 memory = memoryPool.Rent(Constants.HeaderSize);
 
-                async Task Read(Memory<byte> readMem, int readOffset = 0, int readSize = -1) {
+                async Task<bool> Read(Memory<byte> readMem, int readOffset = 0, int readSize = -1) {
                     if (readSize == -1) readSize = Constants.HeaderSize;
                     while (readOffset < readSize) {
                         int size = await socket.ReceiveAsync(readMem[readOffset..readSize], SocketFlags.None);
@@ -128,22 +128,26 @@ public class Server {
                             // treat it as a disconnect and exit
                             Logger.Info($"Socket {socket.RemoteEndPoint} disconnected.");
                             if (socket.Connected) await socket.DisconnectAsync(false);
-                            break;
+                            return false;
                         }
 
                         readOffset += size;
                     }
+
+                    return true;
                 }
 
-                await Read(memory.Memory[..Constants.HeaderSize]);
+                if (!await Read(memory.Memory[..Constants.HeaderSize])) break;
                 PacketHeader header = GetHeader(memory.Memory.Span[..Constants.HeaderSize]);
                 {
                     IMemoryOwner<byte> memTemp = memory;
-                    memory = memoryPool.Rent(Constants.HeaderSize + header.Size);
+                    memory = memoryPool.Rent(Constants.HeaderSize + header.PacketSize);
                     memTemp.Memory.CopyTo(memory.Memory);
                     memTemp.Dispose();
                 }
-                await Read(memory.Memory[Constants.HeaderSize..(Constants.HeaderSize + header.Size)]);
+                if (header.PacketSize > 0 
+                    && !await Read(memory.Memory[Constants.HeaderSize..(Constants.HeaderSize + header.PacketSize)]))
+                    break;
 
                 // connection initialization
                 if (first) {