From aed18e807175afddf87bb838ebd3eb0c1935cd7b Mon Sep 17 00:00:00 2001 From: Alexander Vasin Date: Tue, 14 Sep 2021 05:54:10 +0300 Subject: [PATCH] Add ipv6 support (#339) * Add ipv6 support Signed-off-by: Alexander Vasin * Increase test coverage Signed-off-by: Yuri Shkuro Co-authored-by: Yuri Shkuro Co-authored-by: Yuri Shkuro --- jaeger_client/TUDPTransport.py | 25 +++++++++++++++++----- tests/test_TUDPTransport.py | 38 ++++++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/jaeger_client/TUDPTransport.py b/jaeger_client/TUDPTransport.py index e9a2d4b8..655801da 100644 --- a/jaeger_client/TUDPTransport.py +++ b/jaeger_client/TUDPTransport.py @@ -17,6 +17,7 @@ from thrift.transport.TTransport import TTransportBase import socket + logger = logging.getLogger('jaeger_tracing') @@ -25,16 +26,30 @@ class TUDPTransport(TTransportBase, object): TUDPTransport implements just enough of the tornado transport interface to work for blindly sending UDP packets. """ + + DEFAULT_SOCKET_FAMILY = socket.AF_INET + def __init__(self, host, port, blocking=False): self.transport_host = host self.transport_port = port - self.transport_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - if blocking: - blocking = 1 - else: - blocking = 0 + + self.transport_sock = self._create_socket() self.transport_sock.setblocking(blocking) + def _create_socket(self) -> socket.socket: + family, type, proto = (self.DEFAULT_SOCKET_FAMILY, socket.SOCK_DGRAM, socket.IPPROTO_UDP) + + try: + addrinfo = socket.getaddrinfo( + self.transport_host, self.transport_port, type=socket.SOCK_DGRAM + ) + if addrinfo: + family, type, proto, *_ = addrinfo[0] + except socket.gaierror: + pass + + return socket.socket(family, type, proto) + def write(self, buf): """Raw write to the UDP socket.""" return self.transport_sock.sendto( diff --git a/tests/test_TUDPTransport.py b/tests/test_TUDPTransport.py index a1dfc974..d5c7c2cc 100644 --- a/tests/test_TUDPTransport.py +++ b/tests/test_TUDPTransport.py @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import socket import unittest +from unittest import mock + +import pytest from jaeger_client.TUDPTransport import TUDPTransport @@ -44,3 +48,37 @@ def test_close(self): with self.assertRaises(Exception): # Something bad should happen if we send on a closed socket.. self.t.write(b'hello') + + +def test_created_socket_default_family(): + transport = TUDPTransport('unknown-host', 12345) + sock = transport._create_socket() + assert sock.family == TUDPTransport.DEFAULT_SOCKET_FAMILY + + +@pytest.mark.parametrize('addrinfo,expected_family', ( + ( + ( + socket.AddressFamily.AF_INET6, socket.SocketKind.SOCK_DGRAM, + 17, '', ('aced:a11:7e57', 12345, 0, 0) + ), + socket.AF_INET6 + ), + ( + ( + socket.AddressFamily.AF_INET, socket.SocketKind.SOCK_DGRAM, + 17, '', ('127.0.0.1', 12345) + ), + socket.AF_INET + ), + ( + None, + TUDPTransport.DEFAULT_SOCKET_FAMILY + ) +)) +def test_created_socket_specified_family(addrinfo, expected_family): + return_value = [addrinfo] if addrinfo else [] + with mock.patch('socket.getaddrinfo', return_value=return_value): + transport = TUDPTransport('ipv6-host', 12345) + sock = transport._create_socket() + assert sock.family == expected_family