Cancel connect immediately when requested by token

amphp · Oct 17, 2018 · 2cb9d0e · 2cb9d0e
1 parent c3e4949
commit 2cb9d0e
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 5 deletions.
diff --git a/src/functions.php b/src/functions.php
@@ -109,10 +109,6 @@ function connect(string $uri, ClientConnectContext $socketContext = null, Cancel
         $timeout = $socketContext->getConnectTimeout();
 
         foreach ($uris as $builtUri) {
-            if ($token) {
-                $token->throwIfRequested();
-            }
-
             try {
                 $context = \stream_context_create($socketContext->toStreamContextArray());
 
@@ -130,6 +126,7 @@ function connect(string $uri, ClientConnectContext $socketContext = null, Cancel
 
                 $deferred = new Deferred;
                 $watcher = Loop::onWritable($socket, [$deferred, 'resolve']);
+                $id = $token->subscribe([$deferred, 'fail']);
 
                 try {
                     yield Promise\timeout($deferred->promise(), $timeout);
@@ -142,6 +139,7 @@ function connect(string $uri, ClientConnectContext $socketContext = null, Cancel
                     ), 110); // See ETIMEDOUT in http://www.virtsync.com/c-error-codes-include-errno
                 } finally {
                     Loop::cancel($watcher);
+                    $token->unsubscribe($id);
                 }
 
                 // The following hack looks like the only way to detect connection refused errors with PHP's stream sockets.
@@ -153,7 +151,7 @@ function connect(string $uri, ClientConnectContext $socketContext = null, Cancel
                         $failures ? "; previous attempts: " . \implode($failures) : ""
                     ), 111); // See ECONNREFUSED in http://www.virtsync.com/c-error-codes-include-errno
                 }
-            } catch (\Exception $e) {
+            } catch (ConnectException $e) {
                 // Includes only error codes used in this file, as error codes on other OS families might be different.
                 // In fact, this might show a confusing error message on OS families that return 110 or 111 by itself.
                 $knownReasons = [

diff --git a/test/IntegrationTest.php b/test/IntegrationTest.php
@@ -2,8 +2,12 @@
 
 namespace Amp\Socket\Test;
 
+use Amp\CancelledException;
+use Amp\Socket\ClientConnectContext;
 use Amp\Socket\ClientSocket;
 use Amp\Socket\ClientTlsContext;
+use Amp\Socket\ConnectException;
+use Amp\TimeoutCancellationToken;
 use PHPUnit\Framework\TestCase;
 
 class IntegrationTest extends TestCase {
@@ -23,6 +27,22 @@ public function provideConnectArgs() {
         ];
     }
 
+    public function testConnectFailure() {
+        $this->expectException(ConnectException::class);
+        $promise = \Amp\Socket\connect('8.8.8.8:80', (new ClientConnectContext)->withConnectTimeout(1000));
+        $sock = \Amp\Promise\wait($promise);
+    }
+
+    /**
+     * @depends testConnectFailure
+     */
+    public function testConnectCancellation() {
+        $this->expectException(CancelledException::class);
+        $token = new TimeoutCancellationToken(1000);
+        $promise = \Amp\Socket\connect('8.8.8.8:80', (new ClientConnectContext)->withConnectTimeout(2000), $token);
+        $sock = \Amp\Promise\wait($promise);
+    }
+
     /**
      * @dataProvider provideCryptoConnectArgs
      */