heartbeat/hbd/client/dnsttl.py

"""Re-resolution of the server hostnames at the TTL of their A/AAAA records.

getaddrinfo() hides the TTL, so a minimal stub-resolver query (resolv.conf nameservers, search
list) is made only to learn it. The addresses themselves still come from getaddrinfo(). When
the TTL cannot be learned (IP literal, /etc/hosts name, no resolv.conf, resolver timeout)
DEFAULT_TTL is used.
"""

import ipaddress
import logging
import os
import socket
import struct
from typing import List, Optional, Tuple

logger = logging.getLogger("hbc.dns")

DEFAULT_TTL = 300       # used when the real TTL is unknown
MIN_TTL = 30            # a TTL of 0 must not turn into a resolution loop
FAILED_RETRY = 60       # re-try delay after a failed resolution
QUERY_TIMEOUT = 2.0

TYPE_A, TYPE_CNAME, TYPE_AAAA = 1, 5, 28


def parse_resolv_conf(path: str = "/etc/resolv.conf") -> Tuple[List[str], List[str]]:
    """Return (nameservers, search domains)."""
    servers: List[str] = []
    search: List[str] = []
    try:
        with open(path) as f:
            for line in f:
                parts = line.split()
                if len(parts) < 2:
                    continue
                if parts[0] == "nameserver":
                    servers.append(parts[1].split("%")[0])
                elif parts[0] in ("search", "domain"):
                    search = parts[1:]
    except OSError:
        pass
    return servers, search


def build_query(name: str, qtype: int, qid: int) -> bytes:
    out = struct.pack("!HHHHHH", qid, 0x0100, 1, 0, 0, 0)
    for label in name.rstrip(".").split("."):
        out += bytes([len(label)]) + label.encode("idna")
    return out + b"\x00" + struct.pack("!HH", qtype, 1)


def _skip_name(buf: bytes, off: int) -> int:
    while True:
        n = buf[off]
        if n == 0:
            return off + 1
        if n & 0xC0:
            return off + 2
        off += n + 1


def min_answer_ttl(resp: bytes, qid: int, qtype: int) -> Optional[int]:
    """Smallest TTL among the answer RRs (CNAMEs included), or None if there is no qtype answer."""
    try:
        rid, flags, qd, an = struct.unpack("!HHHH", resp[:8])
        if rid != qid or flags & 0x000F:
            return None
        off = 12
        for _ in range(qd):
            off = _skip_name(resp, off) + 4
        ttls: List[int] = []
        found = False
        for _ in range(an):
            off = _skip_name(resp, off)
            rtype, _cls, ttl, rdlen = struct.unpack("!HHIH", resp[off:off + 10])
            off += 10 + rdlen
            if rtype in (qtype, TYPE_CNAME):
                ttls.append(ttl)
            found = found or rtype == qtype
        return min(ttls) if found else None
    except (struct.error, IndexError):
        return None


def _query_ttl(server: str, name: str, qtype: int) -> Optional[int]:
    qid = struct.unpack("!H", os.urandom(2))[0]
    try:
        af = socket.AF_INET6 if ":" in server else socket.AF_INET
        with socket.socket(af, socket.SOCK_DGRAM) as s:
            s.settimeout(QUERY_TIMEOUT)
            s.sendto(build_query(name, qtype, qid), (server, 53))
            return min_answer_ttl(s.recv(4096), qid, qtype)
    except (OSError, UnicodeError):
        return None


def lookup_ttl(host: str, af_filter: int = 0) -> Optional[int]:
    """Blocking. Smallest A/AAAA TTL for host, or None when it cannot be determined."""
    try:
        ipaddress.ip_address(host)
        return None
    except ValueError:
        pass
    servers, search = parse_resolv_conf()
    if host.endswith("."):
        names = [host]
    elif "." in host:
        names = [host] + [f"{host}.{d}" for d in search]
    else:
        names = [f"{host}.{d}" for d in search] + [host]
    qtypes = [t for t, af in ((TYPE_A, socket.AF_INET), (TYPE_AAAA, socket.AF_INET6))
              if af_filter in (0, af)]
    for name in names:
        ttls = []
        for qtype in qtypes:
            for server in servers:
                ttl = _query_ttl(server, name, qtype)
                if ttl is not None:
                    ttls.append(ttl)
                    break
        if ttls:
            return min(ttls)
    return None


def refresh_delay(ttls: List[Optional[int]]) -> int:
    """Seconds until the next resolution: the smallest known TTL, floored at MIN_TTL."""
    known = [t for t in ttls if t is not None]
    return max(min(known), MIN_TTL) if known else DEFAULT_TTL


def resolve(hosts: List[str], port: int, af_filter: int) -> Optional[List[Tuple[int, str]]]:
    """Blocking. Every (family, address) the hosts resolve to, deduplicated, or None if any lookup fails."""
    found: List[Tuple[int, str]] = []
    for host in hosts:
        try:
            infos = socket.getaddrinfo(host, port, af_filter, 0, socket.SOL_UDP)
        except socket.gaierror as e:
            logger.warning(f"Cannot resolve {host}: {e}")
            return None
        for info in infos:
            entry = (info[0], str(info[4][0]))
            if entry not in found:
                found.append(entry)
    return found or None


def retarget(connections: list, resolved: List[Tuple[int, str]]) -> Tuple[List[Tuple[int, str]], list]:
    """Point the connections at the resolved addresses, in place.

    A connection whose address is gone is re-pointed at an address that is new, in list order, so
    the first connection (the one plugin data uses) survives as long as anything resolves.
    Returns (new addresses still needing a connection, connections to drop).
    """
    current = {c.addr for c in connections}
    wanted = {addr for _, addr in resolved}
    new = [(af, addr) for af, addr in resolved if addr not in current]
    gone = [c for c in connections if c.addr not in wanted]
    paired = min(len(gone), len(new))
    for conn, (af, addr) in zip(gone, new):
        logger.info(f"Server address {conn.addr} -> {addr}")
        conn.set_target(addr, af)
    drop = gone[paired:]
    if connections and connections[0] in drop:
        # plugin data goes out on the first connection: hand it a surviving address and drop that one instead
        keep = next(c for c in connections if c not in gone)
        connections[0].set_target(keep.addr, keep.af)
        drop.remove(connections[0])
        drop.append(keep)
    return new[paired:], drop