From a04783d22bf5121ac58555faf0802f94e745a980 Mon Sep 17 00:00:00 2001
From: Hugo Herter <git@hugoherter.com>
Date: Wed, 14 Feb 2024 12:10:17 +0100
Subject: [PATCH] Fix: `socket.getaddrinfo` does not always return 2 values

A user reported that this function crashed on his host due to `ValueError: too many values to unpack (expected 2)`.

Solution: Process each returned tuple instead of assuming that two of them are always returned.

This fixes the issue both in the orchestrator and in the diagnostic VM.
---
 examples/example_fastapi/main.py              | 35 ++++++++++++++-----
 .../vm/orchestrator/views/host_status.py      | 29 ++++++++++++---
 2 files changed, 51 insertions(+), 13 deletions(-)

diff --git a/examples/example_fastapi/main.py b/examples/example_fastapi/main.py
index c7c5fc161..b6861b307 100644
--- a/examples/example_fastapi/main.py
+++ b/examples/example_fastapi/main.py
@@ -7,13 +7,14 @@
 from datetime import datetime
 from os import listdir
 from pathlib import Path
-from typing import Dict
+from typing import Dict, Optional
 
 import aiohttp
 from fastapi import FastAPI
 from fastapi.responses import PlainTextResponse
 from pip._internal.operations.freeze import freeze
 from pydantic import BaseModel
+from starlette.responses import JSONResponse
 
 from aleph.sdk.chains.remote import RemoteAccount
 from aleph.sdk.client import AlephClient, AuthenticatedAlephClient
@@ -91,13 +92,31 @@ async def read_aleph_messages():
 @app.get("/dns")
 async def resolve_dns_hostname():
     """Check if DNS resolution is working."""
-    info_inet, info_inet6 = socket.getaddrinfo("example.org", 80, proto=socket.IPPROTO_TCP)
-    ipv4 = info_inet[4][0]
-    ipv6 = info_inet6[4][0]
-    return {
-        "ipv4": ipv4,
-        "ipv6": ipv6,
-    }
+    hostname = "example.org"
+    ipv4: Optional[str] = None
+    ipv6: Optional[str] = None
+
+    info = socket.getaddrinfo(hostname, 80, proto=socket.IPPROTO_TCP)
+    if not info:
+        logger.error("DNS resolution failed")
+
+    # Iterate over the results to find the IPv4 and IPv6 addresses they may not all be present.
+    # The function returns a list of 5-tuples with the following structure:
+    # (family, type, proto, canonname, sockaddr)
+    for info_tuple in info:
+        if info_tuple[0] == socket.AF_INET:
+            ipv4 = info_tuple[4][0]
+        elif info_tuple[0] == socket.AF_INET6:
+            ipv6 = info_tuple[4][0]
+
+    if ipv4 and not ipv6:
+        logger.warning(f"DNS resolution for {hostname} returned only an IPv4 address")
+    elif ipv6 and not ipv4:
+        logger.warning(f"DNS resolution for {hostname} returned only an IPv6 address")
+
+    result = {"ipv4": ipv4, "ipv6": ipv6}
+    status_code = 200 if len(info) > 1 else 503
+    return JSONResponse(content=result, status_code=status_code)
 
 
 @app.get("/ip/address")
diff --git a/src/aleph/vm/orchestrator/views/host_status.py b/src/aleph/vm/orchestrator/views/host_status.py
index 5ab80bebc..b429a1e2d 100644
--- a/src/aleph/vm/orchestrator/views/host_status.py
+++ b/src/aleph/vm/orchestrator/views/host_status.py
@@ -1,6 +1,6 @@
 import logging
 import socket
-from typing import Any, Awaitable, Callable, Tuple
+from typing import Any, Awaitable, Callable, Optional, Tuple
 
 import aiohttp
 
@@ -45,10 +45,29 @@ async def check_host_egress_ipv6() -> bool:
     return await check_ip_connectivity(settings.CONNECTIVITY_IPV6_URL)
 
 
-async def resolve_dns(hostname: str) -> Tuple[str, str]:
-    info_inet, info_inet6 = socket.getaddrinfo(hostname, 80, proto=socket.IPPROTO_TCP)
-    ipv4 = info_inet[4][0]
-    ipv6 = info_inet6[4][0]
+async def resolve_dns(hostname: str) -> Tuple[Optional[str], Optional[str]]:
+    """Resolve a hostname to an IPv4 and IPv6 address."""
+    ipv4: Optional[str] = None
+    ipv6: Optional[str] = None
+
+    info = socket.getaddrinfo(hostname, 80, proto=socket.IPPROTO_TCP)
+    if not info:
+        logger.error("DNS resolution failed")
+
+    # Iterate over the results to find the IPv4 and IPv6 addresses they may not all be present.
+    # The function returns a list of 5-tuples with the following structure:
+    # (family, type, proto, canonname, sockaddr)
+    for info_tuple in info:
+        if info_tuple[0] == socket.AF_INET:
+            ipv4 = info_tuple[4][0]
+        elif info_tuple[0] == socket.AF_INET6:
+            ipv6 = info_tuple[4][0]
+
+    if ipv4 and not ipv6:
+        logger.warning(f"DNS resolution for {hostname} returned only an IPv4 address")
+    elif ipv6 and not ipv4:
+        logger.warning(f"DNS resolution for {hostname} returned only an IPv6 address")
+
     return ipv4, ipv6