Skip to content

Commit 0f2c40b

Browse files
authored
Roll back data written to output buffer on packing failure (#641)
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.
1 parent c1db2b9 commit 0f2c40b

File tree

2 files changed

+22
-2
lines changed

2 files changed

+22
-2
lines changed

neo4j/io/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -487,7 +487,8 @@ def _append(self, signature, fields=(), response=None):
487487
:param fields: the fields of the message as a tuple
488488
:param response: a response object to handle callbacks
489489
"""
490-
self.packer.pack_struct(signature, fields)
490+
with self.outbox.tmp_buffer():
491+
self.packer.pack_struct(signature, fields)
491492
self.outbox.wrap_message()
492493
self.responses.append(response)
493494

neo4j/io/_common.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,11 @@
1919
# limitations under the License.
2020

2121

22+
from contextlib import contextmanager
2223
import socket
2324
from struct import pack as struct_pack
2425

2526
from neo4j.exceptions import (
26-
AuthError,
2727
Neo4jError,
2828
ServiceUnavailable,
2929
SessionExpired,
@@ -94,11 +94,14 @@ def __init__(self, max_chunk_size=16384):
9494
self._chunked_data = bytearray()
9595
self._raw_data = bytearray()
9696
self.write = self._raw_data.extend
97+
self._tmp_buffering = 0
9798

9899
def max_chunk_size(self):
99100
return self._max_chunk_size
100101

101102
def clear(self):
103+
if self._tmp_buffering:
104+
raise RuntimeError("Cannot clear while buffering")
102105
self._chunked_data = bytearray()
103106
self._raw_data.clear()
104107

@@ -128,13 +131,29 @@ def _chunk_data(self):
128131
self._raw_data.clear()
129132

130133
def wrap_message(self):
134+
if self._tmp_buffering:
135+
raise RuntimeError("Cannot wrap message while buffering")
131136
self._chunk_data()
132137
self._chunked_data += b"\x00\x00"
133138

134139
def view(self):
140+
if self._tmp_buffering:
141+
raise RuntimeError("Cannot view while buffering")
135142
self._chunk_data()
136143
return memoryview(self._chunked_data)
137144

145+
@contextmanager
146+
def tmp_buffer(self):
147+
self._tmp_buffering += 1
148+
old_len = len(self._raw_data)
149+
try:
150+
yield
151+
except Exception:
152+
del self._raw_data[old_len:]
153+
raise
154+
finally:
155+
self._tmp_buffering -= 1
156+
138157

139158
class ConnectionErrorHandler:
140159
"""

0 commit comments

Comments
 (0)