Message.php

<?php

namespace Utopia\DNS;

use Utopia\DNS\Exception\Message\DecodingException;
use Utopia\DNS\Exception\Message\PartialDecodingException;
use Utopia\DNS\Message\Header;
use Utopia\DNS\Message\Question;
use Utopia\DNS\Message\Record;

final class Message
{
    public const int RCODE_NOERROR = 0;
    public const int RCODE_FORMERR = 1;
    public const int RCODE_SERVFAIL = 2;
    public const int RCODE_NXDOMAIN = 3;
    public const int RCODE_NOTIMP = 4;
    public const int RCODE_REFUSED = 5;
    public const int RCODE_YXDOMAIN = 6;
    public const int RCODE_YXRRSET = 7;
    public const int RCODE_NXRRSET = 8;
    public const int RCODE_NOTAUTH = 9;
    public const int RCODE_NOTZONE = 10;

    /**
     * @param Header $header The header of the message.
     * @param Question[] $questions The question records.
     * @param list<Record> $answers The answer records.
     * @param list<Record> $authority The authority records.
     * @param list<Record> $additional The additional records.
     */
    public function __construct(
        public readonly Header $header,
        /** @var Question[] */
        public readonly array $questions = [],
        /** @var list<Record> */
        public readonly array $answers = [],
        /** @var list<Record> */
        public readonly array $authority = [],
        /** @var list<Record> */
        public readonly array $additional = []
    ) {
        if ($header->questionCount !== count($questions)) {
            throw new \InvalidArgumentException('Invalid DNS response: question count mismatch');
        }
        if ($header->answerCount !== count($answers)) {
            throw new \InvalidArgumentException('Invalid DNS response: answer count mismatch');
        }
        if ($header->authorityCount !== count($authority)) {
            throw new \InvalidArgumentException('Invalid DNS response: authority count mismatch');
        }
        if ($header->additionalCount !== count($additional)) {
            throw new \InvalidArgumentException('Invalid DNS response: additional count mismatch');
        }
        $soaAuthorityCount = count(array_filter(
            $this->authority,
            fn ($record) => $record->type === Record::TYPE_SOA
        ));

        if ($header->isResponse && $header->authoritative && $soaAuthorityCount < 1) {
            if ($header->responseCode === self::RCODE_NXDOMAIN) {
                throw new \InvalidArgumentException('NXDOMAIN requires SOA in authority');
            }
            if ($header->responseCode === self::RCODE_NOERROR && $answers === []) {
                throw new \InvalidArgumentException('NODATA should include SOA in authority');
            }
        }
    }

    public static function query(
        Question $question,
        ?int $id = null,
        bool $recursionDesired = true
    ): self {
        if ($id === null) {
            $id = random_int(0, 0xFFFF);
        }

        $header = new Header(
            id: $id,
            isResponse: false,
            opcode: 0, // QUERY
            authoritative: false,
            truncated: false,
            recursionDesired: $recursionDesired,
            recursionAvailable: false,
            responseCode: 0,
            questionCount: 1,
            answerCount: 0,
            authorityCount: 0,
            additionalCount: 0
        );

        return new self($header, [$question]);
    }

    /**
     * Create a response message.
     *
     * @param Header $header The header of the query message to respond to.
     * @param int $responseCode The response code.
     * @param array<Question> $questions The question records.
     * @param list<Record> $answers The answer records.
     * @param list<Record> $authority The authority records.
     * @param list<Record> $additional The additional records.
     * @param bool $authoritative Whether the response is authoritative.
     * @param bool $truncated Whether the response is truncated.
     * @param bool $recursionAvailable Whether recursion is available.
     * @return self The response message.
     */
    public static function response(
        Header $header,
        int $responseCode,
        array $questions = [],
        array $answers = [],
        array $authority = [],
        array $additional = [],
        bool $authoritative = false,
        bool $truncated = false,
        bool $recursionAvailable = false
    ): self {
        $header = new Header(
            id: $header->id,
            isResponse: true,
            opcode: $header->opcode,
            authoritative: $authoritative,
            truncated: $truncated,
            recursionDesired: $header->recursionDesired,
            recursionAvailable: $recursionAvailable,
            responseCode: $responseCode,
            questionCount: count($questions),
            answerCount: count($answers),
            authorityCount: count($authority),
            additionalCount: count($additional)
        );


        return new self($header, $questions, $answers, $authority, $additional);
    }

    public static function decode(string $packet): self
    {
        if (strlen($packet) < Header::LENGTH) {
            throw new DecodingException('Invalid DNS response: header too short');
        }

        // --- Parse header (12 bytes) ---
        $header = Header::decode($packet);

        // --- Parse Question Section ---
        try {
            $offset = Header::LENGTH;
            $questions = [];
            for ($i = 0; $i < $header->questionCount; $i++) {
                $questions[] = Question::decode($packet, $offset);
            }

            // --- Decode Answer Section ---
            $answers = [];
            for ($i = 0; $i < $header->answerCount; $i++) {
                $answers[] = Record::decode($packet, $offset);
            }

            // --- Decode Authority Section ---
            $authority = [];
            for ($i = 0; $i < $header->authorityCount; $i++) {
                $authority[] = Record::decode($packet, $offset);
            }

            // --- Decode Additional Section ---
            $additional = [];
            for ($i = 0; $i < $header->additionalCount; $i++) {
                $additional[] = Record::decode($packet, $offset);
            }

            if ($offset !== strlen($packet)) {
                throw new DecodingException('Invalid packet length');
            }
        } catch (DecodingException $e) {
            throw new PartialDecodingException($header, $e->getMessage(), $e);
        }

        return new self($header, $questions, $answers, $authority, $additional);
    }

    /**
     * Encode the message to a binary DNS packet.
     *
     * When maxSize is specified, truncation follows RFC 1035 Section 6.2 and RFC 2181 Section 9:
     * - Truncation starts at the end and works forward (additional → authority → answers)
     * - TC flag is only set when required RRSets (answers) couldn't be fully included
     * - Complete RRSets are preserved; partial RRSets are omitted entirely
     * - Questions are always preserved
     *
     * @param int|null $maxSize Maximum packet size (e.g., 512 for UDP per RFC 1035)
     * @return string The encoded DNS packet
     */
    public function encode(?int $maxSize = null): string
    {
        // Build full packet first
        $packet = $this->header->encode();

        foreach ($this->questions as $question) {
            $packet .= $question->encode();
        }

        foreach ($this->answers as $answer) {
            $packet .= $answer->encode($packet);
        }

        foreach ($this->authority as $authority) {
            $packet .= $authority->encode($packet);
        }

        foreach ($this->additional as $additional) {
            $packet .= $additional->encode($packet);
        }

        // No truncation needed
        if ($maxSize === null || strlen($packet) <= $maxSize) {
            return $packet;
        }

        // RFC-compliant truncation: work backward from end
        // Per RFC 1035 Section 6.2 and RFC 2181 Section 9
        return $this->encodeWithTruncation($maxSize);
    }

    /**
     * Encode with RFC-compliant truncation strategy.
     *
     * Truncation order per RFC 1035 Section 6.2:
     * 1. Drop additional section first
     * 2. If still too big, drop authority section
     * 3. If still too big, include as many complete answer RRSets as fit, set TC
     *
     * TC flag is only set when answer section data is truncated (RFC 2181 Section 9).
     */
    private function encodeWithTruncation(int $maxSize): string
    {
        // Step 1: Try without additional section
        $withoutAdditional = self::response(
            $this->header,
            $this->header->responseCode,
            questions: $this->questions,
            answers: $this->answers,
            authority: $this->authority,
            additional: [],
            authoritative: $this->header->authoritative,
            truncated: false,
            recursionAvailable: $this->header->recursionAvailable
        );

        $packet = $withoutAdditional->encode();
        if (strlen($packet) <= $maxSize) {
            return $packet;
        }

        // Step 2: Try without authority section.
        // NODATA (NOERROR + no answers) and NXDOMAIN require SOA in authority per RFC;
        // when we drop authority for size, mark as non-authoritative so validation allows it.
        $isNodataOrNxdomain = ($this->header->responseCode === self::RCODE_NOERROR && $this->answers === [])
            || $this->header->responseCode === self::RCODE_NXDOMAIN;
        $withoutAuthority = self::response(
            $this->header,
            $this->header->responseCode,
            questions: $this->questions,
            answers: $this->answers,
            authority: [],
            additional: [],
            authoritative: $isNodataOrNxdomain ? false : $this->header->authoritative,
            truncated: false,
            recursionAvailable: $this->header->recursionAvailable
        );

        $packet = $withoutAuthority->encode();
        if (strlen($packet) <= $maxSize) {
            return $packet;
        }

        // Step 3: Truncate answers - find how many complete records fit
        // Build base packet with header + questions
        $basePacket = $this->header->encode();
        foreach ($this->questions as $question) {
            $basePacket .= $question->encode();
        }

        $fittingAnswers = [];
        $tempPacket = $basePacket;

        foreach ($this->answers as $answer) {
            $encodedAnswer = $answer->encode($tempPacket);
            if (strlen($tempPacket) + strlen($encodedAnswer) <= $maxSize) {
                $tempPacket .= $encodedAnswer;
                $fittingAnswers[] = $answer;
            } else {
                // This answer doesn't fit, stop here
                break;
            }
        }

        // Determine if we need to set TC flag
        // Per RFC 2181 Section 9: TC is set only when required RRSet data couldn't fit
        $needsTruncation = count($fittingAnswers) < count($this->answers);

        // When authority is empty (dropped for truncation), NODATA/NXDOMAIN must be non-authoritative
        $isNodataOrNxdomainTruncated = ($this->header->responseCode === self::RCODE_NOERROR && $fittingAnswers === [])
            || $this->header->responseCode === self::RCODE_NXDOMAIN;

        $truncatedResponse = self::response(
            $this->header,
            $this->header->responseCode,
            questions: $this->questions,
            answers: $fittingAnswers,
            authority: [],
            additional: [],
            authoritative: $isNodataOrNxdomainTruncated ? false : $this->header->authoritative,
            truncated: $needsTruncation,
            recursionAvailable: $this->header->recursionAvailable
        );

        return $truncatedResponse->encode();
    }
}