certificate_retrieval.py

"""Unauthenticated certificate retrieval and its anti-amplification rule."""

from __future__ import annotations

from typing import Sequence

from constants import (
    CERTIFICATE_RECORD_TTL,
    DNS_CLASS_IN,
    DNS_HEADER_SIZE,
    DNS_TYPE_OPT,
    DNS_TYPE_TXT,
    EDNS_PADDING_OPTION_CODE,
    EDNS_UDP_PAYLOAD_SIZE,
)
from errors import WireFormatError

__all__ = [
    "build_certificate_response",
    "certificate_query",
    "serve_certificates",
]


def _dns_name(name: str) -> bytes:
    """Encode a domain name as a sequence of length-prefixed labels."""

    wire = b""
    for label in name.rstrip(".").split("."):
        encoded = label.encode("ascii")
        if not 1 <= len(encoded) <= 63:
            raise ValueError("each DNS label must be 1 to 63 bytes")
        wire += bytes([len(encoded)]) + encoded
    return wire + b"\x00"


def _edns_padding(padding_len: int) -> bytes:
    """Build an EDNS(0) OPT record carrying a padding option (RFC 7830)."""

    option = (
        EDNS_PADDING_OPTION_CODE.to_bytes(2, "big")
        + padding_len.to_bytes(2, "big")
        + b"\x00" * padding_len
    )
    return (
        b"\x00"  # root owner name
        + DNS_TYPE_OPT.to_bytes(2, "big")
        + EDNS_UDP_PAYLOAD_SIZE.to_bytes(2, "big")
        + b"\x00\x00\x00\x00"  # extended rcode, version, flags
        + len(option).to_bytes(2, "big")
        + option
    )


def _question_end(packet: bytes) -> int:
    """Return the offset just past the single question of a DNS message."""

    if len(packet) < DNS_HEADER_SIZE:
        raise WireFormatError("DNS message is shorter than its header")
    if packet[4:6] != b"\x00\x01":
        raise WireFormatError("a certificate query carries exactly one question")
    offset = DNS_HEADER_SIZE
    while True:
        if offset >= len(packet):
            raise WireFormatError("DNS question name is truncated")
        label_len = packet[offset]
        if label_len & 0xC0:
            raise WireFormatError("unexpected compression pointer in question name")
        offset += 1 + label_len
        if label_len == 0:
            break
    if offset + 4 > len(packet):
        raise WireFormatError("DNS question is missing its type and class")
    return offset + 4


def _txt_rdata(certificate: bytes) -> bytes:
    """Split a certificate into TXT character-strings of at most 255 bytes."""

    chunks = [certificate[i : i + 255] for i in range(0, len(certificate), 255)]
    return b"".join(bytes([len(chunk)]) + chunk for chunk in chunks)


def _txt_answer(certificate: bytes, ttl: int) -> bytes:
    """Encode one certificate as a TXT answer record."""

    rdata = _txt_rdata(certificate)
    return (
        b"\xc0\x0c"  # owner name: compression pointer to the question
        + DNS_TYPE_TXT.to_bytes(2, "big")
        + DNS_CLASS_IN.to_bytes(2, "big")
        + ttl.to_bytes(4, "big")
        + len(rdata).to_bytes(2, "big")
        + rdata
    )


def certificate_query(
    provider_name: str,
    query_id: bytes = b"\x00\x00",
    padded_length: int = 0,
) -> bytes:
    """Build an unencrypted TXT certificate query, optionally EDNS(0)-padded.

    A client that wants the larger PQ certificates over UDP pads the query to at
    least the expected response size, so the response stays within the request and
    passes the anti-amplification check at the resolver and at any relay.
    """

    question = (
        _dns_name(provider_name)
        + DNS_TYPE_TXT.to_bytes(2, "big")
        + DNS_CLASS_IN.to_bytes(2, "big")
    )
    additional = b""
    arcount = b"\x00\x00"
    if padded_length > 0:
        opt_overhead = 15  # OPT record header (11) and padding option header (4)
        padding = max(0, padded_length - DNS_HEADER_SIZE - len(question) - opt_overhead)
        additional = _edns_padding(padding)
        arcount = b"\x00\x01"
    header = (
        query_id
        + b"\x01\x00"  # standard query, recursion desired
        + b"\x00\x01"  # qdcount
        + b"\x00\x00"  # ancount
        + b"\x00\x00"  # nscount
        + arcount
    )
    return header + question + additional


def build_certificate_response(
    request: bytes,
    certificates: Sequence[bytes],
    truncated: bool = False,
    ttl: int = CERTIFICATE_RECORD_TTL,
) -> bytes:
    """Assemble a TXT certificate response that echoes the request's question.

    Each certificate becomes one TXT answer record. The request's OPT record and
    padding are not echoed, exactly as a resolver answers a plain TXT query.
    """

    question = request[DNS_HEADER_SIZE : _question_end(request)]
    answers = b"".join(_txt_answer(certificate, ttl) for certificate in certificates)
    flags = 0x8080  # QR and RA, as sent by deployed resolvers
    flags |= int.from_bytes(request[2:4], "big") & 0x0110  # preserve RD and CD
    if truncated:
        flags |= 0x0200  # TC: the full answer did not fit
    header = (
        request[0:2]
        + flags.to_bytes(2, "big")
        + b"\x00\x01"  # qdcount
        + len(certificates).to_bytes(2, "big")  # ancount
        + b"\x00\x00"  # nscount
        + b"\x00\x00"  # arcount
    )
    return header + question + answers


def serve_certificates(
    request: bytes,
    classical_certificates: Sequence[bytes],
    pq_certificates: Sequence[bytes],
    over_udp: bool,
) -> bytes:
    """Answer a certificate query under the anti-amplification rule.

    The classical certificate is always returned for compatibility with deployed
    DNSCrypt v2 clients and resolvers. The large PQ certificates are added over
    UDP only when the complete response still fits within the request that
    triggered it, because otherwise a spoofed query could amplify traffic. When
    the PQ certificates do not fit, the response carries the classical
    certificates with the TC bit set and a PQ-capable client retries over TCP.
    Over TCP the handshake validates the source, so all certificates are always
    included.
    """

    full = build_certificate_response(
        request, [*classical_certificates, *pq_certificates]
    )
    if not over_udp or not pq_certificates or len(full) <= len(request):
        return full
    return build_certificate_response(request, classical_certificates, truncated=True)