Bound the connection handshake by the connection timeout

src/Config.php

@@ -27,7 +27,7 @@
 
     /**
      * @param non-empty-list<non-empty-string> $urls
-     * @param float $connectionTimeout in seconds
+     * @param float $connectionTimeout in seconds, bounds both the TCP connect and the handshake (INFO, CONNECT, PING/PONG)
      * @param ?non-empty-string $user
      * @param ?non-empty-string $password
      * @param ?non-empty-string $jwt JWT token for authentication

src/Internal/Connection/SocketConnection.php

@@ -9,10 +9,12 @@
 use Amp\DeferredFuture;
 use Amp\Future;
 use Amp\Socket\Socket;
+use Amp\TimeoutCancellation;
 use Revolt\EventLoop;
 use Thesis\Nats\Config;
 use Thesis\Nats\Exception\ConnectionIsNotAvailable;
 use Thesis\Nats\Exception\ConnectionWasClosed;
+use Thesis\Nats\Exception\ServerIsNotAvailable;
 use Thesis\Nats\Internal\Hooks;
 use Thesis\Nats\Internal\Nkey\KeyPair;
 use Thesis\Nats\Internal\Protocol;
@@ -68,48 +70,17 @@ public function __construct(
      */
     public function startup(): void
     {
-        $frame = $this->framer->readFrame() ?? throw new ConnectionIsNotAvailable();
-
-        if (!$frame instanceof Protocol\ServerInfo) {
-            throw new \UnexpectedValueException(
-                \sprintf('An unexpected "%s" startup frame received.', $frame::class),
-            );
-        }
-
-        if ($this->state !== ConnectionState::Alive) {
-            $this->run();
-        }
-
-        $this->info ??= ConnectionInfo::fromServerInfo($frame);
-
-        /** @var DeferredFuture<void> $handshake */
-        $handshake = new DeferredFuture();
-        $this->handshake = $handshake;
+        // The connection timeout bounds the handshake too: a peer that accepts the socket and stays silent
+        // (a hung server, a proxy or a load balancer without a backend) must not make the connection attempt hang.
+        $timeout = $this->config->connectionTimeout;
 
         try {
-            $this->connect(new Protocol\Connect(
-                verbose: $this->config->verbose,
-                pedantic: $this->config->pedantic,
-                tlsRequired: false,
-                name: $this->config->clientName,
-                version: version(self::PACKAGE_NAME),
-                user: $this->config->user,
-                pass: $this->config->password,
-                sig: $this->keypair !== null && $frame->nonce !== null && $frame->nonce !== '' ? $this->keypair->sign($frame->nonce) : null,
-                jwt: $this->config->jwt,
-                noResponders: $this->config->noResponders,
-                headers: $this->info->allowHeaders,
-                nkey: $this->keypair?->publicKey,
-            ));
-
-            // In non-verbose mode the server does not acknowledge CONNECT: it stays silent on success and answers
-            // with -ERR (e.g. Authorization Violation) before closing the socket on failure. A PING forces a reply
-            // (PONG or -ERR), so that the errors surface from here instead of being lost, the same way nats.go does it.
-            $this->framer->writeFrame(Protocol\Ping::Frame);
-            $handshake->getFuture()->await();
-        } finally {
-            $this->handshake = null;
-            $handshake->getFuture()->ignore();
+            $this->handshake(new TimeoutCancellation($timeout));
+        } catch (CancelledException $e) {
+            throw new ServerIsNotAvailable(
+                \sprintf('The nats server did not complete the handshake within %s seconds.', $timeout),
+                previous: $e,
+            );
         }
 
         if (($interval = $this->config->ping) !== null) {
@@ -163,6 +134,58 @@ public function execute(Protocol\Frame|array $frames, ?Cancellation $cancellatio
         }
     }
 
+    /**
+     * Reads INFO, sends CONNECT and waits for the server to answer the PING that follows it.
+     *
+     * @throws \Throwable
+     */
+    private function handshake(Cancellation $cancellation): void
+    {
+        $frame = $this->framer->readFrame($cancellation) ?? throw new ConnectionIsNotAvailable();
+
+        if (!$frame instanceof Protocol\ServerInfo) {
+            throw new \UnexpectedValueException(
+                \sprintf('An unexpected "%s" startup frame received.', $frame::class),
+            );
+        }
+
+        if ($this->state !== ConnectionState::Alive) {
+            $this->run();
+        }
+
+        $this->info ??= ConnectionInfo::fromServerInfo($frame);
+
+        /** @var DeferredFuture<void> $handshake */
+        $handshake = new DeferredFuture();
+        $this->handshake = $handshake;
+
+        try {
+            $this->connect(new Protocol\Connect(
+                verbose: $this->config->verbose,
+                pedantic: $this->config->pedantic,
+                tlsRequired: false,
+                name: $this->config->clientName,
+                version: version(self::PACKAGE_NAME),
+                user: $this->config->user,
+                pass: $this->config->password,
+                sig: $this->keypair !== null && $frame->nonce !== null && $frame->nonce !== '' ? $this->keypair->sign($frame->nonce) : null,
+                jwt: $this->config->jwt,
+                noResponders: $this->config->noResponders,
+                headers: $this->info->allowHeaders,
+                nkey: $this->keypair?->publicKey,
+            ), $cancellation);
+
+            // In non-verbose mode the server does not acknowledge CONNECT: it stays silent on success and answers
+            // with -ERR (e.g. Authorization Violation) before closing the socket on failure. A PING forces a reply
+            // (PONG or -ERR), so that the errors surface from here instead of being lost, the same way nats.go does it.
+            $this->framer->writeFrame(Protocol\Ping::Frame);
+            $handshake->getFuture()->await($cancellation);
+        } finally {
+            $this->handshake = null;
+            $handshake->getFuture()->ignore();
+        }
+    }
+
     public function info(): ConnectionInfo
     {
         return $this->info ?? throw new ConnectionIsNotAvailable();
@@ -270,10 +293,10 @@ private function run(): void
      *
      * @throws \Throwable
      */
-    private function connect(Protocol\Connect $connect): void
+    private function connect(Protocol\Connect $connect, Cancellation $cancellation): void
     {
         try {
-            $this->execute($connect);
+            $this->execute($connect, $cancellation);
         } catch (ConnectionWasClosed $e) {
             // The server closed the connection without acknowledging CONNECT, the same as without answering PING.
             throw new ConnectionIsNotAvailable(previous: $e);

src/Internal/Connection/SocketConnectionFactory.php

@@ -48,7 +48,14 @@ public function connect(): Connection
             $keypair,
         );
 
-        $connection->startup();
+        try {
+            $connection->startup();
+        } catch (\Throwable $e) {
+            // Nobody owns the connection yet: without closing it here a silent or misbehaving peer would keep the socket.
+            $connection->close();
+
+            throw $e;
+        }
 
         return $connection;
     }

tests/ClientTest.php

@@ -322,6 +322,73 @@ public function testRequestReconnectsAfterFailedConnection(): void
         }
     }
 
+    public function testConnectFailsWhenServerAcceptsAndStaysSilent(): void
+    {
+        $server = listen('127.0.0.1:0');
+        $address = $server->getAddress()->toString();
+        \assert($address !== '');
+
+        // The peer accepts the connection, but never sends INFO.
+        $accepted = async($server->accept(...));
+
+        $client = $this->clientWithUrls([$address]);
+        $startedAt = microtime(true);
+
+        try {
+            $client->publish('events.happens', new Message('ok'), cancellation: new TimeoutCancellation(5));
+            self::fail('The server never sends INFO, publish must fail.');
+        } catch (ServerIsNotAvailable) {
+        }
+
+        // The connection timeout of the client is 1 second.
+        self::assertLessThan(3, microtime(true) - $startedAt);
+
+        $peer = $accepted->await(new TimeoutCancellation(1));
+        self::assertNotNull($peer);
+        // The client has closed the socket it gave up on.
+        self::assertNull($peer->read(new TimeoutCancellation(1)));
+
+        $server->close();
+    }
+
+    public function testConnectFailsWhenServerDoesNotAnswerHandshake(): void
+    {
+        $proxy = PausableProxy::start($this->natsAddress());
+
+        try {
+            // The server sends INFO, but never sees CONNECT and the PING after it, so it never answers with PONG.
+            $proxy->pause();
+
+            $client = $this->clientWithUrls([$proxy->address]);
+            $startedAt = microtime(true);
+
+            try {
+                $client->publish('events.happens', new Message('ok'), cancellation: new TimeoutCancellation(5));
+                self::fail('The server never answers the handshake, publish must fail.');
+            } catch (ServerIsNotAvailable) {
+            }
+
+            self::assertLessThan(3, microtime(true) - $startedAt);
+
+            // Once the server answers again, the next call establishes a new connection.
+            $proxy->resume();
+
+            $id = Id\generateUniqueId();
+
+            /** @var DeferredFuture<Delivery> $deferred */
+            $deferred = new DeferredFuture();
+
+            $client->subscribe("{$id}.*", $deferred->complete(...));
+            $client->publish("{$id}.happens", new Message('ok'));
+
+            self::assertSame('ok', $deferred->getFuture()->await(new TimeoutCancellation(5))->message->payload);
+
+            $client->stop();
+        } finally {
+            $proxy->stop();
+        }
+    }
+
     #[DoesNotPerformAssertions]
     public function testClientWithFailedConnectionIsDestroyedQuietly(): void
     {