Fix race in HTTP/2 buffered data writes

src/Driver/Http2Driver.php

@@ -606,25 +606,19 @@ private function writeBufferedData(int $streamId): void
             $this->timeoutTracker->update($streamId);
         }
 
+        // Stream state must be updated before writing, since a write may suspend and let another
+        // fiber send the same buffered data.
         if ($delta >= $length) {
+            $data = $stream->buffer;
+            $stream->buffer = "";
+            $stream->clientWindow -= $length;
             $this->clientWindow -= $length;
 
-            if ($length > $this->maxFrameSize) {
-                $split = \str_split($stream->buffer, $this->maxFrameSize);
-                $stream->buffer = \array_pop($split);
-                foreach ($split as $part) {
-                    $this->writeFrame($part, Http2Parser::DATA, Http2Parser::NO_FLAG, $streamId);
-                }
-            }
-
-            if ($stream->state & Http2Stream::LOCAL_CLOSED) {
-                $this->writeFrame($stream->buffer, Http2Parser::DATA, Http2Parser::END_STREAM, $streamId);
-            } else {
-                $this->writeFrame($stream->buffer, Http2Parser::DATA, Http2Parser::NO_FLAG, $streamId);
-            }
-
-            $stream->clientWindow -= $length;
-            $stream->buffer = "";
+            $this->writeDataFrames(
+                $data,
+                $stream->state & Http2Stream::LOCAL_CLOSED ? Http2Parser::END_STREAM : Http2Parser::NO_FLAG,
+                $streamId,
+            );
 
             if ($stream->deferredFuture) {
                 $stream->deferredFuture->complete();
@@ -635,30 +629,40 @@ private function writeBufferedData(int $streamId): void
         }
 
         if ($delta > 0) {
-            $data = $stream->buffer;
-            $end = $delta - $this->maxFrameSize;
-
+            $data = \substr($stream->buffer, 0, $delta);
+            $stream->buffer = \substr($stream->buffer, $delta);
             $stream->clientWindow -= $delta;
             $this->clientWindow -= $delta;
 
-            for ($off = 0; $off < $end; $off += $this->maxFrameSize) {
-                $this->writeFrame(
-                    \substr($data, $off, $this->maxFrameSize),
-                    Http2Parser::DATA,
-                    Http2Parser::NO_FLAG,
-                    $streamId
-                );
-            }
-
-            $this->writeFrame(\substr($data, $off, $delta - $off), Http2Parser::DATA, Http2Parser::NO_FLAG, $streamId);
-
-            $stream->buffer = \substr($data, $delta);
+            $this->writeDataFrames($data, Http2Parser::NO_FLAG, $streamId);
         }
 
         $stream->deferredFuture ??= new DeferredFuture;
         $stream->deferredFuture->getFuture()->await();
     }
 
+    /**
+     * Writes all frames at once, so frames of concurrent writers cannot interleave.
+     */
+    private function writeDataFrames(string $data, int $flags, int $streamId): void
+    {
+        $frames = "";
+        $length = \strlen($data);
+
+        for ($off = 0; $length - $off > $this->maxFrameSize; $off += $this->maxFrameSize) {
+            $frames .= Http2Parser::compileFrame(
+                \substr($data, $off, $this->maxFrameSize),
+                Http2Parser::DATA,
+                Http2Parser::NO_FLAG,
+                $streamId
+            );
+        }
+
+        $frames .= Http2Parser::compileFrame(\substr($data, $off), Http2Parser::DATA, $flags, $streamId);
+
+        $this->writableStream->write($frames);
+    }
+
     private function writeHeaders(string $headers, int $type, int $flags, int $id): void
     {
         $flags |= Http2Parser::END_HEADERS;

test/Driver/Http2DriverTest.php

@@ -9,6 +9,7 @@
 use Amp\ByteStream\ReadableIterableStream;
 use Amp\ByteStream\ReadableStream;
 use Amp\ByteStream\ReadableStreamChain;
+use Amp\ByteStream\WritableStream;
 use Amp\CancelledException;
 use Amp\Future;
 use Amp\Http\HPack;
@@ -717,6 +718,100 @@ public function testFlowControl(): void
         ], $frames->getValue());
     }
 
+    public function testConcurrentWindowUpdatesWithSuspendedWrites(): void
+    {
+        $body = \random_bytes(100_000);
+
+        $this->givenNextResponse(new Response(
+            HttpStatus::OK,
+            body: new ReadableBuffer($body),
+            trailers: new Trailers(Future::complete(['x-trailer' => 'done']), ['x-trailer']),
+        ));
+
+        // Every write suspends, as a socket write does when the kernel buffer is full.
+        $output = new class($this->output->getSink()) implements WritableStream {
+            public function __construct(private readonly WritableStream $sink)
+            {
+            }
+
+            public function write(string $bytes): void
+            {
+                $this->sink->write($bytes);
+                delay(0);
+            }
+
+            public function end(): void
+            {
+                $this->sink->end();
+            }
+
+            public function isWritable(): bool
+            {
+                return $this->sink->isWritable();
+            }
+
+            public function close(): void
+            {
+                $this->sink->close();
+            }
+
+            public function isClosed(): bool
+            {
+                return $this->sink->isClosed();
+            }
+
+            public function onClose(\Closure $onClose): void
+            {
+                $this->sink->onClose($onClose);
+            }
+        };
+
+        $input = new Queue;
+        $this->givenInput(new ReadableIterableStream($input->pipe()));
+        $frames = $this->whenReceivingFrames($output);
+
+        $input->push(Http2Parser::PREFACE);
+
+        self::assertTrue($frames->continue());
+        self::assertSame(Http2Parser::SETTINGS, $frames->getValue()['type']);
+
+        $input->push(self::packHeader([
+            ":authority" => "localhost",
+            ":path" => "/",
+            ":scheme" => "http",
+            ":method" => "GET",
+        ]));
+
+        self::assertTrue($frames->continue());
+        self::assertSame(Http2Parser::HEADERS, $frames->getValue()['type']);
+
+        $recv = "";
+        while (\strlen($recv) < Http2Driver::DEFAULT_WINDOW_SIZE) {
+            self::assertTrue($frames->continue());
+            self::assertSame(Http2Parser::DATA, $frames->getValue()['type']);
+            $recv .= $frames->getValue()['buffer'];
+        }
+
+        // Both updates are handled at once, so the deferred sends of buffered data run concurrently.
+        $input->push(
+            self::packFrame(\pack("N", Http2Driver::DEFAULT_WINDOW_SIZE), Http2Parser::WINDOW_UPDATE, Http2Parser::NO_FLAG)
+            . self::packFrame(\pack("N", Http2Driver::DEFAULT_WINDOW_SIZE), Http2Parser::WINDOW_UPDATE, Http2Parser::NO_FLAG, 1)
+        );
+
+        do {
+            self::assertTrue($frames->continue(new TimeoutCancellation(1)));
+            $frame = $frames->getValue();
+
+            if ($frame['type'] === Http2Parser::DATA) {
+                $recv .= $frame['buffer'];
+            }
+        } while (!($frame['flags'] & Http2Parser::END_STREAM));
+
+        self::assertSame(Http2Parser::HEADERS, $frame['type']);
+        self::assertSame(\strlen($body), \strlen($recv));
+        self::assertSame($body, $recv);
+    }
+
     public function testPush(): void
     {
         /** @noinspection PhpInternalEntityUsedInspection */
@@ -1041,12 +1136,12 @@ private function givenNextResponse(Response $response)
      *     buffer: string,
      * }>
      */
-    private function whenReceivingFrames(): ConcurrentIterator
+    private function whenReceivingFrames(?WritableStream $output = null): ConcurrentIterator
     {
         async(fn () => $this->driver->handleClient(
             $this->createClientMock(),
             $this->input,
-            $this->output->getSink(),
+            $output ?? $this->output->getSink(),
         ))->ignore();
 
         return Pipeline::fromIterable($this->receiveFrames())->getIterator();