From 45446575cdc4b2dce179b5be50965e54224f61e8 Mon Sep 17 00:00:00 2001
From: Hanne Moa <hanne.moa@sikt.no>
Date: Thu, 6 Jun 2024 14:51:15 +0200
Subject: [PATCH] Raise better exception if losing connection

* Handle disconnects better in old ritz code
* Have the manager raise LostConnectionError for known trouble spots
* Make LostConnectionError a NotConnectedError
---
 src/zinolib/controllers/zino1.py |  8 +++++++-
 src/zinolib/ritz.py              | 17 +++++++++++------
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/src/zinolib/controllers/zino1.py b/src/zinolib/controllers/zino1.py
index 6b4b660..f5b4bec 100644
--- a/src/zinolib/controllers/zino1.py
+++ b/src/zinolib/controllers/zino1.py
@@ -83,7 +83,7 @@
 from .base import EventManager, EventOrId
 from ..compat import StrEnum
 from ..event_types import EventType, Event, HistoryEntry, LogEntry, AdmState
-from ..ritz import ZinoError, ProtocolError, ritz, notifier
+from ..ritz import ZinoError, ProtocolError, ritz, notifier, NotConnectedError
 from ..utils import log_exception_with_params
 
 
@@ -120,6 +120,10 @@ class EventClosedError(Zino1Error):
     pass
 
 
+class LostConnectionError(NotConnectedError):
+    pass
+
+
 def convert_timestamp(timestamp: int) -> datetime:
     return datetime.fromtimestamp(timestamp, timezone.utc)
 
@@ -370,6 +374,8 @@ def get_event_ids(request):
             return request.get_caseids()
         except ProtocolError as e:
             raise RetryError('Zino 1 failed to send a correct response header, retry') from e
+        except BrokenPipeError as e:
+            raise LostConnectionError('Lost connection to Zino 1 server') from e
 
     @staticmethod
     def poll(request, event: EventType) -> bool:
diff --git a/src/zinolib/ritz.py b/src/zinolib/ritz.py
index 321dac3..c401f01 100644
--- a/src/zinolib/ritz.py
+++ b/src/zinolib/ritz.py
@@ -378,11 +378,11 @@ def _request(self, command: bytes, recv_buffer=4096, **_):
         while data:
             try:
                 data = self._sock.recv(recv_buffer)
-            except socket.timeout:
+            except socket.timeout as e:
                 raise TimeoutError(
                     "Timed out waiting for data. command: %s buffer: %s"
                     % (repr(command), repr(buffer))
-                )
+                ) from e
             logger.debug("recv: %s" % data.__repr__())
 
             buffer += data.decode("UTF-8", errors="windows_codepage_cp1252")
@@ -439,8 +439,8 @@ def connect(self):
             self._sock = socket.create_connection(
                 (self.server, self.port), self.timeout
             )
-        except socket.gaierror as E:
-            raise NotConnectedError(E)
+        except socket.gaierror as e:
+            raise NotConnectedError(e) from e
         response = self._request(None)
         if response.header[0] == 200:
             self.authChallenge = response.header[1].split(" ", 1)[0]
@@ -1119,6 +1119,8 @@ def connect(self):
                 (self.zino_session.server, self.port), self.timeout
             )
             self._buff = self._sock.recv(4096)
+            if not self._buff:
+                raise NotConnectedError("Lost connection to server")
             self._sock.setblocking(False)
             rawHeader = self._buff.split(bytes(self.DELIMITER, 'ascii'))[0]
             header = rawHeader.split(b" ", 1)
@@ -1146,8 +1148,11 @@ def poll(self, timeout=0):
             r, _, _ = select.select([self._sock], [], [], timeout)
             if r:
                 try:
-                    self._buff += self._sock.recv(4096).decode()
-                except socket.error as e:
+                    newbuff = self._sock.recv(4096).decode()
+                    if not newbuff:
+                        raise NotConnectedError("Lost connection to server")
+                    self._buff += newbuff
+                except OSError as e:
                     if not (
                         e.args[0] == errno.EAGAIN or e.args[0] == errno.EWOULDBLOCK
                     ):