Only skip past the end boundary if there is a newline character

sebmarkbage · sebmarkbage · commit 1f28f5d7cddc · 2023-06-14T00:20:29.000-04:00
diff --git a/packages/react-client/src/ReactFlightClient.js b/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);
diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js b/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);
   });
 });
