diff --git a/aikido_zen/sinks/socket/__init__.py b/aikido_zen/sinks/socket/__init__.py index 574f2483d..64b1244d9 100644 --- a/aikido_zen/sinks/socket/__init__.py +++ b/aikido_zen/sinks/socket/__init__.py @@ -18,6 +18,9 @@ def _getaddrinfo_after(func, instance, args, kwargs, return_value): host = get_argument(args, kwargs, 0, "host") port = get_argument(args, kwargs, 1, "port") + if isinstance(host, bytes): + host = host.decode("utf-8", errors="replace") + # We want a normalized hostname for reporting & blocking outbound domains # This function decodes the hostname if its written in punycode hostname = normalize_hostname(host) diff --git a/aikido_zen/sinks/tests/socket_test.py b/aikido_zen/sinks/tests/socket_test.py index 6e8bfa45e..352f1b289 100644 --- a/aikido_zen/sinks/tests/socket_test.py +++ b/aikido_zen/sinks/tests/socket_test.py @@ -2,8 +2,10 @@ Test module for socket sink """ +import asyncio import socket import pytest +import httpx from unittest.mock import patch, MagicMock import aikido_zen.sinks.socket # Import to ensure patching from aikido_zen.context import current_context @@ -20,25 +22,65 @@ def run_around_tests(): current_context.set(None) -def test_socket_getaddrinfo_no_blocking(): +@pytest.mark.parametrize( + ("host", "expected_hostname"), + [ + pytest.param("localhost", "localhost", id="str"), + pytest.param(b"localhost", "localhost", id="bytes"), + pytest.param( + b"xn--ssrf-rdirects-ghb.testssandbox.com", + "ssrf-rédirects.testssandbox.com", + id="idn-bytes", + ), + ], +) +def test_socket_getaddrinfo_no_blocking(host, expected_hostname): """Test that getaddrinfo works normally when no blocking is configured""" # Reset cache to ensure clean state get_cache().reset() # Test that allowed domain doesn't throw an error try: - socket.getaddrinfo("localhost", 80) + socket.getaddrinfo(host, 80) except Exception: pytest.fail("getaddrinfo should not throw an error for allowed domains") # Verify hostname was tracked hostnames = get_cache().hostnames.as_array() assert len(hostnames) == 1 - assert hostnames[0]["hostname"] == "localhost" + assert hostnames[0]["hostname"] == expected_hostname assert hostnames[0]["port"] == 80 assert hostnames[0]["hits"] == 1 +@pytest.mark.asyncio +async def test_httpx_async_client_tracks_hostname_as_string(): + async def handle_request(reader, writer): + await reader.readuntil(b"\r\n\r\n") + writer.write( + b"HTTP/1.1 204 No Content\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ) + await writer.drain() + writer.close() + await writer.wait_closed() + + server = await asyncio.start_server(handle_request, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + get_cache().reset() + + try: + async with httpx.AsyncClient(trust_env=False) as client: + response = await client.get(f"http://localhost:{port}") + assert response.status_code == 204 + finally: + server.close() + await server.wait_closed() + + assert get_cache().hostnames.as_array() == [ + {"hostname": "localhost", "port": port, "hits": 1} + ] + + def test_socket_getaddrinfo_block_specific_domain(): """Test that getaddrinfo raises exception when specific domain is blocked""" # Reset cache and set up blocking for specific domain