Fix rare multi-packet string corruption (#3505)

Wraith2 · web-flow · commit 63443f45e22b · 2025-07-24T09:55:51.000-07:00
* add test and simplify code to fix problems

* review feedback

* review feedback

* review feedback
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
@@ -13116,11 +13116,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();
@@ -13136,7 +13135,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)
@@ -13163,7 +13162,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
@@ -13306,11 +13306,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();
@@ -13325,8 +13324,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)
@@ -13353,7 +13352,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
@@ -1529,11 +1529,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();
@@ -1554,7 +1553,7 @@ public TdsOperationStatus TryReadByteArrayWithContinue(int length, out byte[] by
             }
             else if (result == TdsOperationStatus.NeedMoreData)
             {
-                if (isStarting || isContinuing)
+                if (canContinue)
                 {
                     SetSnapshotStorage(temp);
                 }
@@ -1983,11 +1982,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();
@@ -2005,7 +2003,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
`@@ -13116,11 +13116,10 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj`
`13116`	`13116`
`13117`	`13117`	`if (canContinue)`
`13118`	`13118`	`{`
`13119`		`- if (isContinuing \|\| isStarting)`
`13120`		`- {`
`13121`		`- temp = stateObj.TryTakeSnapshotStorage() as char[];`
`13122`		`- 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");`
`13123`		`- }`
	`13119`	`+ temp = stateObj.TryTakeSnapshotStorage() as char[];`
	`13120`	`+ Debug.Assert(temp != null \|\| !isContinuing, "if continuing stored buffer must be present to contain previous data to continue from");`
	`13121`	`+ 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");`
	`13122`	`+`
`13124`	`13123`	`if (temp != null)`
`13125`	`13124`	`{`
`13126`	`13125`	`startOffset = stateObj.GetSnapshotTotalSize();`
`@@ -13136,7 +13135,7 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj`
`13136`	`13135`	`supportRentedBuff: !canContinue, // do not use the arraypool if we are going to keep the buffer in the snapshot`
`13137`	`13136`	`rentedBuff: ref buffIsRented,`
`13138`	`13137`	`startOffset,`
`13139`		`- isStarting \|\| isContinuing`
	`13138`	`+ canContinue`
`13140`	`13139`	`);`
`13141`	`13140`
`13142`	`13141`	`if (result == TdsOperationStatus.Done)`
`@@ -13163,7 +13162,7 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj`
`13163`	`13162`	`}`
`13164`	`13163`	`else if (result == TdsOperationStatus.NeedMoreData)`
`13165`	`13164`	`{`
`13166`		`- if (isStarting \|\| isContinuing)`
	`13165`	`+ if (canContinue)`
`13167`	`13166`	`{`
`13168`	`13167`	`stateObj.SetSnapshotStorage(temp);`
`13169`	`13168`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13306,11 +13306,10 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj`
`13306`	`13306`
`13307`	`13307`	`if (canContinue)`
`13308`	`13308`	`{`
`13309`		`- if (isContinuing \|\| isStarting)`
`13310`		`- {`
`13311`		`- temp = stateObj.TryTakeSnapshotStorage() as char[];`
`13312`		`- 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");`
`13313`		`- }`
	`13309`	`+ temp = stateObj.TryTakeSnapshotStorage() as char[];`
	`13310`	`+ Debug.Assert(temp != null \|\| !isContinuing, "if continuing stored buffer must be present to contain previous data to continue from");`
	`13311`	`+ 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");`
	`13312`	`+`
`13314`	`13313`	`if (temp != null)`
`13315`	`13314`	`{`
`13316`	`13315`	`startOffset = stateObj.GetSnapshotTotalSize();`
`@@ -13325,8 +13324,8 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj`
`13325`	`13324`	`out length,`
`13326`	`13325`	`supportRentedBuff: !canContinue, // do not use the arraypool if we are going to keep the buffer in the snapshot`
`13327`	`13326`	`rentedBuff: ref buffIsRented,`
`13328`		`- startOffset,`
`13329`		`- isStarting \|\| isContinuing`
	`13327`	`+ startOffset,`
	`13328`	`+ canContinue`
`13330`	`13329`	`);`
`13331`	`13330`
`13332`	`13331`	`if (result == TdsOperationStatus.Done)`
`@@ -13353,7 +13352,7 @@ internal TdsOperationStatus TryReadPlpUnicodeCharsWithContinue(TdsParserStateObj`
`13353`	`13352`	`}`
`13354`	`13353`	`else if (result == TdsOperationStatus.NeedMoreData)`
`13355`	`13354`	`{`
`13356`		`- if (isStarting \|\| isContinuing)`
	`13355`	`+ if (canContinue)`
`13357`	`13356`	`{`
`13358`	`13357`	`stateObj.SetSnapshotStorage(temp);`
`13359`	`13358`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1529,11 +1529,10 @@ public TdsOperationStatus TryReadByteArrayWithContinue(int length, out byte[] by`
`1529`	`1529`	`(bool canContinue, bool isStarting, bool isContinuing) = GetSnapshotStatuses();`
`1530`	`1530`	`if (canContinue)`
`1531`	`1531`	`{`
`1532`		`- if (isContinuing \|\| isStarting)`
`1533`		`- {`
`1534`		`- temp = TryTakeSnapshotStorage() as byte[];`
`1535`		`- Debug.Assert(bytes == null \|\| bytes.Length == length, "stored buffer length must be null or must have been created with the correct length");`
`1536`		`- }`
	`1532`	`+ temp = TryTakeSnapshotStorage() as byte[];`
	`1533`	`+ Debug.Assert(temp != null \|\| !isContinuing, "if continuing stored buffer must be present to contain previous data to continue from");`
	`1534`	`+ Debug.Assert(bytes == null \|\| bytes.Length == length, "stored buffer length must be null or must have been created with the correct length");`
	`1535`	`+`
`1537`	`1536`	`if (temp != null)`
`1538`	`1537`	`{`
`1539`	`1538`	`offset = GetSnapshotTotalSize();`
`@@ -1554,7 +1553,7 @@ public TdsOperationStatus TryReadByteArrayWithContinue(int length, out byte[] by`
`1554`	`1553`	`}`
`1555`	`1554`	`else if (result == TdsOperationStatus.NeedMoreData)`
`1556`	`1555`	`{`
`1557`		`- if (isStarting \|\| isContinuing)`
	`1556`	`+ if (canContinue)`
`1558`	`1557`	`{`
`1559`	`1558`	`SetSnapshotStorage(temp);`
`1560`	`1559`	`}`
`@@ -1983,11 +1982,10 @@ internal TdsOperationStatus TryReadStringWithEncoding(int length, System.Text.En`
`1983`	`1982`	`int startOffset = 0;`
`1984`	`1983`	`if (canContinue)`
`1985`	`1984`	`{`
`1986`		`- if (isContinuing \|\| isStarting)`
`1987`		`- {`
`1988`		`- buf = TryTakeSnapshotStorage() as byte[];`
`1989`		`- Debug.Assert(buf == null \|\| buf.Length == length, "stored buffer length must be null or must have been created with the correct length");`
`1990`		`- }`
	`1985`	`+ buf = TryTakeSnapshotStorage() as byte[];`
	`1986`	`+ Debug.Assert(buf != null \|\| !isContinuing, "if continuing stored buffer must be present to contain previous data to continue from");`
	`1987`	`+ Debug.Assert(buf == null \|\| buf.Length == length, "stored buffer length must be null or must have been created with the correct length");`
	`1988`	`+`
`1991`	`1989`	`if (buf != null)`
`1992`	`1990`	`{`
`1993`	`1991`	`startOffset = GetSnapshotTotalSize();`
`@@ -2005,7 +2003,7 @@ internal TdsOperationStatus TryReadStringWithEncoding(int length, System.Text.En`
`2005`	`2003`	`{`
`2006`	`2004`	`if (result == TdsOperationStatus.NeedMoreData)`
`2007`	`2005`	`{`
`2008`		`- if (isStarting \|\| isContinuing)`
	`2006`	`+ if (canContinue)`
`2009`	`2007`	`{`
`2010`	`2008`	`SetSnapshotStorage(buf);`
`2011`	`2009`	`}`