Also use other compressed names in DNS compression

The previous implementation of writeLabels would not include a
compressed name in the label dictionary, so if a packet had
"something.local", "a.service.local" and "b.service.local",
"service.local" would not be compressed because "a.service.local"
already used compression (for .local).

Fix this and add a test.

Bug: 254166302
Test: atest
Change-Id: I41c557d6debd11acb4c0813735ef7af7323f45d7

service/mdns/com/android/server/connectivity/mdns/MdnsPacketWriter.java

@@ -192,22 +192,31 @@ public class MdnsPacketWriter {
             }
         }
 
+        final int[] offsets;
         if (suffixLength > 0) {
-            for (int i = 0; i < (labels.length - suffixLength); ++i) {
+            offsets = writePartialLabelsNoCompression(labels, labels.length - suffixLength);
-                writeString(labels[i]);
-            }
             writePointer(suffixPointer);
         } else {
-            int[] offsets = writeLabelsNoCompression(labels);
+            offsets = writeLabelsNoCompression(labels);
-
-            // Add entries to the label dictionary for each suffix of the label list, including
-            // the whole list itself.
-            for (int i = 0, len = labels.length; i < labels.length; ++i, --len) {
-                String[] value = new String[len];
-                System.arraycopy(labels, i, value, 0, len);
-                labelDictionary.put(offsets[i], value);
-            }
         }
+
+        // Add entries to the label dictionary for each suffix of the label list, including
+        // the whole list itself.
+        // Do not replace the last suffixLength suffixes that already have dictionary entries.
+        for (int i = 0, len = labels.length; i < labels.length - suffixLength; ++i, --len) {
+            String[] value = new String[len];
+            System.arraycopy(labels, i, value, 0, len);
+            labelDictionary.put(offsets[i], value);
+        }
+    }
+
+    private int[] writePartialLabelsNoCompression(String[] labels, int count) throws IOException {
+        int[] offsets = new int[count];
+        for (int i = 0; i < count; ++i) {
+            offsets[i] = getWritePosition();
+            writeString(labels[i]);
+        }
+        return offsets;
     }
 
     /**
@@ -216,11 +225,7 @@ public class MdnsPacketWriter {
      * @return The offsets where each label was written to.
      */
     public int[] writeLabelsNoCompression(String[] labels) throws IOException {
-        int[] offsets = new int[labels.length];
+        final int[] offsets = writePartialLabelsNoCompression(labels, labels.length);
-        for (int i = 0; i < labels.length; ++i) {
-            offsets[i] = getWritePosition();
-            writeString(labels[i]);
-        }
         writeUInt8(0); // NUL terminator
         return offsets;
     }
@@ -246,4 +251,4 @@ public class MdnsPacketWriter {
     public DatagramPacket getPacket(SocketAddress destAddress) throws IOException {
         return new DatagramPacket(data, pos, destAddress);
     }
-}
+}

tests/unit/java/com/android/server/connectivity/mdns/MdnsPacketWriterTest.kt

@@ -0,0 +1,55 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.connectivity.mdns
+
+import android.net.InetAddresses
+import android.os.Build
+import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
+import com.android.testutils.DevSdkIgnoreRunner
+import java.net.InetSocketAddress
+import kotlin.test.assertContentEquals
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(DevSdkIgnoreRunner::class)
+@IgnoreUpTo(Build.VERSION_CODES.S_V2)
+class MdnsPacketWriterTest {
+    @Test
+    fun testNameCompression() {
+        val writer = MdnsPacketWriter(ByteArray(1000))
+        writer.writeLabels(arrayOf("my", "first", "name"))
+        writer.writeLabels(arrayOf("my", "second", "name"))
+        writer.writeLabels(arrayOf("other", "first", "name"))
+        writer.writeLabels(arrayOf("my", "second", "name"))
+        writer.writeLabels(arrayOf("unrelated"))
+
+        val packet = writer.getPacket(
+                InetSocketAddress(InetAddresses.parseNumericAddress("2001:db8::123"), 123))
+
+        // Each label takes length + 1. So "first.name" offset = 3, "name" offset = 9
+        val expected = "my".label() + "first".label() + "name".label() + 0x00.toByte() +
+                // "my.second.name" offset = 15
+                "my".label() + "second".label() + byteArrayOf(0xC0.toByte(), 9) +
+                "other".label() + byteArrayOf(0xC0.toByte(), 3) +
+                byteArrayOf(0xC0.toByte(), 15) +
+                "unrelated".label() + 0x00.toByte()
+
+        assertContentEquals(expected, packet.data.copyOfRange(0, packet.length))
+    }
+}
+
+private fun String.label() = byteArrayOf(length.toByte()) + encodeToByteArray()

tests/unit/java/com/android/server/connectivity/mdns/MdnsProberTest.kt

@@ -160,14 +160,11 @@ class MdnsProberTest {
                 scapy.DNSRR(type='TXT', ttl=120, rrname='testservice._nmt._tcp.local.',
                     rdata='testKey=testValue'))
         )).hex().upper()
-        // NOTE: due to a bug the second "myhostname" is not getting DNS compressed in the current
-        // actual probe, so data below is slightly different. Fix compression so it gets compressed.
          */
         val expected = "0000000000020000000300000B7465737473657276696365045F6E6D74045F746370056C6" +
                 "F63616C0000FF00010C746573747365727669636532C01800FF0001C00C002100010000007800130" +
-                "000000094020A6D79686F73746E616D65C0220C746573747365727669636532C0180021000100000" +
+                "000000094020A6D79686F73746E616D65C022C02D00210001000000780008000000009403C052C00" +
-                "07800130000000094030A6D79686F73746E616D65C022C00C0010000100000078001211746573744" +
+                "C0010000100000078001211746573744B65793D7465737456616C7565"
-                "B65793D7465737456616C7565"
         assertProbesSent(probeInfo, expected)
     }