Pass cancellation token to upgrade request

amphp · Dec 31, 2019 · ade977a · ade977a
1 parent 78f9b32
commit ade977a
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/src/Http1TunnelConnector.php b/src/Http1TunnelConnector.php
@@ -24,9 +24,13 @@ final class Http1TunnelConnector implements Connector
     use ForbidCloning;
     use ForbidSerialization;
 
-    public static function tunnel(EncryptableSocket $socket, string $target, array $customHeaders): Promise
-    {
-        return call(static function () use ($socket, $target, $customHeaders) {
+    public static function tunnel(
+        EncryptableSocket $socket,
+        string $target,
+        array $customHeaders,
+        CancellationToken $cancellationToken
+    ): Promise {
+        return call(static function () use ($socket, $target, $customHeaders, $cancellationToken) {
             $request = new Request('http://' . \str_replace('tcp://', '', $target), 'CONNECT');
             $request->setHeaders($customHeaders);
             $request->setUpgradeHandler(static function (EncryptableSocket $socket) use (&$upgradedSocket) {
@@ -39,7 +43,7 @@ public static function tunnel(EncryptableSocket $socket, string $target, array $
             $stream = yield $connection->getStream($request);
 
             /** @var Response $response */
-            $response = yield $stream->request($request, new NullCancellationToken);
+            $response = yield $stream->request($request, $cancellationToken);
 
             if ($response->getStatus() !== 200) {
                 throw new ConnectException('Failed to connect to proxy: Received a bad status code (' . $response->getStatus() . ')');
@@ -72,7 +76,7 @@ public function connect(string $uri, ?ConnectContext $context = null, ?Cancellat
 
             $socket = yield $connector->connect($this->proxyUri, $context, $token);
 
-            return self::tunnel($socket, $uri, $this->customHeaders);
+            return self::tunnel($socket, $uri, $this->customHeaders, $token ?? new NullCancellationToken);
         });
     }
 }
diff --git a/src/Https1TunnelConnector.php b/src/Https1TunnelConnector.php
@@ -6,6 +6,7 @@
 use Amp\Http\Client\Internal\ForbidCloning;
 use Amp\Http\Client\Internal\ForbidSerialization;
 use Amp\Http\Tunnel\Internal\TunnelSocket;
+use Amp\NullCancellationToken;
 use Amp\Promise;
 use Amp\Socket\ClientTlsContext;
 use Amp\Socket\ConnectContext;
@@ -51,7 +52,7 @@ public function connect(string $uri, ?ConnectContext $context = null, ?Cancellat
 
             yield $remoteSocket->setupTls($token);
 
-            $remoteSocket = yield Http1TunnelConnector::tunnel($remoteSocket, $uri, $this->customHeaders);
+            $remoteSocket = yield Http1TunnelConnector::tunnel($remoteSocket, $uri, $this->customHeaders, $token ?? new NullCancellationToken);
 
             /** @var EncryptableSocket $serverSocket */
             /** @var EncryptableSocket $clientSocket */