PYTHON-4575 Allow valid SRV hostnames with less than 3 parts (#2234)

Author: aclark4life

Diff for: doc/changelog.rst

@@ -24,6 +24,7 @@ PyMongo 4.12 brings a number of changes including:
   :class:`~pymongo.read_preferences.SecondaryPreferred`,
   :class:`~pymongo.read_preferences.Nearest`. Support for ``hedge`` will be removed in PyMongo 5.0.
 - Removed PyOpenSSL support from the asynchronous API due to limitations of the CPython asyncio.Protocol SSL implementation.
+- Allow valid SRV hostnames with less than 3 parts.
 
 Issues Resolved
 ...............

Diff for: pymongo/asynchronous/srv_resolver.py

@@ -90,14 +90,12 @@ def __init__(
             raise ConfigurationError(_INVALID_HOST_MSG % ("an IP address",))
         except ValueError:
             pass
-
         try:
-            self.__plist = self.__fqdn.split(".")[1:]
+            split_fqdn = self.__fqdn.split(".")
+            self.__plist = split_fqdn[1:] if len(split_fqdn) > 2 else split_fqdn
         except Exception:
             raise ConfigurationError(_INVALID_HOST_MSG % (fqdn,)) from None
         self.__slen = len(self.__plist)
-        if self.__slen < 2:
-            raise ConfigurationError(_INVALID_HOST_MSG % (fqdn,))
 
     async def get_options(self) -> Optional[str]:
         from dns import resolver
@@ -139,6 +137,10 @@ async def _get_srv_response_and_hosts(
 
         # Validate hosts
         for node in nodes:
+            if self.__fqdn == node[0].lower():
+                raise ConfigurationError(
+                    "Invalid SRV host: return address is identical to SRV hostname"
+                )
             try:
                 nlist = node[0].lower().split(".")[1:][-self.__slen :]
             except Exception:

Diff for: pymongo/synchronous/srv_resolver.py

@@ -90,14 +90,12 @@ def __init__(
             raise ConfigurationError(_INVALID_HOST_MSG % ("an IP address",))
         except ValueError:
             pass
-
         try:
-            self.__plist = self.__fqdn.split(".")[1:]
+            split_fqdn = self.__fqdn.split(".")
+            self.__plist = split_fqdn[1:] if len(split_fqdn) > 2 else split_fqdn
         except Exception:
             raise ConfigurationError(_INVALID_HOST_MSG % (fqdn,)) from None
         self.__slen = len(self.__plist)
-        if self.__slen < 2:
-            raise ConfigurationError(_INVALID_HOST_MSG % (fqdn,))
 
     def get_options(self) -> Optional[str]:
         from dns import resolver
@@ -139,6 +137,10 @@ def _get_srv_response_and_hosts(
 
         # Validate hosts
         for node in nodes:
+            if self.__fqdn == node[0].lower():
+                raise ConfigurationError(
+                    "Invalid SRV host: return address is identical to SRV hostname"
+                )
             try:
                 nlist = node[0].lower().split(".")[1:][-self.__slen :]
             except Exception:

Diff for: test/asynchronous/test_dns.py

@@ -30,6 +30,7 @@
     unittest,
 )
 from test.utils_shared import async_wait_until
+from unittest.mock import MagicMock, patch
 
 from pymongo.asynchronous.uri_parser import parse_uri
 from pymongo.common import validate_read_preference_tags
@@ -186,12 +187,6 @@ def create_tests(cls):
 
 class TestParsingErrors(AsyncPyMongoTestCase):
     async def test_invalid_host(self):
-        with self.assertRaisesRegex(ConfigurationError, "Invalid URI host: mongodb is not"):
-            client = self.simple_client("mongodb+srv://mongodb")
-            await client.aconnect()
-        with self.assertRaisesRegex(ConfigurationError, "Invalid URI host: mongodb.com is not"):
-            client = self.simple_client("mongodb+srv://mongodb.com")
-            await client.aconnect()
         with self.assertRaisesRegex(ConfigurationError, "Invalid URI host: an IP address is not"):
             client = self.simple_client("mongodb+srv://127.0.0.1")
             await client.aconnect()
@@ -207,5 +202,93 @@ async def test_connect_case_insensitive(self):
         self.assertGreater(len(client.topology_description.server_descriptions()), 1)
 
 
+class TestInitialDnsSeedlistDiscovery(AsyncPyMongoTestCase):
+    """
+    Initial DNS Seedlist Discovery prose tests
+    https://github.com/mongodb/specifications/blob/0a7a8b5/source/initial-dns-seedlist-discovery/tests/README.md#prose-tests
+    """
+
+    async def run_initial_dns_seedlist_discovery_prose_tests(self, test_cases):
+        for case in test_cases:
+            with patch("dns.asyncresolver.resolve") as mock_resolver:
+
+                async def mock_resolve(query, record_type, *args, **kwargs):
+                    mock_srv = MagicMock()
+                    mock_srv.target.to_text.return_value = case["mock_target"]
+                    return [mock_srv]
+
+                mock_resolver.side_effect = mock_resolve
+                domain = case["query"].split("._tcp.")[1]
+                connection_string = f"mongodb+srv://{domain}"
+                try:
+                    await parse_uri(connection_string)
+                except ConfigurationError as e:
+                    self.assertIn(case["expected_error"], str(e))
+                else:
+                    self.fail(f"ConfigurationError was not raised for query: {case['query']}")
+
+    async def test_1_allow_srv_hosts_with_fewer_than_three_dot_separated_parts(self):
+        with patch("dns.asyncresolver.resolve"):
+            await parse_uri("mongodb+srv://localhost/")
+            await parse_uri("mongodb+srv://mongo.local/")
+
+    async def test_2_throw_when_return_address_does_not_end_with_srv_domain(self):
+        test_cases = [
+            {
+                "query": "_mongodb._tcp.localhost",
+                "mock_target": "localhost.mongodb",
+                "expected_error": "Invalid SRV host",
+            },
+            {
+                "query": "_mongodb._tcp.blogs.mongodb.com",
+                "mock_target": "blogs.evil.com",
+                "expected_error": "Invalid SRV host",
+            },
+            {
+                "query": "_mongodb._tcp.blogs.mongo.local",
+                "mock_target": "test_1.evil.com",
+                "expected_error": "Invalid SRV host",
+            },
+        ]
+        await self.run_initial_dns_seedlist_discovery_prose_tests(test_cases)
+
+    async def test_3_throw_when_return_address_is_identical_to_srv_hostname(self):
+        test_cases = [
+            {
+                "query": "_mongodb._tcp.localhost",
+                "mock_target": "localhost",
+                "expected_error": "Invalid SRV host",
+            },
+            {
+                "query": "_mongodb._tcp.mongo.local",
+                "mock_target": "mongo.local",
+                "expected_error": "Invalid SRV host",
+            },
+        ]
+        await self.run_initial_dns_seedlist_discovery_prose_tests(test_cases)
+
+    async def test_4_throw_when_return_address_does_not_contain_dot_separating_shared_part_of_domain(
+        self
+    ):
+        test_cases = [
+            {
+                "query": "_mongodb._tcp.localhost",
+                "mock_target": "test_1.cluster_1localhost",
+                "expected_error": "Invalid SRV host",
+            },
+            {
+                "query": "_mongodb._tcp.mongo.local",
+                "mock_target": "test_1.my_hostmongo.local",
+                "expected_error": "Invalid SRV host",
+            },
+            {
+                "query": "_mongodb._tcp.blogs.mongodb.com",
+                "mock_target": "cluster.testmongodb.com",
+                "expected_error": "Invalid SRV host",
+            },
+        ]
+        await self.run_initial_dns_seedlist_discovery_prose_tests(test_cases)
+
+
 if __name__ == "__main__":
     unittest.main()

Diff for: test/test_dns.py

@@ -30,6 +30,7 @@
     unittest,
 )
 from test.utils_shared import wait_until
+from unittest.mock import MagicMock, patch
 
 from pymongo.common import validate_read_preference_tags
 from pymongo.errors import ConfigurationError
@@ -184,12 +185,6 @@ def create_tests(cls):
 
 class TestParsingErrors(PyMongoTestCase):
     def test_invalid_host(self):
-        with self.assertRaisesRegex(ConfigurationError, "Invalid URI host: mongodb is not"):
-            client = self.simple_client("mongodb+srv://mongodb")
-            client._connect()
-        with self.assertRaisesRegex(ConfigurationError, "Invalid URI host: mongodb.com is not"):
-            client = self.simple_client("mongodb+srv://mongodb.com")
-            client._connect()
         with self.assertRaisesRegex(ConfigurationError, "Invalid URI host: an IP address is not"):
             client = self.simple_client("mongodb+srv://127.0.0.1")
             client._connect()
@@ -205,5 +200,93 @@ def test_connect_case_insensitive(self):
         self.assertGreater(len(client.topology_description.server_descriptions()), 1)
 
 
+class TestInitialDnsSeedlistDiscovery(PyMongoTestCase):
+    """
+    Initial DNS Seedlist Discovery prose tests
+    https://github.com/mongodb/specifications/blob/0a7a8b5/source/initial-dns-seedlist-discovery/tests/README.md#prose-tests
+    """
+
+    def run_initial_dns_seedlist_discovery_prose_tests(self, test_cases):
+        for case in test_cases:
+            with patch("dns.resolver.resolve") as mock_resolver:
+
+                def mock_resolve(query, record_type, *args, **kwargs):
+                    mock_srv = MagicMock()
+                    mock_srv.target.to_text.return_value = case["mock_target"]
+                    return [mock_srv]
+
+                mock_resolver.side_effect = mock_resolve
+                domain = case["query"].split("._tcp.")[1]
+                connection_string = f"mongodb+srv://{domain}"
+                try:
+                    parse_uri(connection_string)
+                except ConfigurationError as e:
+                    self.assertIn(case["expected_error"], str(e))
+                else:
+                    self.fail(f"ConfigurationError was not raised for query: {case['query']}")
+
+    def test_1_allow_srv_hosts_with_fewer_than_three_dot_separated_parts(self):
+        with patch("dns.resolver.resolve"):
+            parse_uri("mongodb+srv://localhost/")
+            parse_uri("mongodb+srv://mongo.local/")
+
+    def test_2_throw_when_return_address_does_not_end_with_srv_domain(self):
+        test_cases = [
+            {
+                "query": "_mongodb._tcp.localhost",
+                "mock_target": "localhost.mongodb",
+                "expected_error": "Invalid SRV host",
+            },
+            {
+                "query": "_mongodb._tcp.blogs.mongodb.com",
+                "mock_target": "blogs.evil.com",
+                "expected_error": "Invalid SRV host",
+            },
+            {
+                "query": "_mongodb._tcp.blogs.mongo.local",
+                "mock_target": "test_1.evil.com",
+                "expected_error": "Invalid SRV host",
+            },
+        ]
+        self.run_initial_dns_seedlist_discovery_prose_tests(test_cases)
+
+    def test_3_throw_when_return_address_is_identical_to_srv_hostname(self):
+        test_cases = [
+            {
+                "query": "_mongodb._tcp.localhost",
+                "mock_target": "localhost",
+                "expected_error": "Invalid SRV host",
+            },
+            {
+                "query": "_mongodb._tcp.mongo.local",
+                "mock_target": "mongo.local",
+                "expected_error": "Invalid SRV host",
+            },
+        ]
+        self.run_initial_dns_seedlist_discovery_prose_tests(test_cases)
+
+    def test_4_throw_when_return_address_does_not_contain_dot_separating_shared_part_of_domain(
+        self
+    ):
+        test_cases = [
+            {
+                "query": "_mongodb._tcp.localhost",
+                "mock_target": "test_1.cluster_1localhost",
+                "expected_error": "Invalid SRV host",
+            },
+            {
+                "query": "_mongodb._tcp.mongo.local",
+                "mock_target": "test_1.my_hostmongo.local",
+                "expected_error": "Invalid SRV host",
+            },
+            {
+                "query": "_mongodb._tcp.blogs.mongodb.com",
+                "mock_target": "cluster.testmongodb.com",
+                "expected_error": "Invalid SRV host",
+            },
+        ]
+        self.run_initial_dns_seedlist_discovery_prose_tests(test_cases)
+
+
 if __name__ == "__main__":
     unittest.main()

Diff for: test/test_uri_parser.py

@@ -24,6 +24,7 @@
 sys.path[0:0] = [""]
 
 from test import unittest
+from unittest.mock import patch
 
 from bson.binary import JAVA_LEGACY
 from pymongo import ReadPreference

Diff for: tools/synchro.py

@@ -133,6 +133,7 @@
     "async_joinall": "joinall",
     "_async_create_connection": "_create_connection",
     "pymongo.asynchronous.srv_resolver._SrvResolver.get_hosts": "pymongo.synchronous.srv_resolver._SrvResolver.get_hosts",
+    "dns.asyncresolver.resolve": "dns.resolver.resolve",
 }
 
 docstring_replacements: dict[tuple[str, str], str] = {