from __future__ import annotations
import argparse
import os
from pathlib import Path
import signal
import socket
import struct
import subprocess
import tempfile
import threading
import time
from typing import Final
TEST_NAME: Final = "example.test"
TEST_ADDRESS: Final = "192.0.2.123"
LOOPBACK: Final = "127.0.0.1"
def question_end(packet: bytes) -> tuple[int, int, int]:
if len(packet) < 12:
raise ValueError("short DNS packet")
if struct.unpack_from("!H", packet, 4)[0] != 1:
raise ValueError("expected one DNS question")
offset = 12
labels: list[bytes] = []
while True:
if offset >= len(packet):
raise ValueError("truncated DNS name")
length = packet[offset]
offset += 1
if length == 0:
break
if length & 0xC0:
raise ValueError("compressed DNS question")
if length > 63 or offset + length > len(packet):
raise ValueError("invalid DNS label")
labels.append(packet[offset : offset + length])
offset += length
if offset + 4 > len(packet):
raise ValueError("truncated DNS question fields")
qtype, qclass = struct.unpack_from("!HH", packet, offset)
name = b".".join(labels).decode("ascii").lower()
if name != TEST_NAME:
raise ValueError(f"unexpected DNS name {name!r}")
return offset + 4, qtype, qclass
def make_response(query: bytes) -> bytes:
end, qtype, qclass = question_end(query)
identifier, query_flags = struct.unpack_from("!HH", query, 0)
response_flags = 0x8000 | 0x0080 | (query_flags & (0x0100 | 0x0010))
answers = 1 if qtype == 1 and qclass == 1 else 0
header = struct.pack("!HHHHHH", identifier, response_flags, 1, answers, 0, 0)
response = bytearray(header)
response.extend(query[12:end])
if answers:
response.extend(b"\xc0\x0c")
response.extend(struct.pack("!HHIH", 1, 1, 60, 4))
response.extend(socket.inet_aton(TEST_ADDRESS))
return bytes(response)
def read_exact(stream: socket.socket, length: int) -> bytes:
output = bytearray()
while len(output) < length:
chunk = stream.recv(length - len(output))
if not chunk:
raise ConnectionError("unexpected EOF")
output.extend(chunk)
return bytes(output)
class DeterministicUpstream:
def __init__(self) -> None:
self.stop = threading.Event()
self.udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
self.udp.bind((LOOPBACK, 0))
self.port = int(self.udp.getsockname()[1])
self.udp.settimeout(0.2)
self.tcp = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.tcp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
self.tcp.bind((LOOPBACK, self.port))
self.tcp.listen(16)
self.tcp.settimeout(0.2)
self.threads = [
threading.Thread(target=self._serve_udp, name="test-upstream-udp", daemon=True),
threading.Thread(target=self._serve_tcp, name="test-upstream-tcp", daemon=True),
]
def start(self) -> None:
for thread in self.threads:
thread.start()
def close(self) -> None:
self.stop.set()
self.udp.close()
self.tcp.close()
for thread in self.threads:
thread.join(timeout=2)
def _serve_udp(self) -> None:
while not self.stop.is_set():
try:
query, peer = self.udp.recvfrom(65535)
except socket.timeout:
continue
except OSError:
return
try:
response = make_response(query)
self.udp.sendto(response, peer)
except (OSError, ValueError):
continue
def _serve_tcp(self) -> None:
while not self.stop.is_set():
try:
client, _ = self.tcp.accept()
except socket.timeout:
continue
except OSError:
return
threading.Thread(
target=self._serve_tcp_client,
args=(client,),
name="test-upstream-tcp-client",
daemon=True,
).start()
@staticmethod
def _serve_tcp_client(client: socket.socket) -> None:
with client:
client.settimeout(5)
try:
while True:
length_bytes = client.recv(2)
if not length_bytes:
return
if len(length_bytes) != 2:
length_bytes += read_exact(client, 2 - len(length_bytes))
length = struct.unpack("!H", length_bytes)[0]
query = read_exact(client, length)
response = make_response(query)
client.sendall(struct.pack("!H", len(response)) + response)
except (ConnectionError, OSError, ValueError):
return
def dual_protocol_port() -> int:
for _ in range(100):
tcp = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
try:
tcp.bind((LOOPBACK, 0))
port = int(tcp.getsockname()[1])
udp.bind((LOOPBACK, port))
return port
except OSError:
continue
finally:
tcp.close()
udp.close()
raise RuntimeError("could not reserve a dual-protocol port")
def make_query(identifier: int, qtype: int = 1) -> bytes:
packet = bytearray(struct.pack("!HHHHHH", identifier, 0x0100, 1, 0, 0, 0))
for label in TEST_NAME.split("."):
encoded = label.encode("ascii")
packet.append(len(encoded))
packet.extend(encoded)
packet.append(0)
packet.extend(struct.pack("!HH", qtype, 1))
return bytes(packet)
def validate_answer(packet: bytes, identifier: int) -> None:
if len(packet) < 12:
raise AssertionError("short DNS response")
response_id, flags, qdcount, ancount = struct.unpack_from("!HHHH", packet, 0)
if response_id != identifier:
raise AssertionError("DNS transaction ID was not preserved")
if flags & 0x8000 == 0 or flags & 0x000F:
raise AssertionError(f"unexpected DNS response flags 0x{flags:04x}")
if qdcount != 1 or ancount < 1:
raise AssertionError("DNS answer is missing")
if socket.inet_aton(TEST_ADDRESS) not in packet:
raise AssertionError("expected A record is missing")
def query_udp(port: int, identifier: int) -> None:
with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as client:
client.settimeout(2)
client.sendto(make_query(identifier), (LOOPBACK, port))
response, _ = client.recvfrom(65535)
validate_answer(response, identifier)
def query_tcp(port: int, identifier: int) -> None:
query = make_query(identifier)
with socket.create_connection((LOOPBACK, port), timeout=2) as client:
client.settimeout(2)
client.sendall(struct.pack("!H", len(query)) + query)
length = struct.unpack("!H", read_exact(client, 2))[0]
response = read_exact(client, length)
validate_answer(response, identifier)
def wait_for_stub(process: subprocess.Popen[str], port: int) -> None:
deadline = time.monotonic() + 15
last_error: BaseException | None = None
while time.monotonic() < deadline:
if process.poll() is not None:
raise RuntimeError(f"resolver exited with status {process.returncode}")
try:
query_udp(port, 0x4100)
return
except (AssertionError, OSError) as error:
last_error = error
time.sleep(0.1)
raise RuntimeError(f"resolver did not become ready: {last_error}")
def terminate(process: subprocess.Popen[str]) -> None:
if process.poll() is not None:
return
process.send_signal(signal.SIGTERM)
try:
process.wait(timeout=10)
except subprocess.TimeoutExpired:
process.kill()
process.wait(timeout=5)
def run(binary: Path, resolvectl: Path) -> None:
if not binary.is_file() or not os.access(binary, os.X_OK):
raise FileNotFoundError(binary)
if not resolvectl.is_file() or not os.access(resolvectl, os.X_OK):
raise FileNotFoundError(resolvectl)
upstream = DeterministicUpstream()
upstream.start()
stub_port = dual_protocol_port()
proxy_port = dual_protocol_port()
try:
with tempfile.TemporaryDirectory(prefix="resolved-live-") as temporary:
root = Path(temporary)
run_dir = root / "run"
varlink = run_dir / "io.rustd.Resolve"
config = root / "resolved.conf"
log = root / "daemon.log"
config.write_text(
"[Resolve]\n"
f"DNS={LOOPBACK}:{upstream.port}\n"
"FallbackDNS=\n"
"DNSSEC=no\n"
"DNSOverTLS=no\n"
"LLMNR=no\n"
"MulticastDNS=no\n",
encoding="utf-8",
)
check = subprocess.run(
[
str(binary),
"--config",
str(config),
"--listen",
f"{LOOPBACK}:{stub_port}",
"--proxy-listen",
f"{LOOPBACK}:{proxy_port}",
"--runtime-directory",
str(run_dir),
"--varlink",
str(varlink),
"--workers",
"2",
"--no-dbus",
"--check-config",
],
text=True,
capture_output=True,
timeout=10,
check=True,
)
if "configuration is valid" not in check.stdout:
raise AssertionError("configuration validation output is missing")
with log.open("w", encoding="utf-8") as log_file:
process = subprocess.Popen(
[
str(binary),
"--config",
str(config),
"--listen",
f"{LOOPBACK}:{stub_port}",
"--proxy-listen",
f"{LOOPBACK}:{proxy_port}",
"--runtime-directory",
str(run_dir),
"--varlink",
str(varlink),
"--workers",
"2",
"--no-dbus",
],
stdout=log_file,
stderr=subprocess.STDOUT,
text=True,
)
try:
wait_for_stub(process, stub_port)
query_udp(stub_port, 0x4101)
query_tcp(stub_port, 0x4102)
query_udp(proxy_port, 0x4103)
query_tcp(proxy_port, 0x4104)
result = subprocess.run(
[str(resolvectl), "--socket", str(varlink), "query", TEST_NAME],
text=True,
capture_output=True,
timeout=15,
check=True,
)
if TEST_ADDRESS not in result.stdout:
raise AssertionError(
f"Varlink lookup did not return {TEST_ADDRESS}: {result.stdout!r}"
)
for name in ("stub-resolv.conf", "resolv.conf"):
path = run_dir / name
if not path.is_file() or not path.read_text(encoding="utf-8"):
raise AssertionError(f"runtime resolver file is missing: {path}")
except BaseException:
log_file.flush()
print(log.read_text(encoding="utf-8"), end="")
raise
finally:
terminate(process)
if process.returncode != 0:
log_file.flush()
print(log.read_text(encoding="utf-8"), end="")
raise RuntimeError(f"resolver exited with status {process.returncode}")
finally:
upstream.close()
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("binary", type=Path)
parser.add_argument("resolvectl", type=Path)
arguments = parser.parse_args()
run(arguments.binary.resolve(), arguments.resolvectl.resolve())
if __name__ == "__main__":
main()