Skip to content

Commit 4c55825

Browse files
author
Zhen
committed
Close connection directly when we failed to pack a data type by not sending the corrupted data to server
1 parent 60966e7 commit 4c55825

File tree

4 files changed

+34
-9
lines changed

4 files changed

+34
-9
lines changed

neo4j/bolt/connection.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -239,7 +239,12 @@ def append(self, signature, fields=(), response=None):
239239
log_info("C: INIT (%r, {...})", fields[0])
240240
else:
241241
raise ValueError("Unknown message signature")
242-
self.packer.pack_struct(signature, fields)
242+
try:
243+
self.packer.pack_struct(signature, fields)
244+
except ValueError as e:
245+
# We failed to pack this message, therefore we close this connection to avoid sending corrupted data
246+
self.close()
247+
raise e
243248
self.output_buffer.chunk()
244249
self.output_buffer.chunk()
245250
self.responses.append(response)

neo4j/v1/api.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -612,7 +612,8 @@ def rollback(self):
612612
def close(self):
613613
""" Close this transaction, triggering either a COMMIT or a ROLLBACK.
614614
"""
615-
if not self.closed():
615+
conn_closed = self.session._connection.closed()
616+
if not self.closed() and not conn_closed:
616617
try:
617618
self.sync()
618619
except CypherError:

test/integration/test_session.py

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1818
# See the License for the specific language governing permissions and
1919
# limitations under the License.
20-
2120
from unittest import SkipTest
2221
from uuid import uuid4
2322

@@ -596,7 +595,21 @@ def test_should_sync_after_rollback(self):
596595
assert len(buffer) == 1
597596
assert buffer[0][0] == 1
598597

599-
def test_errors_on_run(self):
598+
def test_errors_on_write_transaction(self):
599+
session = self.driver.session()
600+
with self.assertRaises(ValueError):
601+
session.write_transaction(lambda tx, uuid : tx.run("CREATE (a:Thing {uuid:$uuid})", uuid=uuid), uuid4())
602+
session.close()
603+
604+
def test_errors_on_run_transaction(self):
605+
session = self.driver.session()
606+
tx = session.begin_transaction()
607+
with self.assertRaises(ValueError):
608+
tx.run("CREATE (a:Thing {uuid:$uuid})", uuid=uuid4())
609+
tx.rollback()
610+
session.close()
611+
612+
def test_errors_on_run_session(self):
600613
session = self.driver.session()
601614
session.close()
602615
with self.assertRaises(SessionError):

test/unit/test_packstream.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from io import BytesIO
2525
from math import pi
2626
from unittest import TestCase
27+
from uuid import uuid4
2728

2829
from neo4j.bolt.io import MessageFrame as PyMessageFrame
2930
from neo4j.packstream.packer import Packer as PyPacker
@@ -273,19 +274,24 @@ def test_map_stream(self):
273274
(unpacked, unpacked_value))
274275

275276
def test_illegal_signature(self):
276-
try:
277+
with self.assertRaises(ValueError):
277278
self.assert_packable((b"XXX", ()), b"\xB0XXX")
278-
except ValueError:
279-
assert True
280-
else:
281-
assert False
282279

283280
def test_empty_struct(self):
284281
self.assert_packable((b"X", ()), b"\xB0X")
285282

286283
def test_tiny_struct(self):
287284
self.assert_packable((b"Z", (u"A", 1)), b"\xB2Z\x81A\x01")
288285

286+
def test_illegal_uuid(self):
287+
self.assert_packable(uuid4(), b"\xB0XXX")
288+
try:
289+
print(uuid4())
290+
291+
except ValueError:
292+
assert True
293+
else:
294+
assert False
289295

290296
try:
291297
from neo4j.bolt._io import MessageFrame as CMessageFrame

0 commit comments

Comments
 (0)