Close the connection on an authentication exception

Author: mbabker

src/Authentication/Authenticator.php

@@ -3,11 +3,14 @@
 namespace BabDev\WebSocketBundle\Authentication;
 
 use BabDev\WebSocket\Server\Connection;
+use BabDev\WebSocketBundle\Authentication\Exception\AuthenticationException;
 
 interface Authenticator
 {
     /**
      * Attempts to authenticate the current connection.
+     *
+     * @throws AuthenticationException if there was an error while trying to authenticate the user
      */
     public function authenticate(Connection $connection): void;
 }

src/Authentication/DefaultAuthenticator.php

@@ -3,6 +3,7 @@
 namespace BabDev\WebSocketBundle\Authentication;
 
 use BabDev\WebSocket\Server\Connection;
+use BabDev\WebSocketBundle\Authentication\Exception\AuthenticationException;
 use BabDev\WebSocketBundle\Authentication\Provider\AuthenticationProvider;
 use BabDev\WebSocketBundle\Authentication\Storage\TokenStorage;
 use Psr\Log\LoggerAwareInterface;
@@ -22,6 +23,8 @@ public function __construct(
 
     /**
      * Attempts to authenticate the current connection.
+     *
+     * @throws AuthenticationException if there was an error while trying to authenticate the user
      */
     public function authenticate(Connection $connection): void
     {

src/Authentication/Exception/AuthenticationException.php

@@ -0,0 +1,7 @@
+<?php declare(strict_types=1);
+
+namespace BabDev\WebSocketBundle\Authentication\Exception;
+
+use BabDev\WebSocket\Server\WebSocketException;
+
+class AuthenticationException extends \RuntimeException implements WebSocketException {}

src/Authentication/Provider/AuthenticationProvider.php

@@ -3,6 +3,7 @@
 namespace BabDev\WebSocketBundle\Authentication\Provider;
 
 use BabDev\WebSocket\Server\Connection;
+use BabDev\WebSocketBundle\Authentication\Exception\AuthenticationException;
 use Symfony\Component\Security\Core\Authentication\Token\TokenInterface;
 
 interface AuthenticationProvider
@@ -16,6 +17,8 @@ public function supports(Connection $connection): bool;
      * Attempts to authenticate the current connection.
      *
      * Implementations can assume this method will only be executed when supports() is true.
+     *
+     * @throws AuthenticationException if there was an error while trying to authenticate the user
      */
     public function authenticate(Connection $connection): TokenInterface;
 }

src/Authentication/Provider/SessionAuthenticationProvider.php

@@ -3,6 +3,8 @@
 namespace BabDev\WebSocketBundle\Authentication\Provider;
 
 use BabDev\WebSocket\Server\Connection;
+use BabDev\WebSocket\Server\WebSocketException;
+use BabDev\WebSocketBundle\Authentication\Exception\AuthenticationException;
 use BabDev\WebSocketBundle\Authentication\Storage\TokenStorage;
 use Psr\Log\LoggerAwareInterface;
 use Psr\Log\LoggerAwareTrait;
@@ -36,9 +38,19 @@ public function supports(Connection $connection): bool
         return $attributeStore->has('session') && $attributeStore->get('session') instanceof SessionInterface;
     }
 
+    /**
+     * @throws AuthenticationException if there was an error while trying to authenticate the user
+     */
     public function authenticate(Connection $connection): TokenInterface
     {
-        $token = $this->getToken($connection);
+        try {
+            $token = $this->getToken($connection);
+        } catch (WebSocketException $exception) {
+            // Out-of-the-box, we'll get a WebSocketException from our read-only session handler if there was an issue grabbing the session data, so focus only on the component's exceptions
+            $this->logger?->error('Could not authenticate user.', ['exception' => $exception]);
+
+            throw new AuthenticationException('Could not authenticate user.', previous: $exception);
+        }
 
         $storageId = $this->tokenStorage->generateStorageId($connection);
 

src/Server/Middleware/AuthenticateUser.php

@@ -3,8 +3,10 @@
 namespace BabDev\WebSocketBundle\Server\Middleware;
 
 use BabDev\WebSocket\Server\Connection;
+use BabDev\WebSocket\Server\Connection\ClosesConnectionWithResponse;
 use BabDev\WebSocket\Server\ServerMiddleware;
 use BabDev\WebSocketBundle\Authentication\Authenticator;
+use BabDev\WebSocketBundle\Authentication\Exception\AuthenticationException;
 use BabDev\WebSocketBundle\Authentication\Storage\Exception\StorageError;
 use BabDev\WebSocketBundle\Authentication\Storage\Exception\TokenNotFound;
 use BabDev\WebSocketBundle\Authentication\Storage\TokenStorage;
@@ -16,6 +18,7 @@
  */
 final class AuthenticateUser implements ServerMiddleware, LoggerAwareInterface
 {
+    use ClosesConnectionWithResponse;
     use LoggerAwareTrait;
 
     public function __construct(
@@ -29,7 +32,13 @@ public function __construct(
      */
     public function onOpen(Connection $connection): void
     {
-        $this->authenticator->authenticate($connection);
+        try {
+            $this->authenticator->authenticate($connection);
+        } catch (AuthenticationException $exception) {
+            $this->close($connection, 401);
+
+            throw $exception;
+        }
 
         $this->middleware->onOpen($connection);
     }