From 1a6231472e0fb60bb65a7cfbbc5734ac61c0880b Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Tue, 22 Aug 2017 23:12:32 +0200
Subject: [PATCH] Ensure the connection is flagged as closed on errors, move
 #201

---
 telethon/extensions/tcp_client.py | 38 +++++++++++++++++++------------
 telethon/telegram_client.py       |  8 ++-----
 2 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/telethon/extensions/tcp_client.py b/telethon/extensions/tcp_client.py
index 27cc0d93..5abf620d 100644
--- a/telethon/extensions/tcp_client.py
+++ b/telethon/extensions/tcp_client.py
@@ -4,6 +4,7 @@ import time
 from datetime import datetime, timedelta
 from io import BytesIO, BufferedWriter
 from threading import Event, Lock
+import errno
 
 from ..errors import ReadCancelledError
 
@@ -43,8 +44,13 @@ class TcpClient:
     def close(self):
         """Closes the connection"""
         if self.connected:
-            self._socket.shutdown(socket.SHUT_RDWR)
-            self._socket.close()
+            try:
+                self._socket.shutdown(socket.SHUT_RDWR)
+                self._socket.close()
+            except OSError as e:
+                if e.errno != errno.ENOTCONN:
+                    raise
+
             self.connected = False
             self._recreate_socket()
 
@@ -53,18 +59,22 @@ class TcpClient:
 
         # Ensure that only one thread can send data at once
         with self._lock:
-            view = memoryview(data)
-            total_sent, total = 0, len(data)
-            while total_sent < total:
-                try:
-                    sent = self._socket.send(view[total_sent:])
-                    if sent == 0:
-                        raise ConnectionResetError(
-                            'The server has closed the connection.')
-                    total_sent += sent
+            try:
+                view = memoryview(data)
+                total_sent, total = 0, len(data)
+                while total_sent < total:
+                    try:
+                        sent = self._socket.send(view[total_sent:])
+                        if sent == 0:
+                            raise ConnectionResetError(
+                                'The server has closed the connection.')
+                        total_sent += sent
 
-                except BlockingIOError:
-                    time.sleep(self.delay)
+                    except BlockingIOError:
+                        time.sleep(self.delay)
+            except BrokenPipeError:
+                self.close()
+                raise
 
     def read(self, size, timeout=timedelta(seconds=5)):
         """Reads (receives) a whole block of 'size bytes
@@ -96,7 +106,7 @@ class TcpClient:
                     try:
                         partial = self._socket.recv(bytes_left)
                         if len(partial) == 0:
-                            self.connected = False
+                            self.close()
                             raise ConnectionResetError(
                                 'The server has closed the connection.')
 
diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py
index 249d2a7b..0d179be6 100644
--- a/telethon/telegram_client.py
+++ b/telethon/telegram_client.py
@@ -321,13 +321,9 @@ class TelegramClient(TelegramBareClient):
 
         try:
             self(LogOutRequest())
+            # The server may have already disconnected us, we still
+            # try to disconnect to make sure.
             self.disconnect()
-        except OSError as e:
-            # macOS issue: https://github.com/veusz/veusz/issues/54
-            # Socket has been already closed (Errno 57)
-            # Fail on any other error
-            if e.errno != errno.ENOTCONN:
-                raise
         except (RPCError, ConnectionError):
             # Something happened when logging out, restore the state back
             self._sender.logging_out = False