From 22e71c463d7c1dbccc31893afebd200937a27e6d Mon Sep 17 00:00:00 2001
From: Rouven Bauer <rouven.bauer@neo4j.com>
Date: Wed, 12 Jan 2022 10:25:06 +0100
Subject: [PATCH] Roll back data written to output buffer on packing failure

While packing data to packstream, several errors can occur (integers that
are out of bounds, unknown data types, etc.). On packing failure, the driver
should never send the half-finished packed data over the wire. This will most
likely cause the server to close the connection as the data will be corrupt.
---
 neo4j/io/__init__.py |  3 ++-
 neo4j/io/_common.py  | 21 ++++++++++++++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/neo4j/io/__init__.py b/neo4j/io/__init__.py
index 348bbce8..d7ed2323 100644
--- a/neo4j/io/__init__.py
+++ b/neo4j/io/__init__.py
@@ -487,7 +487,8 @@ def _append(self, signature, fields=(), response=None):
         :param fields: the fields of the message as a tuple
         :param response: a response object to handle callbacks
         """
-        self.packer.pack_struct(signature, fields)
+        with self.outbox.tmp_buffer():
+            self.packer.pack_struct(signature, fields)
         self.outbox.wrap_message()
         self.responses.append(response)
 
diff --git a/neo4j/io/_common.py b/neo4j/io/_common.py
index becb7db4..8fa25ee8 100644
--- a/neo4j/io/_common.py
+++ b/neo4j/io/_common.py
@@ -19,11 +19,11 @@
 # limitations under the License.
 
 
+from contextlib import contextmanager
 import socket
 from struct import pack as struct_pack
 
 from neo4j.exceptions import (
-    AuthError,
     Neo4jError,
     ServiceUnavailable,
     SessionExpired,
@@ -94,11 +94,14 @@ def __init__(self, max_chunk_size=16384):
         self._chunked_data = bytearray()
         self._raw_data = bytearray()
         self.write = self._raw_data.extend
+        self._tmp_buffering = 0
 
     def max_chunk_size(self):
         return self._max_chunk_size
 
     def clear(self):
+        if self._tmp_buffering:
+            raise RuntimeError("Cannot clear while buffering")
         self._chunked_data = bytearray()
         self._raw_data.clear()
 
@@ -128,13 +131,29 @@ def _chunk_data(self):
         self._raw_data.clear()
 
     def wrap_message(self):
+        if self._tmp_buffering:
+            raise RuntimeError("Cannot wrap message while buffering")
         self._chunk_data()
         self._chunked_data += b"\x00\x00"
 
     def view(self):
+        if self._tmp_buffering:
+            raise RuntimeError("Cannot view while buffering")
         self._chunk_data()
         return memoryview(self._chunked_data)
 
+    @contextmanager
+    def tmp_buffer(self):
+        self._tmp_buffering += 1
+        old_len = len(self._raw_data)
+        try:
+            yield
+        except Exception:
+            del self._raw_data[old_len:]
+            raise
+        finally:
+            self._tmp_buffering -= 1
+
 
 class ConnectionErrorHandler:
     """