Allow homebrew store DNS lookups

docker/guide_gateway.py

@@ -15,6 +15,17 @@ HTTPS_PORT = int(os.environ.get("HTTPS_PORT", "8443"))
 GATEWAY_IP = ipaddress.IPv4Address(os.environ.get("GATEWAY_IP", "192.168.20.30"))
 REDIRECT_URL = os.environ.get("REDIRECT_URL", f"http://{GATEWAY_IP}:8000/")
 TARGET_DOMAIN = "manuals.playstation.net"
+# Exact names forwarded to the upstream resolver (homebrew store: catalog + downloads).
+FORWARDED_DOMAINS = frozenset(
+    name.strip().lower()
+    for name in os.environ.get(
+        "FORWARDED_DOMAINS",
+        "homebrew.page,github.com,release-assets.githubusercontent.com",
+    ).split(",")
+    if name.strip()
+)
+UPSTREAM_DNS = str(ipaddress.ip_address(os.environ.get("UPSTREAM_DNS", "1.1.1.1")))
+UPSTREAM_TIMEOUT = 3.0
 CERTIFICATE = "/srv/certificates/manuals.crt"
 PRIVATE_KEY = "/srv/certificates/manuals.key"
 
@@ -43,7 +54,54 @@ def parse_question(packet):
     return ".".join(labels).lower(), qtype, qclass, offset + 4
 
 
-def dns_response(packet):
+def error_response(packet, rcode):
+    transaction_id, request_flags = struct.unpack_from("!HH", packet, 0)
+    response_flags = 0x8000 | (request_flags & 0x0100) | rcode
+    try:
+        _, _, _, question_end = parse_question(packet)
+        question = packet[12:question_end]
+        question_count = 1
+    except (UnicodeDecodeError, ValueError):
+        question = b""
+        question_count = 0
+    return (
+        struct.pack("!HHHHHH", transaction_id, response_flags, question_count, 0, 0, 0)
+        + question
+    )
+
+
+def receive_exactly(connection, wanted):
+    chunks = []
+    received = 0
+    while received < wanted:
+        chunk = connection.recv(wanted - received)
+        if not chunk:
+            raise OSError("upstream DNS connection closed early")
+        chunks.append(chunk)
+        received += len(chunk)
+    return b"".join(chunks)
+
+
+def forward_query(packet, use_tcp=False):
+    family = socket.AF_INET6 if ":" in UPSTREAM_DNS else socket.AF_INET
+    socket_type = socket.SOCK_STREAM if use_tcp else socket.SOCK_DGRAM
+    with socket.socket(family, socket_type) as upstream:
+        upstream.settimeout(UPSTREAM_TIMEOUT)
+        upstream.connect((UPSTREAM_DNS, 53))
+        if use_tcp:
+            upstream.sendall(struct.pack("!H", len(packet)) + packet)
+            response_length = struct.unpack("!H", receive_exactly(upstream, 2))[0]
+            response = receive_exactly(upstream, response_length)
+        else:
+            upstream.send(packet)
+            response = upstream.recv(65535)
+
+    if len(response) < 12 or response[:2] != packet[:2]:
+        raise ValueError("invalid upstream DNS response")
+    return response
+
+
+def dns_response(packet, use_tcp=False):
     if len(packet) < 12:
         return b""
 
@@ -57,6 +115,13 @@ def dns_response(packet):
         return struct.pack("!HHHHHH", transaction_id, response_flags | 1, 0, 0, 0, 0)
 
     question = packet[12:question_end]
+    if name in FORWARDED_DOMAINS:
+        try:
+            return forward_query(packet, use_tcp=use_tcp)
+        except (OSError, ValueError) as error:
+            print(f"DNS upstream failure for {name}: {error}", flush=True)
+            return error_response(packet, 2)  # SERVFAIL
+
     if name != TARGET_DOMAIN:
         return struct.pack("!HHHHHH", transaction_id, response_flags | 3, 1, 0, 0, 0) + question
 
@@ -109,20 +174,15 @@ class TCPDNSHandler(socketserver.BaseRequestHandler):
             print(f"DNS TCP {self.client_address[0]} {name} type={qtype}", flush=True)
         except (UnicodeDecodeError, ValueError):
             print(f"DNS TCP {self.client_address[0]} malformed", flush=True)
-        response = dns_response(packet)
+        response = dns_response(packet, use_tcp=True)
         if response:
             self.request.sendall(struct.pack("!H", len(response)) + response)
 
     def _receive_exactly(self, wanted):
-        chunks = []
-        received = 0
-        while received < wanted:
-            chunk = self.request.recv(wanted - received)
-            if not chunk:
-                return None
-            chunks.append(chunk)
-            received += len(chunk)
-        return b"".join(chunks)
+        try:
+            return receive_exactly(self.request, wanted)
+        except OSError:
+            return None
 
 
 class RedirectHandler(http.server.BaseHTTPRequestHandler):
@@ -187,7 +247,8 @@ if __name__ == "__main__":
     serve(https, "https-redirect")
 
     print(
-        f"PS5 guide gateway: {TARGET_DOMAIN} -> {GATEWAY_IP}; redirect -> {REDIRECT_URL}",
+        f"PS5 guide gateway: {TARGET_DOMAIN} -> {GATEWAY_IP}; redirect -> {REDIRECT_URL}; "
+        f"forward {','.join(sorted(FORWARDED_DOMAINS))} -> {UPSTREAM_DNS}",
         flush=True,
     )