From b616b8e3b3d7026d77bfc62ca7ac59002d24c833 Mon Sep 17 00:00:00 2001 From: Fanny Jiang Date: Wed, 3 Aug 2022 16:09:03 -0400 Subject: [PATCH] Add IPv4-only configuration option (#12632) * add ipv4 config option * add tests --- tcp_check/assets/configuration/spec.yaml | 6 ++ .../tcp_check/config_models/defaults.py | 4 ++ .../tcp_check/config_models/instance.py | 1 + .../tcp_check/data/conf.yaml.example | 5 ++ .../datadog_checks/tcp_check/tcp_check.py | 15 +++-- tcp_check/tests/common.py | 4 +- tcp_check/tests/test_tcp_check.py | 64 +++++++++++++++---- 7 files changed, 83 insertions(+), 16 deletions(-) diff --git a/tcp_check/assets/configuration/spec.yaml b/tcp_check/assets/configuration/spec.yaml index 6fc30e7c63b41..a16ee964bdb7a 100644 --- a/tcp_check/assets/configuration/spec.yaml +++ b/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 diff --git a/tcp_check/datadog_checks/tcp_check/config_models/defaults.py b/tcp_check/datadog_checks/tcp_check/config_models/defaults.py index d9d975709e9d7..913d5e7668896 100644 --- a/tcp_check/datadog_checks/tcp_check/config_models/defaults.py +++ b/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) diff --git a/tcp_check/datadog_checks/tcp_check/config_models/instance.py b/tcp_check/datadog_checks/tcp_check/config_models/instance.py index 594fcfe105b12..fd07477213d97 100644 --- a/tcp_check/datadog_checks/tcp_check/config_models/instance.py +++ b/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] diff --git a/tcp_check/datadog_checks/tcp_check/data/conf.yaml.example b/tcp_check/datadog_checks/tcp_check/data/conf.yaml.example index 991d2d71e1a13..ad173497e409b 100644 --- a/tcp_check/datadog_checks/tcp_check/data/conf.yaml.example +++ b/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. ## diff --git a/tcp_check/datadog_checks/tcp_check/tcp_check.py b/tcp_check/datadog_checks/tcp_check/tcp_check.py index b22255335a1a4..81bcd31590198 100644 --- a/tcp_check/datadog_checks/tcp_check/tcp_check.py +++ b/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] diff --git a/tcp_check/tests/common.py b/tcp_check/tests/common.py index 24715f5e8b9f7..432d927f84f35 100644 --- a/tcp_check/tests/common.py +++ b/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 = [ diff --git a/tcp_check/tests/test_tcp_check.py b/tcp_check/tests/test_tcp_check.py index b21321f88c6f4..8b4f62e5bca24 100644 --- a/tcp_check/tests/test_tcp_check.py +++ b/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,15 +127,21 @@ 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( @@ -142,6 +149,7 @@ def test_response_time(aggregator): 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,9 +191,11 @@ 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( @@ -190,6 +203,7 @@ def test_response_time(aggregator): 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)