vpn-loopback-firewall/channel/proxy_dns.py

"""L2 本地 DNS 分流代理(软路由侧)。

接入点: 软路由 smartdns/dnsmasq 的上游指向本代理(127.0.0.1#8053),
全量 DNS 查询先到这里, 按域名分流:
  - 内网域名(配置 internal_suffixes, 如 .internal.example)-> 本地转发权威 DNS
    (内网记录只有本地权威 DNS 有, 云端解析不到, 认证域名另走 hosts)
  - 其余公网域名 -> 走云通道, 云端递归解析(223.5.5.5)
query 全程原样透传(不改 ID), 响应沿原路返回。

2026-09-10 降级加固(此前云服卡死时通道断, 本代理静默丢包 -> 全屋 DNS 干等超时):
  1. 云通道不可用 -> 直连 fallback_dns 解析(默认 school_dns), DNS 不断;
  2. 直连也失败 -> 回 SERVFAIL, 客户端立即切备用 DNS 而不是干等超时;
  3. 通道异常时的重连转后台异步(限频), 不在 DNS 查询路径上阻塞;
  4. 失败/降级计数用于观测(每 50 次降级打一条日志)。
"""
from __future__ import annotations

import socket
import sys
import threading
import time
from typing import Any

from .client import ChannelClient, ChannelError
from .dnsmsg import parse_qname

DNS_TIMEOUT = 4.0
FALLBACK_TIMEOUT = 3.0
RECONNECT_MIN_INTERVAL = 10.0


def _suffix_match(name: str, suffixes: list[str]) -> bool:
    for suffix in suffixes:
        s = suffix.strip().lower()
        if not s:
            continue
        if s.startswith("."):
            if name == s[1:] or name.endswith(s):
                return True
        elif name == s or name.endswith("." + s):
            return True
    return False


def _servfail(query: bytes) -> bytes:
    """构造 SERVFAIL 响应(保留 ID 与 question 段), 让客户端立即换备用 DNS。"""
    if len(query) < 12:
        return b""
    head = bytearray(query[:12])
    head[2] = head[2] | 0x80  # QR=1(保留 opcode/RD)
    head[3] = (head[3] & 0xF0) | 0x02  # RCODE=2 SERVFAIL
    head[6:8] = b"\x00\x00"
    head[8:10] = b"\x00\x00"
    head[10:12] = b"\x00\x00"
    i = 12
    while i < len(query) and query[i] != 0:
        i += 1 + int(query[i])
    if i >= len(query):
        return bytes(head)
    end = i + 5  # qname 终止符 + qtype(2) + qclass(2)
    if end > len(query):
        return bytes(head)
    return bytes(head) + query[12:end]


class DnsSplitProxy:
    """UDP 53 分流代理(带通道降级)。"""

    def __init__(
        self,
        client: ChannelClient,
        listen_port: int,
        school_dns: str,
        internal_suffixes: list[str],
        fallback_dns: str = "",
    ) -> None:
        if listen_port < 1 or listen_port > 65535:
            raise ValueError(f"bad listen port: {listen_port}")
        self.client = client
        self.port = listen_port
        self.school_dns = school_dns
        self.internal_suffixes = internal_suffixes
        self.fallback_dns = fallback_dns or school_dns
        self.local_resolved = 0
        self.channel_resolved = 0
        self.fallback_resolved = 0
        self.failed = 0
        self._reconnect_lock = threading.Lock()
        self._last_reconnect = 0.0

    def serve(self) -> None:
        srv = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        srv.bind(("0.0.0.0", self.port))
        print(
            f"[dns-proxy] listen 0.0.0.0:{self.port} "
            f"school_dns={self.school_dns} fallback={self.fallback_dns}",
            flush=True,
        )
        while True:
            data, client_addr = srv.recvfrom(4096)
            threading.Thread(
                target=self._handle,
                args=(srv, data, client_addr),
                name="dns-proxy-query",
                daemon=True,
            ).start()

    def _handle(
        self,
        srv: socket.socket,
        query: bytes,
        client_addr: tuple[str, int],
    ) -> None:
        name = parse_qname(query)
        try:
            if name is not None and _suffix_match(name, self.internal_suffixes):
                self._local_forward(srv, query, client_addr)
                self.local_resolved += 1
            else:
                self._channel_forward(srv, query, client_addr)
        except OSError:
            self.failed += 1

    def _local_forward(
        self, srv: socket.socket, query: bytes, client_addr: tuple[str, int]
    ) -> None:
        """直连本地权威 DNS(内网, 不走通道)。"""
        resp = self._direct_query(query, self.school_dns, DNS_TIMEOUT)
        if resp is not None:
            srv.sendto(resp, client_addr)

    def _channel_forward(
        self, srv: socket.socket, query: bytes, client_addr: tuple[str, int]
    ) -> None:
        resp = self._channel_send(query)
        if resp is not None:
            self.channel_resolved += 1
            srv.sendto(resp, client_addr)
            return
        # 云通道不可用 -> 降级直连上游, 保证全屋 DNS 不断
        resp = self._direct_query(query, self.fallback_dns, FALLBACK_TIMEOUT)
        if resp is not None:
            self.fallback_resolved += 1
            if self.fallback_resolved == 1 or self.fallback_resolved % 50 == 0:
                print(
                    f"[dns-proxy] 云通道不可用, 降级直连 {self.fallback_dns} "
                    f"(累计 {self.fallback_resolved} 次)",
                    flush=True,
                )
            srv.sendto(resp, client_addr)
            return
        # 直连也失败: 回 SERVFAIL, 让客户端立即切备用 DNS
        self.failed += 1
        sf = _servfail(query)
        if sf:
            srv.sendto(sf, client_addr)

    def _direct_query(
        self, query: bytes, server: str, timeout: float
    ) -> bytes | None:
        try:
            up = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            try:
                up.settimeout(timeout)
                up.sendto(query, (server, 53))
                resp, _ = up.recvfrom(4096)
                return resp
            finally:
                up.close()
        except OSError:
            return None

    def _channel_send(self, query: bytes) -> bytes | None:
        try:
            return self.client.send_dns_query(query)
        except ChannelError:
            # 不在此阻塞重连(会把 DNS 查询卡住): 后台限频重连, 本次直接降级
            self._trigger_reconnect()
            return None

    def _trigger_reconnect(self) -> None:
        now = time.time()
        with self._reconnect_lock:
            if now - self._last_reconnect < RECONNECT_MIN_INTERVAL:
                return
            self._last_reconnect = now
        threading.Thread(
            target=self._safe_reconnect, name="dns-reconnect", daemon=True
        ).start()

    def _safe_reconnect(self) -> None:
        try:
            self.client.ensure_connected()
        except ChannelError:
            pass


def run_dns_proxy(cfg: dict[str, Any]) -> None:
    cloud: dict[str, Any] = cfg["cloud"]
    proxy_cfg: dict[str, Any] = cfg["dns_proxy"]
    client = ChannelClient(
        host=str(cloud["host"]),
        port=int(cloud["port"]),
        token=str(cloud["token"]),
        ca_path=str(cloud.get("ca_path", "")),
        heartbeat=float(cfg.get("heartbeat", 30)),
    )
    client.connect()
    DnsSplitProxy(
        client,
        listen_port=int(proxy_cfg["listen_port"]),
        school_dns=str(proxy_cfg["school_dns"]),
        internal_suffixes=[str(s) for s in proxy_cfg.get("internal_suffixes", [])],
        fallback_dns=str(proxy_cfg.get("fallback_dns", "")),
    ).serve()


def main() -> int:
    if len(sys.argv) < 2:
        print("usage: python -m channel.proxy_dns <config.json>", file=sys.stderr)
        return 2
    import json
    from pathlib import Path

    cfg = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8"))
    run_dns_proxy(cfg)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())