Decode hostname to string if passed as bytes

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)

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