whitelist of prefix when doing hello messsage (#442)

Author: martin-neotech

neo4j/api.py

@@ -149,6 +149,7 @@ def values(self):
         """
         return self._values
 
+
 class ServerInfo:
 
     def __init__(self, address, protocol_version):
@@ -178,9 +179,9 @@ def version_info(self):
         # Note: Confirm that the server agent string begins with "Neo4j/" and fail gracefully if not.
         # This is intended to help prevent drivers working for non-genuine Neo4j instances.
 
-        neo4j, _, value = self.agent.partition("/")
+        prefix, _, value = self.agent.partition("/")
         try:
-            assert neo4j == "Neo4j"
+            assert prefix in ["Neo4j"]
         except AssertionError:
             raise DriverError("Server name does not start with Neo4j/")
 
@@ -198,6 +199,15 @@ def version_info(self):
                 pass
         return tuple(value)
 
+    def _update_metadata(self, metadata):
+        """Internal, update the metadata and perform check that the prefix is whitelisted by calling self.version()
+
+        :param metadata: metadata from the server
+        :type metadata: dict
+        """
+        self.metadata.update(metadata)
+        _ = self.version_info()
+
 
 class Version(tuple):
 

neo4j/io/__init__.py

@@ -218,7 +218,14 @@ def open(cls, address, *, auth=None, timeout=None, routing_context=None, **pool_
             supported_versions = Bolt.protocol_handlers().keys()
             raise BoltHandshakeError("The Neo4J server does not support communication with this driver. This driver have support for Bolt Protocols {}".format(supported_versions), address=address, request_data=handshake, response_data=data)
 
-        connection.hello()
+        try:
+            connection.hello()
+        except Exception as error:
+            log.debug("[#%04X]  C: <CLOSE> %s", s.getsockname()[1], str(error))
+            s.shutdown(SHUT_RDWR)
+            s.close()
+            raise error
+
         return connection
 
     @property

neo4j/io/_bolt3.py

@@ -143,7 +143,7 @@ def hello(self):
             logged_headers["credentials"] = "*******"
         log.debug("[#%04X]  C: HELLO %r", self.local_port, logged_headers)
         self._append(b"\x01", (headers,),
-                     response=InitResponse(self, on_success=self.server_info.metadata.update))
+                     response=InitResponse(self, on_success=self.server_info._update_metadata))
         self.send_all()
         self.fetch_all()
 

neo4j/io/_bolt4x0.py

@@ -143,7 +143,7 @@ def hello(self):
             logged_headers["credentials"] = "*******"
         log.debug("[#%04X]  C: HELLO %r", self.local_port, logged_headers)
         self._append(b"\x01", (headers,),
-                     response=InitResponse(self, on_success=self.server_info.metadata.update))
+                     response=InitResponse(self, on_success=self.server_info._update_metadata))
         self.send_all()
         self.fetch_all()
 

neo4j/io/_bolt4x1.py

@@ -143,8 +143,7 @@ def hello(self):
         if "credentials" in logged_headers:
             logged_headers["credentials"] = "*******"
         log.debug("[#%04X]  C: HELLO %r", self.local_port, logged_headers)
-        self._append(b"\x01", (headers,),
-                     response=InitResponse(self, on_success=self.server_info.metadata.update))
+        self._append(b"\x01", (headers,), response=InitResponse(self, on_success=self.server_info._update_metadata))
         self.send_all()
         self.fetch_all()
 

neo4j/work/simple.py

@@ -192,7 +192,7 @@ def run(self, query, parameters=None, **kwparameters):
         :type parameters: dict
         :param kwparameters: additional keyword parameters
         :returns: a new :class:`neo4j.Result` object
-        :type: :class:`neo4j.Result`
+        :rtype: :class:`neo4j.Result`
         """
         if not query:
             raise ValueError("Cannot run an empty query")

tests/stub/scripts/v4x1/return_1_port_9001_bogus_server.script

@@ -0,0 +1,10 @@
+!: BOLT 4.1
+!: PORT 9001
+
+C: HELLO {"user_agent": "test", "scheme": "basic", "principal": "test", "credentials": "test", "routing": {"address": "localhost:9001"}}
+S: SUCCESS {"server": "Bogus/4.1.0", "connection_id": "123e4567-e89b-12d3-a456-426655440000"}
+C: RUN "RETURN 1 AS x" {} {"mode": "r"}
+   PULL {"n": -1}
+S: SUCCESS {"fields": ["x"]}
+   RECORD [1]
+   SUCCESS {"bookmark": "neo4j:bookmark-test-1", "type": "r", "t_last": 5, "db": "neo4j"}

tests/stub/test_directdriver.py

@@ -24,6 +24,7 @@
 from neo4j.exceptions import (
     ServiceUnavailable,
     ConfigurationError,
+    DriverError,
 )
 from neo4j._exceptions import (
     BoltHandshakeError,
@@ -88,27 +89,47 @@ def test_bolt_uri_constructs_bolt_driver(driver_info, test_script):
 
 
 @pytest.mark.parametrize(
-    "test_script, test_expected",
+    "test_script",
     [
-        # ("v1/empty_explicit_hello_goodbye.script", ServiceUnavailable), # skip: cant close stub server gracefully
-        # ("v2/empty_explicit_hello_goodbye.script", ServiceUnavailable), # skip: cant close stub server gracefully
-        ("v3/empty_explicit_hello_goodbye.script", None),
-        ("v4x0/empty_explicit_hello_goodbye.script", None),
-        ("v4x1/empty_explicit_hello_goodbye.script", None),
+        "v3/empty_explicit_hello_goodbye.script",
+        "v4x0/empty_explicit_hello_goodbye.script",
+        "v4x1/empty_explicit_hello_goodbye.script",
     ]
 )
-def test_direct_driver_handshake_negotiation(driver_info, test_script, test_expected):
+def test_direct_driver_handshake_negotiation(driver_info, test_script):
     # python -m pytest tests/stub/test_directdriver.py -s -v -k test_direct_driver_handshake_negotiation
     with StubCluster(test_script):
         uri = "bolt://localhost:9001"
-        if test_expected:
-            with pytest.raises(test_expected) as error:
-                driver = GraphDatabase.driver(uri, auth=driver_info["auth_token"], **driver_config)
-            assert isinstance(error.value.__cause__, BoltHandshakeError)
-        else:
-            driver = GraphDatabase.driver(uri, auth=driver_info["auth_token"], **driver_config)
+        driver = GraphDatabase.driver(uri, auth=driver_info["auth_token"], **driver_config)
+        assert isinstance(driver, BoltDriver)
+        driver.close()
+
+
+@pytest.mark.parametrize(
+    "test_script, test_expected",
+    [
+        ("v3/return_1_port_9001.script", "Neo4j/3.0.0"),
+        ("v4x0/return_1_port_9001.script", "Neo4j/4.0.0"),
+        ("v4x1/return_1_port_9001_bogus_server.script", DriverError),
+    ]
+)
+def test_return_1_as_x(driver_info, test_script, test_expected):
+    # python -m pytest tests/stub/test_directdriver.py -s -v -k test_return_1_as_x
+    with StubCluster(test_script):
+        uri = "bolt://localhost:9001"
+        try:
+            driver = GraphDatabase.driver(uri, auth=driver_info["auth_token"], user_agent="test")
             assert isinstance(driver, BoltDriver)
+            with driver.session(default_access_mode=READ_ACCESS, fetch_size=-1) as session:
+                result = session.run("RETURN 1 AS x")
+                value = result.single().value()
+                assert value == 1
+                summary = result.consume()
+                assert summary.server.agent == test_expected
+                assert summary.server.agent.startswith("Neo4j")
             driver.close()
+        except DriverError as error:
+            assert isinstance(error, test_expected)
 
 
 def test_direct_driver_with_wrong_port(driver_info):
@@ -128,7 +149,7 @@ def test_direct_driver_with_wrong_port(driver_info):
 def test_direct_verify_connectivity(driver_info, test_script, test_expected):
     # python -m pytest tests/stub/test_directdriver.py -s -v -k test_direct_verify_connectivity
     with StubCluster(test_script):
-        uri = "bolt://127.0.0.1:9001"
+        uri = "bolt://localhost:9001"
         with GraphDatabase.driver(uri, auth=driver_info["auth_token"], **driver_config) as driver:
             assert isinstance(driver, BoltDriver)
             assert driver.verify_connectivity(default_access_mode=READ_ACCESS) == test_expected