[Flight] Only skip past the end boundary if there is a newline charac…

…ter (facebook#26945)

Follow up to facebook#26932

For regular rows, we're increasing the index by one to skip past the
last trailing newline character which acts a boundary. For length
encoded rows we shouldn't skip an extra byte because it'll leave us
missing one.

This only accidentally worked because this was also the end of the
current chunk which tests don't account for since we're just passing
through the chunks. So I added some noise by splitting and joining the
chunks so that this gets tested.

packages/react-client/src/ReactFlightClient.js

@@ -950,29 +950,33 @@ export function processBinaryChunk(
       }
       case ROW_CHUNK_BY_LENGTH: {
         // We're looking for the remaining byte length
-        if (i + rowLength <= chunk.length) {
-          lastIdx = i + rowLength;
+        lastIdx = i + rowLength;
+        if (lastIdx > chunk.length) {
+          lastIdx = -1;
         }
         break;
       }
     }
+    const offset = chunk.byteOffset + i;
     if (lastIdx > -1) {
       // We found the last chunk of the row
-      const offset = chunk.byteOffset + i;
       const length = lastIdx - i;
       const lastChunk = new Uint8Array(chunk.buffer, offset, length);
       processFullRow(response, rowID, rowTag, buffer, lastChunk);
       // Reset state machine for a new row
+      i = lastIdx;
+      if (rowState === ROW_CHUNK_BY_NEWLINE) {
+        // If we're trailing by a newline we need to skip it.
+        i++;
+      }
       rowState = ROW_ID;
       rowTag = 0;
       rowID = 0;
       rowLength = 0;
       buffer.length = 0;
-      i = lastIdx + 1;
     } else {
       // The rest of this row is in a future chunk. We stash the rest of the
       // current chunk until we can process the full row.
-      const offset = chunk.byteOffset + i;
       const length = chunk.byteLength - i;
       const remainingSlice = new Uint8Array(chunk.buffer, offset, length);
       buffer.push(remainingSlice);

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js

@@ -42,6 +42,37 @@ describe('ReactFlightDOMEdge', () => {
     use = React.use;
   });
 
+  function passThrough(stream) {
+    // Simulate more realistic network by splitting up and rejoining some chunks.
+    // This lets us test that we don't accidentally rely on particular bounds of the chunks.
+    return new ReadableStream({
+      async start(controller) {
+        const reader = stream.getReader();
+        let prevChunk = new Uint8Array(0);
+        function push() {
+          reader.read().then(({done, value}) => {
+            if (done) {
+              controller.enqueue(prevChunk);
+              controller.close();
+              return;
+            }
+            const chunk = new Uint8Array(prevChunk.length + value.length);
+            chunk.set(prevChunk, 0);
+            chunk.set(value, prevChunk.length);
+            if (chunk.length > 50) {
+              controller.enqueue(chunk.subarray(0, chunk.length - 50));
+              prevChunk = chunk.subarray(chunk.length - 50);
+            } else {
+              prevChunk = chunk;
+            }
+            push();
+          });
+        }
+        push();
+      },
+    });
+  }
+
   async function readResult(stream) {
     const reader = stream.getReader();
     let result = '';
@@ -101,15 +132,17 @@ describe('ReactFlightDOMEdge', () => {
 
   it('should encode long string in a compact format', async () => {
     const testString = '"\n\t'.repeat(500) + '🙃';
+    const testString2 = 'hello'.repeat(400);
 
     const stream = ReactServerDOMServer.renderToReadableStream({
       text: testString,
+      text2: testString2,
     });
-    const [stream1, stream2] = stream.tee();
+    const [stream1, stream2] = passThrough(stream).tee();
 
     const serializedContent = await readResult(stream1);
     // The content should be compact an unescaped
-    expect(serializedContent.length).toBeLessThan(2000);
+    expect(serializedContent.length).toBeLessThan(4000);
     expect(serializedContent).not.toContain('\\n');
     expect(serializedContent).not.toContain('\\t');
     expect(serializedContent).not.toContain('\\"');
@@ -118,5 +151,6 @@ describe('ReactFlightDOMEdge', () => {
     const result = await ReactServerDOMClient.createFromReadableStream(stream2);
     // Should still match the result when parsed
     expect(result.text).toBe(testString);
+    expect(result.text2).toBe(testString2);
   });
 });