From 104e8447a4944295ab3c10271de456f536fda92c Mon Sep 17 00:00:00 2001
From: Badrish Chandramouli <badrishc@microsoft.com>
Date: Wed, 15 Dec 2021 13:17:40 -0800
Subject: [PATCH] Fix object serialization boundary condition (#619)

---
 cs/src/core/Allocator/GenericAllocator.cs | 28 +++++++++++------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/cs/src/core/Allocator/GenericAllocator.cs b/cs/src/core/Allocator/GenericAllocator.cs
index c00d2811d..fe6f074ea 100644
--- a/cs/src/core/Allocator/GenericAllocator.cs
+++ b/cs/src/core/Allocator/GenericAllocator.cs
@@ -413,6 +413,7 @@ private void WriteAsync<TContext>(long flushPage, ulong alignedDestinationAddres
 
             for (int i=start/recordSize; i<end/recordSize; i++)
             {
+                long endPosition = 0;
                 if (!src[i].info.Invalid)
                 {
                     if (KeyHasObjects())
@@ -423,6 +424,7 @@ private void WriteAsync<TContext>(long flushPage, ulong alignedDestinationAddres
                         key_address->Address = pos;
                         key_address->Size = (int)(ms.Position - pos);
                         addr.Add((long)key_address);
+                        endPosition = pos + key_address->Size;
                     }
 
                     if (ValueHasObjects() && !src[i].info.Tombstone)
@@ -433,10 +435,11 @@ private void WriteAsync<TContext>(long flushPage, ulong alignedDestinationAddres
                         value_address->Address = pos;
                         value_address->Size = (int)(ms.Position - pos);
                         addr.Add((long)value_address);
+                        endPosition = pos + value_address->Size;
                     }
                 }
 
-                if (ms.Position > ObjectBlockSize || i == (end / recordSize) - 1)
+                if (endPosition > ObjectBlockSize || i == (end / recordSize) - 1)
                 {
                     var memoryStreamLength = (int)ms.Position;
 
@@ -835,8 +838,9 @@ public void GetObjectInfo(byte* raw, ref long ptr, long untilptr, int objectBloc
         {
             long minObjAddress = long.MaxValue;
             long maxObjAddress = long.MinValue;
+            bool done = false;
 
-            while (ptr < untilptr)
+            while (!done && (ptr < untilptr))
             {
                 ref Record<Key, Value> record = ref Unsafe.AsRef<Record<Key, Value>>(raw + ptr);
 
@@ -847,15 +851,13 @@ public void GetObjectInfo(byte* raw, ref long ptr, long untilptr, int objectBloc
                         var key_addr = GetKeyAddressInfo((long)raw + ptr);
                         var addr = key_addr->Address;
 
-                        // If object pointer is greater than kObjectSize from starting object pointer
-                        if (minObjAddress != long.MaxValue && (addr - minObjAddress > objectBlockSize))
-                        {
-                            break;
-                        }
-
                         if (addr < minObjAddress) minObjAddress = addr;
                         addr += key_addr->Size;
                         if (addr > maxObjAddress) maxObjAddress = addr;
+
+                        // If object pointer is greater than kObjectSize from starting object pointer
+                        if (minObjAddress != long.MaxValue && (addr - minObjAddress > objectBlockSize))
+                            done = true;
                     }
 
 
@@ -864,15 +866,13 @@ public void GetObjectInfo(byte* raw, ref long ptr, long untilptr, int objectBloc
                         var value_addr = GetValueAddressInfo((long)raw + ptr);
                         var addr = value_addr->Address;
 
-                        // If object pointer is greater than kObjectSize from starting object pointer
-                        if (minObjAddress != long.MaxValue && (addr - minObjAddress > objectBlockSize))
-                        {
-                            break;
-                        }
-
                         if (addr < minObjAddress) minObjAddress = addr;
                         addr += value_addr->Size;
                         if (addr > maxObjAddress) maxObjAddress = addr;
+
+                        // If object pointer is greater than kObjectSize from starting object pointer
+                        if (minObjAddress != long.MaxValue && (addr - minObjAddress > objectBlockSize))
+                            done = true;
                     }
                 }
                 ptr += GetRecordSize(ptr).Item2;