Add IPv4-only configuration option (#12632)

* add ipv4 config option

* add tests

tcp_check/assets/configuration/spec.yaml

@@ -53,4 +53,10 @@ files:
         default: false
         example: false
         type: boolean
+    - name: ipv4_only
+      description: Enable to run the check against only IPv4-formatted addresses on every check run.
+      value:
+        default: false
+        example: false
+        type: boolean
     - template: instances/default

tcp_check/datadog_checks/tcp_check/config_models/defaults.py

@@ -30,6 +30,10 @@ def instance_ip_cache_duration(field, value):
     return get_default_field_value(field, value)
 
 
+def instance_ipv4_only(field, value):
+    return False
+
+
 def instance_metric_patterns(field, value):
     return get_default_field_value(field, value)
 

tcp_check/datadog_checks/tcp_check/config_models/instance.py

@@ -36,6 +36,7 @@ class Config:
     empty_default_hostname: Optional[bool]
     host: str
     ip_cache_duration: Optional[float]
+    ipv4_only: Optional[bool]
     metric_patterns: Optional[MetricPatterns]
     min_collection_interval: Optional[float]
     multiple_ips: Optional[bool]

tcp_check/datadog_checks/tcp_check/data/conf.yaml.example

@@ -53,6 +53,11 @@ instances:
     #
     # multiple_ips: false
 
+    ## @param ipv4_only - boolean - optional - default: false
+    ## Enable to run the check against only IPv4-formatted addresses on every check run.
+    #
+    # ipv4_only: false
+
     ## @param tags - list of strings - optional
     ## A list of tags to attach to every metric and service check emitted by this instance.
     ##

tcp_check/datadog_checks/tcp_check/tcp_check.py

@@ -32,6 +32,7 @@ def __init__(self, name, init_config, instances):
         self.ip_cache_last_ts = 0
         self.ip_cache_duration = self.DEFAULT_IP_CACHE_DURATION
         self.multiple_ips = instance.get('multiple_ips', False)
+        self.ipv4_only = instance.get('ipv4_only', False)
 
         ip_cache_duration = instance.get('ip_cache_duration', None)
         if ip_cache_duration is not None:
@@ -76,10 +77,16 @@ def addrs(self):
 
     def resolve_ips(self):
         # type: () -> None
-        self._addrs = [
-            AddrTuple(sockaddr[0], socket_type)
-            for (socket_type, _, _, _, sockaddr) in socket.getaddrinfo(self.host, self.port, 0, 0, socket.IPPROTO_TCP)
-        ]
+        if self.ipv4_only:
+            _, _, ipv4_list = socket.gethostbyname_ex(self.host)
+            self._addrs = [AddrTuple(ipv4_addr, socket.AF_INET) for ipv4_addr in ipv4_list]
+        else:
+            self._addrs = [
+                AddrTuple(sockaddr[0], socket_type)
+                for (socket_type, _, _, _, sockaddr) in socket.getaddrinfo(
+                    self.host, self.port, 0, 0, socket.IPPROTO_TCP
+                )
+            ]
         if not self.multiple_ips:
             self._addrs = self._addrs[:1]
 

tcp_check/tests/common.py

@@ -16,7 +16,7 @@
 INSTANCE_IPV6 = {
     'host': 'ip-ranges.datadoghq.com',
     'port': 80,
-    'timeout': 1.5,
+    'timeout': 5,
     'name': 'UpService',
     'tags': ["foo:bar"],
     'multiple_ips': True,
@@ -29,6 +29,7 @@
 DUAL_STACK_GETADDRINFO_LOCALHOST_IPV6 = [
     (socket.AF_INET6, socket.SOCK_STREAM, 6, '', ('::1', 80, 0, 0)),
     (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('127.0.0.1', 80)),
+    (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('ip3', 80)),
 ]
 
 DUAL_STACK_GETADDRINFO_LOCALHOST_IPV4 = [
@@ -51,6 +52,7 @@
     (socket.AF_INET6, socket.SOCK_STREAM, 6, '', ('ip1', 80, 0, 0)),
     (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('ip2', 80)),
     (socket.AF_INET6, socket.SOCK_STREAM, 6, '', ('ip3', 80, 0, 0)),
+    (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('ip4', 80)),
 ]
 
 SINGLE_STACK_GETADDRINFO_IPV4 = [

tcp_check/tests/test_tcp_check.py

@@ -111,13 +111,14 @@ def test_response_time(aggregator):
 
 
 @pytest.mark.parametrize(
-    'hostname, getaddrinfo, expected_addrs, multiple_ips, ip_count',
+    'hostname, getaddrinfo, expected_addrs, multiple_ips, ipv4_only, ip_count',
     [
         pytest.param(
             'localhost',
             common.DUAL_STACK_GETADDRINFO_LOCALHOST_IPV6,
             [AddrTuple('::1', socket.AF_INET6)],
             False,
+            False,
             1,
             id='Dual IPv4/IPv6: localhost, IPv6 resolution, multiple_ips False',
         ),
@@ -126,22 +127,29 @@ def test_response_time(aggregator):
             common.DUAL_STACK_GETADDRINFO_LOCALHOST_IPV4,
             [AddrTuple('127.0.0.1', socket.AF_INET)],
             False,
+            False,
             1,
             id='Dual IPv4/IPv6: localhost, IPv4 resolution, multiple_ips False',
         ),
         pytest.param(
             'localhost',
             common.DUAL_STACK_GETADDRINFO_LOCALHOST_IPV6,
-            [AddrTuple('::1', socket.AF_INET6), AddrTuple('127.0.0.1', socket.AF_INET)],
+            [
+                AddrTuple('::1', socket.AF_INET6),
+                AddrTuple('127.0.0.1', socket.AF_INET),
+                AddrTuple('ip3', socket.AF_INET),
+            ],
             True,
-            2,
+            False,
+            3,
             id='Dual IPv4/IPv6: localhost, IPv6 resolution, multiple_ips True',
         ),
         pytest.param(
             'localhost',
             common.DUAL_STACK_GETADDRINFO_LOCALHOST_IPV4,
             [AddrTuple('127.0.0.1', socket.AF_INET), AddrTuple('::1', socket.AF_INET6)],
             True,
+            False,
             2,
             id='Dual IPv4/IPv6: localhost, IPv4 resolution, multiple_ips True',
         ),
@@ -150,6 +158,7 @@ def test_response_time(aggregator):
             common.DUAL_STACK_GETADDRINFO_IPV4,
             [AddrTuple('ip1', socket.AF_INET)],
             False,
+            False,
             1,
             id='Dual IPv4/IPv6: string hostname, IPv4 resolution, multiple_ips False',
         ),
@@ -158,6 +167,7 @@ def test_response_time(aggregator):
             common.DUAL_STACK_GETADDRINFO_IPV6,
             [AddrTuple('ip1', socket.AF_INET6)],
             False,
+            False,
             1,
             id='Dual IPv4/IPv6: string hostname, IPv6 resolution, multiple_ips False',
         ),
@@ -170,6 +180,7 @@ def test_response_time(aggregator):
                 AddrTuple('ip3', socket.AF_INET6),
             ],
             True,
+            False,
             3,
             id='Dual IPv4/IPv6: string hostname, IPv4 resolution, multiple_ips True',
         ),
@@ -180,16 +191,19 @@ def test_response_time(aggregator):
                 AddrTuple('ip1', socket.AF_INET6),
                 AddrTuple('ip2', socket.AF_INET),
                 AddrTuple('ip3', socket.AF_INET6),
+                AddrTuple('ip4', socket.AF_INET),
             ],
             True,
-            3,
+            False,
+            4,
             id='Dual IPv4/IPv6: string hostname, IPv6 resolution, multiple_ips True',
         ),
         pytest.param(
             'localhost',
             common.SINGLE_STACK_GETADDRINFO_LOCALHOST_IPV4,
             [AddrTuple('127.0.0.1', socket.AF_INET)],
             False,
+            False,
             1,
             id='Single stack IPv4: localhost, IPv4 resolution, multiple_ips False',
         ),
@@ -198,6 +212,7 @@ def test_response_time(aggregator):
             common.SINGLE_STACK_GETADDRINFO_LOCALHOST_IPV4,
             [AddrTuple('127.0.0.1', socket.AF_INET), AddrTuple('ip2', socket.AF_INET)],
             True,
+            False,
             2,
             id='Single stack IPv4: localhost, IPv4 resolution, multiple_ips True',
         ),
@@ -206,6 +221,7 @@ def test_response_time(aggregator):
             common.SINGLE_STACK_GETADDRINFO_IPV4,
             [AddrTuple('ip1', socket.AF_INET)],
             False,
+            False,
             1,
             id='Single stack IPv4: string hostname, IPv4 resolution, multiple_ips False',
         ),
@@ -218,23 +234,51 @@ def test_response_time(aggregator):
                 AddrTuple('ip3', socket.AF_INET),
             ],
             True,
+            False,
             3,
             id='Single stack IPv4: string hostname, IPv4 resolution, multiple_ips True',
         ),
+        pytest.param(
+            'localhost',
+            common.DUAL_STACK_GETADDRINFO_LOCALHOST_IPV6,
+            [AddrTuple('127.0.0.1', socket.AF_INET), AddrTuple('ip3', socket.AF_INET)],
+            True,
+            True,
+            2,
+            id='Dual IPv4/IPv6: localhost, ipv4_only True, multiple_ips True',
+        ),
+        pytest.param(
+            'another-hostname',
+            common.DUAL_STACK_GETADDRINFO_IPV6,
+            [
+                AddrTuple('ip2', socket.AF_INET),
+                AddrTuple('ip4', socket.AF_INET),
+            ],
+            True,
+            True,
+            2,
+            id='Dual IPv4/IPv6: string hostname, ipv4_only True, multiple_ips True',
+        ),
     ],
 )
-def test_hostname_resolution(aggregator, monkeypatch, hostname, getaddrinfo, expected_addrs, multiple_ips, ip_count):
+def test_hostname_resolution(
+    aggregator, monkeypatch, hostname, getaddrinfo, expected_addrs, multiple_ips, ipv4_only, ip_count
+):
     """
     Test that string hostnames get resolved to ipv4 address format properly.
     """
     instance = deepcopy(common.INSTANCE)
     instance['host'] = hostname
     instance['multiple_ips'] = multiple_ips
+    instance['ipv4_only'] = ipv4_only
     check = TCPCheck(common.CHECK_NAME, {}, [instance])
 
     monkeypatch.setattr('socket.getaddrinfo', mock.Mock(return_value=getaddrinfo))
     monkeypatch.setattr(check, 'connect', mock.Mock())
-
+    if ipv4_only:
+        monkeypatch.setattr(
+            'socket.gethostbyname_ex', mock.Mock(return_value=(hostname, [], [addr.address for addr in expected_addrs]))
+        )
     expected_tags = [
         "instance:UpService",
         "target_host:{}".format(hostname),
@@ -293,7 +337,9 @@ def test_multiple(aggregator):
 
 def has_ipv6_connectivity():
     try:
-        for sockaddr in socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET6, 0, socket.IPPROTO_TCP):
+        for _, _, _, _, sockaddr in socket.getaddrinfo(
+            socket.gethostname(), None, socket.AF_INET6, 0, socket.IPPROTO_TCP
+        ):
             if not sockaddr[0].startswith('fe80:'):
                 return True
         return False
@@ -318,10 +364,6 @@ def test_ipv6(aggregator, check):
             if has_ipv6_connectivity():
                 aggregator.assert_service_check('tcp.can_connect', status=check.OK, tags=expected_tags)
                 aggregator.assert_metric('network.tcp.can_connect', value=1, tags=expected_tags)
-            elif platform.system() == 'Darwin':
-                # IPv6 connectivity varies when running test locally on macOS, so we do not check status or metric value
-                aggregator.assert_service_check('tcp.can_connect', tags=expected_tags)
-                aggregator.assert_metric('network.tcp.can_connect', tags=expected_tags)
             else:
                 aggregator.assert_service_check('tcp.can_connect', status=check.CRITICAL, tags=expected_tags)
                 aggregator.assert_metric('network.tcp.can_connect', value=0, tags=expected_tags)