Return response instead of throwing if server responds early

amphp · Jan 12, 2025 · 32c715f · 32c715f
1 parent 1bcb4dc
commit 32c715f
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 8 deletions.
diff --git a/src/Connection/Internal/Http2ConnectionProcessor.php b/src/Connection/Internal/Http2ConnectionProcessor.php
@@ -981,6 +981,11 @@ public function request(Request $request, Cancellation $cancellation, Stream $st
             fn (CancelledException $exception) => $this->releaseStream($streamId, $exception, false),
         );
 
+        $cancellation = $http2stream->cancellation; // Use CompositeCancellation from Http2Stream.
+
+        \assert($http2stream->pendingResponse !== null);
+        $responseFuture = $http2stream->pendingResponse->getFuture();
+
         \assert($http2stream->trailers !== null);
         $http2stream->trailers->getFuture()
             ->finally(static fn () => $cancellation->unsubscribe($cancellationId))
@@ -1015,9 +1020,6 @@ public function request(Request $request, Cancellation $cancellation, Stream $st
 
             events()->requestHeaderEnd($request, $stream);
 
-            \assert($http2stream->pendingResponse !== null);
-            $responseFuture = $http2stream->pendingResponse->getFuture();
-
             events()->requestBodyStart($request, $stream);
 
             if ($chunk === null) {
@@ -1048,8 +1050,6 @@ public function request(Request $request, Cancellation $cancellation, Stream $st
             }
 
             events()->requestBodyEnd($request, $stream);
-
-            return $responseFuture->await();
         } catch (\Throwable $exception) {
             $cancellation->unsubscribe($cancellationId);
 
@@ -1060,9 +1060,9 @@ public function request(Request $request, Cancellation $cancellation, Stream $st
             }
 
             $this->releaseStream($streamId, $exception, false);
-
-            throw $exception;
         }
+
+        return $responseFuture->await();
     }
 
     public function isClosed(): bool
@@ -1367,6 +1367,8 @@ private function releaseStream(int $streamId, ?\Throwable $exception, bool $unpr
             }
         }
 
+        $stream->cancel();
+
         if ($this->streams) {
             return;
         }

diff --git a/src/Connection/Internal/Http2Stream.php b/src/Connection/Internal/Http2Stream.php
@@ -3,6 +3,8 @@
 namespace Amp\Http\Client\Connection\Internal;
 
 use Amp\Cancellation;
+use Amp\CompositeCancellation;
+use Amp\DeferredCancellation;
 use Amp\DeferredFuture;
 use Amp\ForbidCloning;
 use Amp\ForbidSerialization;
@@ -55,11 +57,15 @@ final class Http2Stream
 
     public ?DeferredFuture $windowSizeIncrease = null;
 
+    private readonly DeferredCancellation $deferredCancellation;
+
+    public readonly Cancellation $cancellation;
+
     public function __construct(
         public readonly int $id,
         public readonly Request $request,
         public readonly Stream $stream,
-        public readonly Cancellation $cancellation,
+        Cancellation $cancellation,
         public readonly ?string $transferWatcher,
         public readonly ?string $inactivityWatcher,
         public int $serverWindow,
@@ -69,11 +75,19 @@ public function __construct(
         $this->requestBodyCompletion = new DeferredFuture();
         $this->body = new Queue();
 
+        $this->deferredCancellation = new DeferredCancellation();
+        $this->cancellation = new CompositeCancellation($cancellation, $this->deferredCancellation->getCancellation());
+
         // Trailers future may never be exposed to the user if the request fails, so ignore.
         $this->trailers = new DeferredFuture();
         $this->trailers->getFuture()->ignore();
     }
 
+    public function cancel(): void
+    {
+        $this->deferredCancellation->cancel();
+    }
+
     public function __destruct()
     {
         if ($this->transferWatcher !== null) {
@@ -84,6 +98,8 @@ public function __destruct()
             EventLoop::cancel($this->inactivityWatcher);
         }
 
+        $this->deferredCancellation->cancel();
+
         // Setting these to null due to PHP's random destruct order on shutdown to avoid errors from double completion.
         $this->pendingResponse = null;
         $this->body = null;

diff --git a/test/Connection/Http2ConnectionTest.php b/test/Connection/Http2ConnectionTest.php
@@ -5,6 +5,7 @@
 use Amp\ByteStream\StreamException;
 use Amp\CancelledException;
 use Amp\Future;
+use Amp\Http\Client\BufferedContent;
 use Amp\Http\Client\HttpException;
 use Amp\Http\Client\InvalidRequestException;
 use Amp\Http\Client\Request;
@@ -587,4 +588,25 @@ public function testServerGoAwayFrame(): void
             $this->connection->close();
         }
     }
+
+    public function testServerEarlyResponse(): void
+    {
+        $request = new Request('http://localhost/', 'POST', BufferedContent::fromString(str_repeat('a', 2 ** 10)));
+        events()->requestStart($request);
+        $stream = $this->connection->getStream($request);
+
+        $responseFuture = async(fn () => $stream->request($request, new NullCancellation));
+
+        EventLoop::delay(0.1, function (): void {
+            $this->server->write(self::packFrame($this->hpack->encode([
+                [":status", (string) HttpStatus::PAYLOAD_TOO_LARGE],
+                ["date", formatDateHeader()],
+            ]), Http2Parser::HEADERS, Http2Parser::END_HEADERS | Http2Parser::END_STREAM, 1));
+
+            $this->server->close();
+        });
+
+        $response = $responseFuture->await();
+        self::assertSame(HttpStatus::PAYLOAD_TOO_LARGE, $response->getStatus());
+    }
 }