Fix #10: Require compression context factory in acceptor.

amphp · Sep 14, 2023 · 01be6bf · 01be6bf
1 parent 8a57456
commit 01be6bf
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 8 deletions.
diff --git a/src/Rfc6455Acceptor.php b/src/Rfc6455Acceptor.php
@@ -4,19 +4,27 @@
 
 use Amp\ForbidCloning;
 use Amp\ForbidSerialization;
+use Amp\Http;
 use Amp\Http\HttpStatus;
 use Amp\Http\Server\ErrorHandler;
 use Amp\Http\Server\Request;
 use Amp\Http\Server\Response;
+use Amp\Websocket\Compression\WebsocketCompressionContextFactory;
 use function Amp\Websocket\generateAcceptFromKey;
 
 final class Rfc6455Acceptor implements WebsocketAcceptor
 {
     use ForbidCloning;
     use ForbidSerialization;
 
-    public function __construct(private readonly ErrorHandler $errorHandler = new Internal\UpgradeErrorHandler())
-    {
+    /**
+     * @param WebsocketCompressionContextFactory|null $compressionContextFactory Provide the same instance to the
+     *      {@see WebsocketClientFactory}. Use null to disable compression.
+     */
+    public function __construct(
+        private readonly ErrorHandler $errorHandler = new Internal\UpgradeErrorHandler(),
+        private readonly ?WebsocketCompressionContextFactory $compressionContextFactory = null,
+    ) {
     }
 
     public function handleHandshake(Request $request): Response
@@ -81,10 +89,24 @@ public function handleHandshake(Request $request): Response
             return $response;
         }
 
-        return new Response(HttpStatus::SWITCHING_PROTOCOLS, [
+        $response = new Response(HttpStatus::SWITCHING_PROTOCOLS, [
             'connection' => 'upgrade',
             'upgrade' => 'websocket',
             'sec-websocket-accept' => generateAcceptFromKey($acceptKey),
         ]);
+
+        if ($this->compressionContextFactory) {
+            $extensions = Http\splitHeader($request, 'sec-websocket-extensions') ?? [];
+
+            foreach ($extensions as $extension) {
+                if ($this->compressionContextFactory->fromClientHeader($extension, $headerLine)) {
+                    /** @psalm-suppress PossiblyNullArgument */
+                    $response->setHeader('sec-websocket-extensions', $headerLine);
+                    break;
+                }
+            }
+        }
+
+        return $response;
     }
 }
diff --git a/src/Rfc6455ClientFactory.php b/src/Rfc6455ClientFactory.php
@@ -5,6 +5,7 @@
 use Amp\ByteStream\ResourceStream;
 use Amp\ForbidCloning;
 use Amp\ForbidSerialization;
+use Amp\Http;
 use Amp\Http\Server\Request;
 use Amp\Http\Server\Response;
 use Amp\Socket\Socket;
@@ -24,7 +25,8 @@ final class Rfc6455ClientFactory implements WebsocketClientFactory
     use ForbidSerialization;
 
     /**
-     * @param WebsocketCompressionContextFactory|null $compressionContextFactory Use null to disable compression.
+     * @param WebsocketCompressionContextFactory|null $compressionContextFactory Use the same instance which is
+     *      provided to the {@see WebsocketAcceptor}. 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.
      */
@@ -66,12 +68,9 @@ public function createClient(
 
         $compressionContext = null;
         if ($this->compressionContextFactory) {
-            $extensions = \array_map('trim', \explode(',', (string) $request->getHeader('sec-websocket-extensions')));
-
+            $extensions = Http\splitHeader($request, 'sec-websocket-extensions') ?? [];
             foreach ($extensions as $extension) {
                 if ($compressionContext = $this->compressionContextFactory->fromClientHeader($extension, $headerLine)) {
-                    /** @psalm-suppress PossiblyNullArgument */
-                    $response->setHeader('sec-websocket-extensions', $headerLine);
                     break;
                 }
             }