Refactor Message::encode() to unify truncation path

Collapse encode() and encodeWithTruncation() into a single linear pass
that streams answers with partial truncation, then attempts authority
and additional all-or-nothing, and only rebuilds the header when counts
or the TC flag actually change. Add tests covering section priority,
answer-truncation clearing populated non-answer sections, TC preservation
on re-encode, the exact maxSize boundary, and NODATA-style drops.

Author: loks0n

src/DNS/Message.php

@@ -186,131 +186,89 @@ public static function decode(string $packet): self
     /**
      * 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
+     * When maxSize is specified, truncation follows RFC 1035 Section 6.2 and
+     * RFC 2181 Section 9:
+     * - Sections are dropped from the end first (additional → authority → answers)
+     * - Authority and additional are all-or-nothing; answers allow partial inclusion
+     * - TC flag is only set when answer records couldn't all fit
      * - 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();
         }
 
+        // Answers: include as many complete records as fit (partial allowed).
+        $answerCount = 0;
         foreach ($this->answers as $answer) {
-            $packet .= $answer->encode($packet);
+            $encoded = $answer->encode($packet);
+            if ($maxSize !== null && strlen($packet) + strlen($encoded) > $maxSize) {
+                break;
+            }
+            $packet .= $encoded;
+            $answerCount++;
         }
-
-        foreach ($this->authority as $authority) {
-            $packet .= $authority->encode($packet);
+        $answersTruncated = $answerCount < count($this->answers);
+
+        // Authority then additional: all-or-nothing, and only once answers all fit.
+        // Order matches RFC 1035 Section 6.2 (drop additional before authority).
+        $authorityCount = 0;
+        $additionalCount = 0;
+        if (!$answersTruncated) {
+            $withAuthority = $this->appendRecords($packet, $this->authority);
+            if ($maxSize === null || strlen($withAuthority) <= $maxSize) {
+                $packet = $withAuthority;
+                $authorityCount = count($this->authority);
+
+                $withAdditional = $this->appendRecords($packet, $this->additional);
+                if ($maxSize === null || strlen($withAdditional) <= $maxSize) {
+                    $packet = $withAdditional;
+                    $additionalCount = count($this->additional);
+                }
+            }
         }
 
-        foreach ($this->additional as $additional) {
-            $packet .= $additional->encode($packet);
-        }
+        $sectionsUnchanged =
+            $answerCount === count($this->answers)
+            && $authorityCount === count($this->authority)
+            && $additionalCount === count($this->additional);
 
-        // No truncation needed
-        if ($maxSize === null || strlen($packet) <= $maxSize) {
+        if ($sectionsUnchanged) {
             return $packet;
         }
 
-        // RFC-compliant truncation: work backward from end
-        // Per RFC 1035 Section 6.2 and RFC 2181 Section 9
-        return $this->encodeWithTruncation($maxSize);
+        // Per RFC 2181 Section 9, TC signals truncated required data (answers).
+        $header = new Header(
+            id: $this->header->id,
+            isResponse: $this->header->isResponse,
+            opcode: $this->header->opcode,
+            authoritative: $this->header->authoritative,
+            truncated: $answersTruncated,
+            recursionDesired: $this->header->recursionDesired,
+            recursionAvailable: $this->header->recursionAvailable,
+            responseCode: $this->header->responseCode,
+            questionCount: count($this->questions),
+            answerCount: $answerCount,
+            authorityCount: $authorityCount,
+            additionalCount: $additionalCount,
+        );
+
+        return $header->encode() . substr($packet, Header::LENGTH);
     }
 
     /**
-     * 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).
+     * @param list<Record> $records
      */
-    private function encodeWithTruncation(int $maxSize): string
+    private function appendRecords(string $packet, array $records): 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
-        $withoutAuthority = self::response(
-            $this->header,
-            $this->header->responseCode,
-            questions: $this->questions,
-            answers: $this->answers,
-            authority: [],
-            additional: [],
-            authoritative: $this->header->authoritative,
-            truncated: false,
-            recursionAvailable: $this->header->recursionAvailable
-        );
-
-        $packet = $withoutAuthority->encode();
-        if (strlen($packet) <= $maxSize) {
-            return $packet;
+        foreach ($records as $record) {
+            $packet .= $record->encode($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);
-
-        $truncatedResponse = self::response(
-            $this->header,
-            $this->header->responseCode,
-            questions: $this->questions,
-            answers: $fittingAnswers,
-            authority: [],
-            additional: [],
-            authoritative: $this->header->authoritative,
-            truncated: $needsTruncation,
-            recursionAvailable: $this->header->recursionAvailable
-        );
-
-        return $truncatedResponse->encode();
+        return $packet;
     }
 }

tests/unit/DNS/MessageTest.php

@@ -455,4 +455,190 @@ public function testEncodeWithoutMaxSizeDoesNotTruncate(): void
         // Verify all answers are preserved
         $this->assertCount(5, $decoded->answers);
     }
+
+    /**
+     * When authority doesn't fit, additional is dropped too — even if it
+     * would have fit alongside answers on its own. Locks in the section
+     * priority (additional depends on authority being included first).
+     */
+    public function testTruncationDropsAdditionalWhenAuthorityOverflows(): void
+    {
+        $question = new Question('example.com', Record::TYPE_A);
+        $query = Message::query($question, id: 0xAB01);
+
+        $answers = [
+            new Record('example.com', Record::TYPE_A, Record::CLASS_IN, 60, '192.168.1.1'),
+        ];
+
+        // Oversized authority — will not fit
+        $authority = [];
+        for ($i = 0; $i < 30; $i++) {
+            $authority[] = new Record('example.com', Record::TYPE_NS, Record::CLASS_IN, 3600, 'ns' . $i . '.example.com');
+        }
+
+        // Tiny additional — would fit on its own with just the answers
+        $additional = [
+            new Record('glue.example.com', Record::TYPE_A, Record::CLASS_IN, 60, '192.168.1.2'),
+        ];
+
+        $response = Message::response(
+            $query->header,
+            Message::RCODE_NOERROR,
+            questions: $query->questions,
+            answers: $answers,
+            authority: $authority,
+            additional: $additional
+        );
+
+        $truncated = $response->encode(512);
+        $decoded = Message::decode($truncated);
+
+        $this->assertFalse($decoded->header->truncated, 'TC should not be set when only authority/additional dropped');
+        $this->assertCount(1, $decoded->answers);
+        $this->assertCount(0, $decoded->authority);
+        $this->assertCount(0, $decoded->additional, 'Additional must be dropped whenever authority is dropped');
+        $this->assertLessThanOrEqual(512, strlen($truncated));
+    }
+
+    /**
+     * When answers are partially truncated, authority and additional are
+     * always cleared regardless of how much room remains. Prior tests used
+     * empty authority/additional for this path — this one populates both.
+     */
+    public function testAnswerTruncationDropsPopulatedAuthorityAndAdditional(): void
+    {
+        $question = new Question('example.com', Record::TYPE_A);
+        $query = Message::query($question, id: 0xCD02);
+
+        $answers = [];
+        for ($i = 0; $i < 100; $i++) {
+            $answers[] = new Record('example.com', Record::TYPE_A, Record::CLASS_IN, 60, '10.0.' . ($i % 256) . '.' . ($i % 256));
+        }
+
+        $authority = [];
+        for ($i = 0; $i < 5; $i++) {
+            $authority[] = new Record('example.com', Record::TYPE_NS, Record::CLASS_IN, 3600, 'ns' . $i . '.example.com');
+        }
+
+        $additional = [];
+        for ($i = 0; $i < 5; $i++) {
+            $additional[] = new Record('ns' . $i . '.example.com', Record::TYPE_A, Record::CLASS_IN, 60, '192.168.2.' . $i);
+        }
+
+        $response = Message::response(
+            $query->header,
+            Message::RCODE_NOERROR,
+            questions: $query->questions,
+            answers: $answers,
+            authority: $authority,
+            additional: $additional
+        );
+
+        $truncated = $response->encode(512);
+        $decoded = Message::decode($truncated);
+
+        $this->assertTrue($decoded->header->truncated, 'TC must be set when answers are partial');
+        $this->assertGreaterThan(0, count($decoded->answers));
+        $this->assertLessThan(100, count($decoded->answers));
+        $this->assertCount(0, $decoded->authority, 'Authority cleared under answer truncation');
+        $this->assertCount(0, $decoded->additional, 'Additional cleared under answer truncation');
+        $this->assertLessThanOrEqual(512, strlen($truncated));
+    }
+
+    /**
+     * Re-encoding a message whose header already has TC=1 must preserve
+     * the flag when nothing new is dropped. The encode() short-circuit
+     * relies on the original header bytes being returned verbatim.
+     */
+    public function testReEncodePreservesOriginalTruncatedFlag(): void
+    {
+        $question = new Question('example.com', Record::TYPE_A);
+        $query = Message::query($question, id: 0xEF03);
+
+        $answers = [
+            new Record('example.com', Record::TYPE_A, Record::CLASS_IN, 60, '192.168.1.1'),
+        ];
+
+        $response = Message::response(
+            $query->header,
+            Message::RCODE_NOERROR,
+            questions: $query->questions,
+            answers: $answers,
+            authority: [],
+            additional: [],
+            truncated: true
+        );
+
+        $encoded = $response->encode();
+        $decoded = Message::decode($encoded);
+
+        $this->assertTrue($decoded->header->truncated, 'TC flag must survive re-encoding when nothing is dropped');
+        $this->assertSame($encoded, $decoded->encode(), 'Second round-trip must be byte-identical');
+    }
+
+    /**
+     * maxSize equal to the natural encoded length must not trigger truncation.
+     * Guards against an off-by-one where `>` would incorrectly become `>=`.
+     */
+    public function testEncodeFitsExactlyAtMaxSizeBoundary(): void
+    {
+        $question = new Question('example.com', Record::TYPE_A);
+        $query = Message::query($question, id: 0x1104);
+
+        $answers = [
+            new Record('example.com', Record::TYPE_A, Record::CLASS_IN, 60, '192.168.1.1'),
+            new Record('example.com', Record::TYPE_A, Record::CLASS_IN, 60, '192.168.1.2'),
+        ];
+
+        $response = Message::response(
+            $query->header,
+            Message::RCODE_NOERROR,
+            questions: $query->questions,
+            answers: $answers,
+            authority: [],
+            additional: []
+        );
+
+        $natural = $response->encode();
+        $exactSize = strlen($natural);
+
+        $atBoundary = $response->encode($exactSize);
+        $this->assertSame($natural, $atBoundary, 'Encoding at exact size must match unconstrained output');
+
+        $belowBoundary = $response->encode($exactSize - 1);
+        $this->assertLessThan(count($answers), count(Message::decode($belowBoundary)->answers));
+    }
+
+    /**
+     * NODATA-style response: zero answers but populated authority. When the
+     * authority section can't fit, it's dropped without setting TC — there
+     * are no answer records that failed to transmit.
+     */
+    public function testNoAnswersWithOversizedAuthorityDropsWithoutTruncation(): void
+    {
+        $question = new Question('example.com', Record::TYPE_A);
+        $query = Message::query($question, id: 0x2205);
+
+        $authority = [];
+        for ($i = 0; $i < 30; $i++) {
+            $authority[] = new Record('example.com', Record::TYPE_NS, Record::CLASS_IN, 3600, 'ns' . $i . '.example.com');
+        }
+
+        $response = Message::response(
+            $query->header,
+            Message::RCODE_NOERROR,
+            questions: $query->questions,
+            answers: [],
+            authority: $authority,
+            additional: []
+        );
+
+        $truncated = $response->encode(512);
+        $decoded = Message::decode($truncated);
+
+        $this->assertFalse($decoded->header->truncated, 'TC must remain unset when there were no answers to truncate');
+        $this->assertCount(0, $decoded->answers);
+        $this->assertCount(0, $decoded->authority, 'Oversized authority is dropped');
+        $this->assertLessThanOrEqual(512, strlen($truncated));
+    }
 }