Fix rare multi-packet string corruption (#3505) (#3513)

cheenamalhotra · Wraith2 · web-flow · commit 1af7327454f9 · 2025-07-25T10:57:46.000-03:00
* add test and simplify code to fix problems

* review feedback

* review feedback

* review feedback

Co-authored-by: Wraith &lt;wraith2@gmail.com&gt;
diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParser.cs b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParser.cs
@@ -13085,11 +13085,10 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj
 
             if (canContinue)
             {
-                if (isContinuing || isStarting)
-                {
-                    temp = stateObj.TryTakeSnapshotStorage() as char[];
-                    Debug.Assert(temp == null || length == int.MaxValue || temp.Length == length, "stored buffer length must be null or must have been created with the correct length");
-                }
+                temp = stateObj.TryTakeSnapshotStorage() as char[];
+                Debug.Assert(temp != null || !isContinuing, "if continuing stored buffer must be present to contain previous data to continue from");
+                Debug.Assert(temp == null || length == int.MaxValue || temp.Length == length, "stored buffer length must be null or must have been created with the correct length");
+                
                 if (temp != null)
                 {
                     startOffset = stateObj.GetSnapshotTotalSize();
@@ -13105,7 +13104,7 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj
                 supportRentedBuff: !canContinue, // do not use the arraypool if we are going to keep the buffer in the snapshot
                 rentedBuff: ref buffIsRented,
                 startOffset,
-                isStarting || isContinuing
+                canContinue
             );
 
             if (result == TdsOperationStatus.Done)
@@ -13132,7 +13131,7 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj
             }
             else if (result == TdsOperationStatus.NeedMoreData)
             {
-                if (isStarting || isContinuing)
+                if (canContinue)
                 {
                     stateObj.SetSnapshotStorage(temp);
                 }
diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParser.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParser.cs
@@ -13266,11 +13266,10 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj
 
             if (canContinue)
             {
-                if (isContinuing || isStarting)
-                {
-                    temp = stateObj.TryTakeSnapshotStorage() as char[];
-                    Debug.Assert(temp == null || length == int.MaxValue || temp.Length == length, "stored buffer length must be null or must have been created with the correct length");
-                }
+                temp = stateObj.TryTakeSnapshotStorage() as char[];
+                Debug.Assert(temp != null || !isContinuing, "if continuing stored buffer must be present to contain previous data to continue from");
+                Debug.Assert(temp == null || length == int.MaxValue || temp.Length == length, "stored buffer length must be null or must have been created with the correct length");
+
                 if (temp != null)
                 {
                     startOffset = stateObj.GetSnapshotTotalSize();
@@ -13285,8 +13284,8 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj
                 out length, 
                 supportRentedBuff: !canContinue, // do not use the arraypool if we are going to keep the buffer in the snapshot
                 rentedBuff: ref buffIsRented, 
-                startOffset, 
-                isStarting || isContinuing
+                startOffset,
+                canContinue
             );
 
             if (result == TdsOperationStatus.Done)
@@ -13313,7 +13312,7 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj
             }
             else if (result == TdsOperationStatus.NeedMoreData)
             {
-                if (isStarting || isContinuing)
+                if (canContinue)
                 {
                     stateObj.SetSnapshotStorage(temp);
                 }
diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsParserStateObject.cs
@@ -1527,11 +1527,10 @@ public TdsOperationStatus TryReadByteArrayWithContinue(int length, out byte[] by
             (bool canContinue, bool isStarting, bool isContinuing) = GetSnapshotStatuses();
             if (canContinue)
             {
-                if (isContinuing || isStarting)
-                {
-                    temp = TryTakeSnapshotStorage() as byte[];
-                    Debug.Assert(bytes == null || bytes.Length == length, "stored buffer length must be null or must have been created with the correct length");
-                }
+                temp = TryTakeSnapshotStorage() as byte[];
+                Debug.Assert(temp != null || !isContinuing, "if continuing stored buffer must be present to contain previous data to continue from");
+                Debug.Assert(bytes == null || bytes.Length == length, "stored buffer length must be null or must have been created with the correct length");
+                
                 if (temp != null)
                 {
                     offset = GetSnapshotTotalSize();
@@ -1552,7 +1551,7 @@ public TdsOperationStatus TryReadByteArrayWithContinue(int length, out byte[] by
             }
             else if (result == TdsOperationStatus.NeedMoreData)
             {
-                if (isStarting || isContinuing)
+                if (canContinue)
                 {
                     SetSnapshotStorage(temp);
                 }
@@ -1981,11 +1980,10 @@ internal TdsOperationStatus TryReadStringWithEncoding(int length, System.Text.En
                     int startOffset = 0;
                     if (canContinue)
                     {
-                        if (isContinuing || isStarting)
-                        {
-                            buf = TryTakeSnapshotStorage() as byte[];
-                            Debug.Assert(buf == null || buf.Length == length, "stored buffer length must be null or must have been created with the correct length");
-                        }
+                        buf = TryTakeSnapshotStorage() as byte[];
+                        Debug.Assert(buf != null || !isContinuing, "if continuing stored buffer must be present to contain previous data to continue from");
+                        Debug.Assert(buf == null || buf.Length == length, "stored buffer length must be null or must have been created with the correct length");
+                        
                         if (buf != null)
                         {
                             startOffset = GetSnapshotTotalSize();
@@ -2003,7 +2001,7 @@ internal TdsOperationStatus TryReadStringWithEncoding(int length, System.Text.En
                     {
                         if (result == TdsOperationStatus.NeedMoreData)
                         {
-                            if (isStarting || isContinuing)
+                            if (canContinue)
                             {
                                 SetSnapshotStorage(buf);
                             }
diff --git a/src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/DataReaderTest/DataReaderTest.cs b/src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/DataReaderTest/DataReaderTest.cs

Original file line number	Diff line number	Diff line change
`@@ -13085,11 +13085,10 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj`
`13085`	`13085`
`13086`	`13086`	`if (canContinue)`
`13087`	`13087`	`{`
`13088`		`- if (isContinuing \|\| isStarting)`
`13089`		`- {`
`13090`		`- temp = stateObj.TryTakeSnapshotStorage() as char[];`
`13091`		`- Debug.Assert(temp == null \|\| length == int.MaxValue \|\| temp.Length == length, "stored buffer length must be null or must have been created with the correct length");`
`13092`		`- }`
	`13088`	`+ temp = stateObj.TryTakeSnapshotStorage() as char[];`
	`13089`	`+ Debug.Assert(temp != null \|\| !isContinuing, "if continuing stored buffer must be present to contain previous data to continue from");`
	`13090`	`+ Debug.Assert(temp == null \|\| length == int.MaxValue \|\| temp.Length == length, "stored buffer length must be null or must have been created with the correct length");`
	`13091`	`+`
`13093`	`13092`	`if (temp != null)`
`13094`	`13093`	`{`
`13095`	`13094`	`startOffset = stateObj.GetSnapshotTotalSize();`
`@@ -13105,7 +13104,7 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj`
`13105`	`13104`	`supportRentedBuff: !canContinue, // do not use the arraypool if we are going to keep the buffer in the snapshot`
`13106`	`13105`	`rentedBuff: ref buffIsRented,`
`13107`	`13106`	`startOffset,`
`13108`		`- isStarting \|\| isContinuing`
	`13107`	`+ canContinue`
`13109`	`13108`	`);`
`13110`	`13109`
`13111`	`13110`	`if (result == TdsOperationStatus.Done)`
`@@ -13132,7 +13131,7 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj`
`13132`	`13131`	`}`
`13133`	`13132`	`else if (result == TdsOperationStatus.NeedMoreData)`
`13134`	`13133`	`{`
`13135`		`- if (isStarting \|\| isContinuing)`
	`13134`	`+ if (canContinue)`
`13136`	`13135`	`{`
`13137`	`13136`	`stateObj.SetSnapshotStorage(temp);`
`13138`	`13137`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13266,11 +13266,10 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj`
`13266`	`13266`
`13267`	`13267`	`if (canContinue)`
`13268`	`13268`	`{`
`13269`		`- if (isContinuing \|\| isStarting)`
`13270`		`- {`
`13271`		`- temp = stateObj.TryTakeSnapshotStorage() as char[];`
`13272`		`- Debug.Assert(temp == null \|\| length == int.MaxValue \|\| temp.Length == length, "stored buffer length must be null or must have been created with the correct length");`
`13273`		`- }`
	`13269`	`+ temp = stateObj.TryTakeSnapshotStorage() as char[];`
	`13270`	`+ Debug.Assert(temp != null \|\| !isContinuing, "if continuing stored buffer must be present to contain previous data to continue from");`
	`13271`	`+ Debug.Assert(temp == null \|\| length == int.MaxValue \|\| temp.Length == length, "stored buffer length must be null or must have been created with the correct length");`
	`13272`	`+`
`13274`	`13273`	`if (temp != null)`
`13275`	`13274`	`{`
`13276`	`13275`	`startOffset = stateObj.GetSnapshotTotalSize();`
`@@ -13285,8 +13284,8 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj`
`13285`	`13284`	`out length,`
`13286`	`13285`	`supportRentedBuff: !canContinue, // do not use the arraypool if we are going to keep the buffer in the snapshot`
`13287`	`13286`	`rentedBuff: ref buffIsRented,`
`13288`		`- startOffset,`
`13289`		`- isStarting \|\| isContinuing`
	`13287`	`+ startOffset,`
	`13288`	`+ canContinue`
`13290`	`13289`	`);`
`13291`	`13290`
`13292`	`13291`	`if (result == TdsOperationStatus.Done)`
`@@ -13313,7 +13312,7 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj`
`13313`	`13312`	`}`
`13314`	`13313`	`else if (result == TdsOperationStatus.NeedMoreData)`
`13315`	`13314`	`{`
`13316`		`- if (isStarting \|\| isContinuing)`
	`13315`	`+ if (canContinue)`
`13317`	`13316`	`{`
`13318`	`13317`	`stateObj.SetSnapshotStorage(temp);`
`13319`	`13318`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1527,11 +1527,10 @@ public TdsOperationStatus TryReadByteArrayWithContinue(int length, out byte[] by`
`1527`	`1527`	`(bool canContinue, bool isStarting, bool isContinuing) = GetSnapshotStatuses();`
`1528`	`1528`	`if (canContinue)`
`1529`	`1529`	`{`
`1530`		`- if (isContinuing \|\| isStarting)`
`1531`		`- {`
`1532`		`- temp = TryTakeSnapshotStorage() as byte[];`
`1533`		`- Debug.Assert(bytes == null \|\| bytes.Length == length, "stored buffer length must be null or must have been created with the correct length");`
`1534`		`- }`
	`1530`	`+ temp = TryTakeSnapshotStorage() as byte[];`
	`1531`	`+ Debug.Assert(temp != null \|\| !isContinuing, "if continuing stored buffer must be present to contain previous data to continue from");`
	`1532`	`+ Debug.Assert(bytes == null \|\| bytes.Length == length, "stored buffer length must be null or must have been created with the correct length");`
	`1533`	`+`
`1535`	`1534`	`if (temp != null)`
`1536`	`1535`	`{`
`1537`	`1536`	`offset = GetSnapshotTotalSize();`
`@@ -1552,7 +1551,7 @@ public TdsOperationStatus TryReadByteArrayWithContinue(int length, out byte[] by`
`1552`	`1551`	`}`
`1553`	`1552`	`else if (result == TdsOperationStatus.NeedMoreData)`
`1554`	`1553`	`{`
`1555`		`- if (isStarting \|\| isContinuing)`
	`1554`	`+ if (canContinue)`
`1556`	`1555`	`{`
`1557`	`1556`	`SetSnapshotStorage(temp);`
`1558`	`1557`	`}`
`@@ -1981,11 +1980,10 @@ internal TdsOperationStatus TryReadStringWithEncoding(int length, System.Text.En`
`1981`	`1980`	`int startOffset = 0;`
`1982`	`1981`	`if (canContinue)`
`1983`	`1982`	`{`
`1984`		`- if (isContinuing \|\| isStarting)`
`1985`		`- {`
`1986`		`- buf = TryTakeSnapshotStorage() as byte[];`
`1987`		`- Debug.Assert(buf == null \|\| buf.Length == length, "stored buffer length must be null or must have been created with the correct length");`
`1988`		`- }`
	`1983`	`+ buf = TryTakeSnapshotStorage() as byte[];`
	`1984`	`+ Debug.Assert(buf != null \|\| !isContinuing, "if continuing stored buffer must be present to contain previous data to continue from");`
	`1985`	`+ Debug.Assert(buf == null \|\| buf.Length == length, "stored buffer length must be null or must have been created with the correct length");`
	`1986`	`+`
`1989`	`1987`	`if (buf != null)`
`1990`	`1988`	`{`
`1991`	`1989`	`startOffset = GetSnapshotTotalSize();`
`@@ -2003,7 +2001,7 @@ internal TdsOperationStatus TryReadStringWithEncoding(int length, System.Text.En`
`2003`	`2001`	`{`
`2004`	`2002`	`if (result == TdsOperationStatus.NeedMoreData)`
`2005`	`2003`	`{`
`2006`		`- if (isStarting \|\| isContinuing)`
	`2004`	`+ if (canContinue)`
`2007`	`2005`	`{`
`2008`	`2006`	`SetSnapshotStorage(buf);`
`2009`	`2007`	`}`