From 310b489dfb35274616bf897ef2493951e903c4da Mon Sep 17 00:00:00 2001
From: Daniil Gentili <daniil@daniil.it>
Date: Fri, 13 Sep 2019 18:03:18 +0200
Subject: [PATCH] Avoid issues with bad connections

---
 src/danog/MadelineProto/Connection.php        |  2 +-
 src/danog/MadelineProto/DataCenter.php        | 11 +++++++++--
 .../MadelineProto/DataCenterConnection.php    |  5 ++++-
 .../Loop/Connection/ReadLoop.php              |  9 +++++----
 .../Loop/Connection/WriteLoop.php             | 19 ++++++++++++++++---
 5 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/src/danog/MadelineProto/Connection.php b/src/danog/MadelineProto/Connection.php
index 1d8cdec6..e48ca547 100644
--- a/src/danog/MadelineProto/Connection.php
+++ b/src/danog/MadelineProto/Connection.php
@@ -310,7 +310,7 @@ class Connection extends Session
 
         $ctx->setReadCallback([$this, 'haveRead']);
         $this->stream = yield $ctx->getStream();
-        
+
         if (isset($this->old)) {
             unset($this->old);
         }
diff --git a/src/danog/MadelineProto/DataCenter.php b/src/danog/MadelineProto/DataCenter.php
index 259d9857..0ad84df5 100644
--- a/src/danog/MadelineProto/DataCenter.php
+++ b/src/danog/MadelineProto/DataCenter.php
@@ -495,19 +495,26 @@ class DataCenter
 
     public function dcConnectAsync(string $dc_number, int $id = -1): \Generator
     {
-        if (isset($this->sockets[$dc_number]) && !isset($this->sockets[$dc_number]->old)) {
+        $old = isset($this->sockets[$dc_number]) && (
+            isset($this->sockets[$dc_number]->old) ||
+            ($id !== -1 && isset($this->sockets[$dc_number]->getConnection($id)->old))
+        );
+        if (isset($this->sockets[$dc_number]) && !$old) {
             return false;
         }
         $ctxs = $this->generateContexts($dc_number);
+
         if (empty($ctxs)) {
             return false;
         }
         foreach ($ctxs as $ctx) {
             try {
-                if (isset($this->sockets[$dc_number]->old)) {
+                if ($old) {
+                    $this->API->logger->logger("Reconnecting to DC $dc_number ($id) from existing", \danog\MadelineProto\Logger::WARNING);
                     $this->sockets[$dc_number]->setExtra($this->API);
                     yield $this->sockets[$dc_number]->connect($ctx, $id);
                 } else {
+                    $this->API->logger->logger("Connecting to DC $dc_number from scratch", \danog\MadelineProto\Logger::WARNING);
                     $this->sockets[$dc_number] = new DataCenterConnection();
                     $this->sockets[$dc_number]->setExtra($this->API);
                     yield $this->sockets[$dc_number]->connect($ctx);
diff --git a/src/danog/MadelineProto/DataCenterConnection.php b/src/danog/MadelineProto/DataCenterConnection.php
index 1969a670..03de48f1 100644
--- a/src/danog/MadelineProto/DataCenterConnection.php
+++ b/src/danog/MadelineProto/DataCenterConnection.php
@@ -450,8 +450,11 @@ class DataCenterConnection implements JsonSerializable
      *
      * @return Connection
      */
-    public function getConnection(): Connection
+    public function getConnection(int $id = -1): Connection
     {
+        if ($id >= 0) {
+            return $this->connections[$id];
+        }
         if (\count($this->availableConnections) <= 1) {
             return $this->connections[0];
         }
diff --git a/src/danog/MadelineProto/Loop/Connection/ReadLoop.php b/src/danog/MadelineProto/Loop/Connection/ReadLoop.php
index 2bdb988e..b3bff8b4 100644
--- a/src/danog/MadelineProto/Loop/Connection/ReadLoop.php
+++ b/src/danog/MadelineProto/Loop/Connection/ReadLoop.php
@@ -62,7 +62,6 @@ class ReadLoop extends SignalLoop
     {
         $this->connection = $connection;
         $this->API = $connection->getExtra();
-        $ctx = $connection->getCtx();
         $this->datacenter = $connection->getDatacenterID();
         $this->datacenterConnection = $connection->getShared();
     }
@@ -81,9 +80,11 @@ class ReadLoop extends SignalLoop
                 if (isset($connection->old)) {
                     return;
                 }
-                $API->logger->logger($e);
-                $API->logger->logger("Got nothing in the socket in DC {$datacenter}, reconnecting...", Logger::ERROR);
-                Tools::callForkDefer($connection->reconnect());
+                Tools::callForkDefer((function () use ($API, $connection, $datacenter, $e) {
+                    $API->logger->logger($e);
+                    $API->logger->logger("Got nothing in the socket in DC {$datacenter}, reconnecting...", Logger::ERROR);
+                    yield $connection->reconnect();
+                })());
                 return;
             }
 
diff --git a/src/danog/MadelineProto/Loop/Connection/WriteLoop.php b/src/danog/MadelineProto/Loop/Connection/WriteLoop.php
index 46f8720c..b853f6b9 100644
--- a/src/danog/MadelineProto/Loop/Connection/WriteLoop.php
+++ b/src/danog/MadelineProto/Loop/Connection/WriteLoop.php
@@ -74,12 +74,23 @@ class WriteLoop extends ResumableSignalLoop
         $please_wait = false;
         while (true) {
             while (empty($connection->pending_outgoing) || $please_wait) {
+                if (isset($connection->old)) {
+                    $API->logger->logger('Not writing because connection is old');
+                    return;
+                }
+
                 $please_wait = false;
                 $API->logger->logger("Waiting in $this", Logger::ULTRA_VERBOSE);
                 if (yield $this->waitSignal($this->pause())) {
+                    $API->logger->logger("Exiting $this", Logger::ULTRA_VERBOSE);
                     return;
                 }
                 $API->logger->logger("Done waiting in $this", Logger::ULTRA_VERBOSE);
+
+                if (isset($connection->old)) {
+                    $API->logger->logger('Not writing because connection is old');
+                    return;
+                }
             }
 
             $connection->writing(true);
@@ -89,9 +100,11 @@ class WriteLoop extends ResumableSignalLoop
                 if (isset($connection->old)) {
                     return;
                 }
-                $API->logger($e);
-                $API->logger->logger("Got nothing in the socket in DC {$datacenter}, reconnecting...", Logger::ERROR);
-                Tools::callForkDefer($connection->reconnect());
+                Tools::callForkDefer((function () use ($API, $connection, $datacenter, $e) {
+                    $API->logger->logger($e);
+                    $API->logger->logger("Got nothing in the socket in DC {$datacenter}, reconnecting...", Logger::ERROR);
+                    yield $connection->reconnect();
+                })());
                 return;
             } finally {
                 $connection->writing(false);