Move compression context creation into client factory

amphp · Sep 9, 2023 · 3235078 · 3235078
1 parent fe925cf
commit 3235078
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 31 deletions.
diff --git a/src/Rfc6455ClientFactory.php b/src/Rfc6455ClientFactory.php
@@ -8,7 +8,7 @@
 use Amp\Http\Server\Request;
 use Amp\Http\Server\Response;
 use Amp\Socket\Socket;
-use Amp\Websocket\Compression\WebsocketCompressionContext;
+use Amp\Websocket\Compression\WebsocketCompressionContextFactory;
 use Amp\Websocket\ConstantRateLimit;
 use Amp\Websocket\Parser\Rfc6455ParserFactory;
 use Amp\Websocket\Parser\WebsocketParserFactory;
@@ -23,7 +23,13 @@ final class Rfc6455ClientFactory implements WebsocketClientFactory
     use ForbidCloning;
     use ForbidSerialization;
 
+    /**
+     * @param WebsocketCompressionContextFactory|null $compressionContextFactory Use null to disable compression.
+     * @param WebsocketHeartbeatQueue|null $heartbeatQueue Use null to disable automatic heartbeats (pings).
+     * @param WebsocketRateLimit|null $rateLimit Use null to disable client rate limits.
+     */
     public function __construct(
+        private readonly ?WebsocketCompressionContextFactory $compressionContextFactory = null,
         private readonly ?WebsocketHeartbeatQueue $heartbeatQueue = new PeriodicHeartbeatQueue(),
         private readonly ?WebsocketRateLimit $rateLimit = new ConstantRateLimit(),
         private readonly WebsocketParserFactory $parserFactory = new Rfc6455ParserFactory(),
@@ -36,7 +42,6 @@ public function createClient(
         Request $request,
         Response $response,
         Socket $socket,
-        ?WebsocketCompressionContext $compressionContext = null,
     ): WebsocketClient {
         if ($socket instanceof ResourceStream) {
             $socketResource = $socket->getResource();
@@ -59,6 +64,19 @@ public function createClient(
             }
         }
 
+        $compressionContext = null;
+        if ($this->compressionContextFactory) {
+            $extensions = \array_map('trim', \explode(',', (string) $request->getHeader('sec-websocket-extensions')));
+
+            foreach ($extensions as $extension) {
+                if ($compressionContext = $this->compressionContextFactory->fromClientHeader($extension, $headerLine)) {
+                    /** @psalm-suppress PossiblyNullArgument */
+                    $response->setHeader('sec-websocket-extensions', $headerLine);
+                    break;
+                }
+            }
+        }
+
         return new Rfc6455Client(
             socket: $socket,
             masked: false,

diff --git a/src/Websocket.php b/src/Websocket.php
@@ -11,8 +11,6 @@
 use Amp\Http\Server\Request;
 use Amp\Http\Server\RequestHandler;
 use Amp\Http\Server\Response;
-use Amp\Websocket\Compression\WebsocketCompressionContext;
-use Amp\Websocket\Compression\WebsocketCompressionContextFactory;
 use Amp\Websocket\WebsocketClient;
 use Amp\Websocket\WebsocketCloseCode;
 use Amp\Websocket\WebsocketClosedException;
@@ -29,16 +27,12 @@ final class Websocket implements RequestHandler
     /** @var \WeakMap<WebsocketClient, true> */
     private \WeakMap $clients;
 
-    /**
-     * @param WebsocketCompressionContextFactory|null $compressionContextFactory Use null to disable compression.
-     */
     public function __construct(
         HttpServer $httpServer,
         private readonly PsrLogger $logger,
         private readonly RequestHandler $acceptor,
         private readonly WebsocketClientHandler $clientHandler,
         private readonly WebsocketClientFactory $clientFactory = new Rfc6455ClientFactory(),
-        private readonly ?WebsocketCompressionContextFactory $compressionContextFactory = null,
     ) {
         /** @psalm-suppress PropertyTypeCoercion */
         $this->clients = new \WeakMap();
@@ -57,24 +51,10 @@ public function handleRequest(Request $request): Response
             return $response;
         }
 
-        $compressionContext = null;
-        if ($this->compressionContextFactory) {
-            $extensions = \array_map('trim', \explode(',', (string) $request->getHeader('sec-websocket-extensions')));
-
-            foreach ($extensions as $extension) {
-                if ($compressionContext = $this->compressionContextFactory->fromClientHeader($extension, $headerLine)) {
-                    /** @psalm-suppress PossiblyNullArgument */
-                    $response->setHeader('sec-websocket-extensions', $headerLine);
-                    break;
-                }
-            }
-        }
-
         $response->upgrade(fn (UpgradedSocket $socket) => $this->reapClient(
             socket: $socket,
             request: $request,
             response: $response,
-            compressionContext: $compressionContext,
         ));
 
         return $response;
@@ -84,9 +64,8 @@ private function reapClient(
         UpgradedSocket $socket,
         Request $request,
         Response $response,
-        ?WebsocketCompressionContext $compressionContext
     ): void {
-        $client = $this->clientFactory->createClient($request, $response, $socket, $compressionContext);
+        $client = $this->clientFactory->createClient($request, $response, $socket);
 
         /** @psalm-suppress RedundantCondition */
         \assert($this->logger->debug(\sprintf(

diff --git a/src/WebsocketClientFactory.php b/src/WebsocketClientFactory.php
@@ -5,18 +5,12 @@
 use Amp\Http\Server\Request;
 use Amp\Http\Server\Response;
 use Amp\Socket\Socket;
-use Amp\Websocket\Compression\WebsocketCompressionContext;
 use Amp\Websocket\WebsocketClient;
 
 interface WebsocketClientFactory
 {
     /**
      * Creates a Client object based on the given Request.
      */
-    public function createClient(
-        Request $request,
-        Response $response,
-        Socket $socket,
-        ?WebsocketCompressionContext $compressionContext = null,
-    ): WebsocketClient;
+    public function createClient(Request $request, Response $response, Socket $socket): WebsocketClient;
 }