Skip to content

Commit

Permalink
BUG Handle additional errors that can come from unpickling
Browse files Browse the repository at this point in the history
  • Loading branch information
gadorlhiac committed Mar 27, 2024
1 parent cd69d7d commit 6124b92
Showing 1 changed file with 13 additions and 9 deletions.
22 changes: 13 additions & 9 deletions lute/execution/ipc.py
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ def read(self, proc: subprocess.Popen) -> Message:
if self._use_pickle:
try:
contents = pickle.loads(raw_contents)
except pickle.UnpicklingError as err:
except (pickle.UnpicklingError, ValueError, EOFError) as err:
logger.debug("PipeCommunicator (Executor) - Set _use_pickle=False")
self._use_pickle = False
contents = self._safe_unpickle_decode(raw_contents)
Expand Down Expand Up @@ -237,9 +237,18 @@ def _safe_unpickle_decode(self, maybe_mixed: bytes) -> Optional[str]:
repickled: bytes = pickle.dumps(contents)
if len(repickled) < len(maybe_mixed):
# Successful unpickling, but pickle stops even if there are more bytes
additional_data: str = maybe_mixed[len(repickled) :].decode()
contents = f"{contents}{additional_data}"
except pickle.UnpicklingError as err:
try:
additional_data: str = maybe_mixed[len(repickled) :].decode()
contents = f"{contents}{additional_data}"
except UnicodeDecodeError:
# Can't decode the bytes left by pickle, so they are lost
missing_bytes: int = len(maybe_mixed) - len(repickled)
logger.debug(
f"PipeCommunicator has truncated message. Unable to retrieve {missing_bytes} bytes."
)
except (pickle.UnpicklingError, ValueError, EOFError) as err:
# Pickle may also throw a ValueError, e.g. this bytes: b"Found! \n"
# Pickle may also throw an EOFError, eg. this bytes: b"F0\n"
try:
contents = maybe_mixed.decode()
except UnicodeDecodeError as err2:
Expand All @@ -251,11 +260,6 @@ def _safe_unpickle_decode(self, maybe_mixed: bytes) -> Optional[str]:
f"PipeCommunicator unable to decode/parse data! {err3}"
)
contents = None
except UnicodeDecodeError as err3:
missing_bytes: int = len(maybe_mixed) - len(repickled)
logger.debug(
f"PipeCommunicator has truncated message. Unable to retrieve {missing_bytes} bytes."
)
return contents

def write(self, msg: Message) -> None:
Expand Down

0 comments on commit 6124b92

Please sign in to comment.