Build MdnsResponse in ServiceTypeClient

Instead of building MdnsResponse lists in MdnsSocketClient and sending
each one to ServiceTypeClient, only parse packets in MdnsSocketClient,
and let each ServiceTypeClient build the MdnsResponse from them.

MdnsResponseDecoder already implements the same filtering that
MdnsDiscoveryManager#onResponseReceived did on the service types so this
is a no-op.

This will be useful to allow MdnsServiceTypeClient to update its
responses using records that are not necessarily part of a full
MdnsResponse; for example if a reply packet only contains address
records, or only SRV/TXT records without a PTR record.

Bug: 267570781
Test: atest
Change-Id: I8157c3732c425aae497d44c44c64b640859c13dc
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
index cc6b98b..67059e7 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
@@ -19,6 +19,7 @@
 import android.Manifest.permission;
 import android.annotation.NonNull;
 import android.annotation.RequiresPermission;
+import android.net.Network;
 import android.text.TextUtils;
 import android.util.ArrayMap;
 import android.util.Log;
@@ -27,7 +28,6 @@
 import com.android.server.connectivity.mdns.util.MdnsLogger;
 
 import java.io.IOException;
-import java.util.Arrays;
 import java.util.Map;
 
 /**
@@ -119,22 +119,10 @@
     }
 
     @Override
-    public synchronized void onResponseReceived(@NonNull MdnsResponse response) {
-        String[] name =
-                response.getPointerRecords().isEmpty()
-                        ? null
-                        : response.getPointerRecords().get(0).getName();
-        if (name != null) {
-            for (MdnsServiceTypeClient serviceTypeClient : serviceTypeClients.values()) {
-                String[] serviceType = serviceTypeClient.getServiceTypeLabels();
-                if ((Arrays.equals(name, serviceType)
-                        || ((name.length == (serviceType.length + 2))
-                        && name[1].equals(MdnsConstants.SUBTYPE_LABEL)
-                        && MdnsRecord.labelsAreSuffix(serviceType, name)))) {
-                    serviceTypeClient.processResponse(response);
-                    return;
-                }
-            }
+    public synchronized void onResponseReceived(@NonNull MdnsPacket packet,
+            int interfaceIndex, Network network) {
+        for (MdnsServiceTypeClient serviceTypeClient : serviceTypeClients.values()) {
+            serviceTypeClient.processResponse(packet, interfaceIndex, network);
         }
     }
 
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
index d959065..93972d9 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
@@ -33,7 +33,6 @@
 import java.net.Inet4Address;
 import java.net.Inet6Address;
 import java.net.InetSocketAddress;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
@@ -48,7 +47,6 @@
 
     @NonNull private final Handler mHandler;
     @NonNull private final MdnsSocketProvider mSocketProvider;
-    @NonNull private final MdnsResponseDecoder mResponseDecoder;
 
     private final Map<MdnsServiceBrowserListener, InterfaceSocketCallback> mRequestedNetworks =
             new ArrayMap<>();
@@ -62,8 +60,6 @@
             @NonNull MdnsSocketProvider provider) {
         mHandler = new Handler(looper);
         mSocketProvider = provider;
-        mResponseDecoder = new MdnsResponseDecoder(
-                new MdnsResponseDecoder.Clock(), null /* serviceType */);
     }
 
     private class InterfaceSocketCallback implements MdnsSocketProvider.SocketCallback {
@@ -170,19 +166,21 @@
             @NonNull Network network) {
         int packetNumber = ++mReceivedPacketNumber;
 
-        final List<MdnsResponse> responses = new ArrayList<>();
-        final int errorCode = mResponseDecoder.decode(
-                recvbuf, length, responses, interfaceIndex, network);
-        if (errorCode == MdnsResponseDecoder.SUCCESS) {
-            for (MdnsResponse response : responses) {
+        final MdnsPacket response;
+        try {
+            response = MdnsResponseDecoder.parseResponse(recvbuf, length);
+        } catch (MdnsPacket.ParseException e) {
+            if (e.code != MdnsResponseErrorCode.ERROR_NOT_RESPONSE_MESSAGE) {
+                Log.e(TAG, e.getMessage(), e);
                 if (mCallback != null) {
-                    mCallback.onResponseReceived(response);
+                    mCallback.onFailedToParseMdnsResponse(packetNumber, e.code);
                 }
             }
-        } else if (errorCode != MdnsResponseErrorCode.ERROR_NOT_RESPONSE_MESSAGE) {
-            if (mCallback != null) {
-                mCallback.onFailedToParseMdnsResponse(packetNumber, errorCode);
-            }
+            return;
+        }
+
+        if (mCallback != null) {
+            mCallback.onResponseReceived(response, interfaceIndex, network);
         }
     }
 
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsResponseDecoder.java b/service-t/src/com/android/server/connectivity/mdns/MdnsResponseDecoder.java
index 82da2e4..31527c0 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsResponseDecoder.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsResponseDecoder.java
@@ -24,7 +24,6 @@
 import com.android.server.connectivity.mdns.util.MdnsLogger;
 
 import java.io.EOFException;
-import java.net.DatagramPacket;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -82,34 +81,16 @@
 
     /**
      * Decodes all mDNS responses for the desired service type from a packet. The class does not
-     * check
-     * the responses for completeness; the caller should do that.
-     *
-     * @param packet The packet to read from.
-     * @param interfaceIndex the network interface index (or {@link
-     *     MdnsSocket#INTERFACE_INDEX_UNSPECIFIED} if not known) at which the packet was received
-     * @param network the network at which the packet was received, or null if it is unknown.
-     * @return A list of mDNS responses, or null if the packet contained no appropriate responses.
-     */
-    public int decode(@NonNull DatagramPacket packet, @NonNull List<MdnsResponse> responses,
-            int interfaceIndex, @Nullable Network network) {
-        return decode(packet.getData(), packet.getLength(), responses, interfaceIndex, network);
-    }
-
-    /**
-     * Decodes all mDNS responses for the desired service type from a packet. The class does not
-     * check
-     * the responses for completeness; the caller should do that.
+     * check the responses for completeness; the caller should do that.
      *
      * @param recvbuf The received data buffer to read from.
      * @param length The length of received data buffer.
-     * @param interfaceIndex the network interface index (or {@link
-     *     MdnsSocket#INTERFACE_INDEX_UNSPECIFIED} if not known) at which the packet was received
-     * @param network the network at which the packet was received, or null if it is unknown.
-     * @return A list of mDNS responses, or null if the packet contained no appropriate responses.
+     * @return A decoded {@link MdnsPacket}.
+     * @throws MdnsPacket.ParseException if a response packet could not be parsed.
      */
-    public int decode(@NonNull byte[] recvbuf, int length, @NonNull List<MdnsResponse> responses,
-            int interfaceIndex, @Nullable Network network) {
+    @NonNull
+    public static MdnsPacket parseResponse(@NonNull byte[] recvbuf, int length)
+            throws MdnsPacket.ParseException {
         MdnsPacketReader reader = new MdnsPacketReader(recvbuf, length);
 
         final MdnsPacket mdnsPacket;
@@ -117,21 +98,35 @@
             reader.readUInt16(); // transaction ID (not used)
             int flags = reader.readUInt16();
             if ((flags & MdnsConstants.FLAGS_RESPONSE_MASK) != MdnsConstants.FLAGS_RESPONSE) {
-                return MdnsResponseErrorCode.ERROR_NOT_RESPONSE_MESSAGE;
+                throw new MdnsPacket.ParseException(
+                        MdnsResponseErrorCode.ERROR_NOT_RESPONSE_MESSAGE, "Not a response", null);
             }
 
             mdnsPacket = MdnsPacket.parseRecordsSection(reader, flags);
             if (mdnsPacket.answers.size() < 1) {
-                return MdnsResponseErrorCode.ERROR_NO_ANSWERS;
+                throw new MdnsPacket.ParseException(
+                        MdnsResponseErrorCode.ERROR_NOT_RESPONSE_MESSAGE, "Response has no answers",
+                        null);
             }
+            return mdnsPacket;
         } catch (EOFException e) {
-            LOGGER.e("Reached the end of the mDNS response unexpectedly.", e);
-            return MdnsResponseErrorCode.ERROR_END_OF_FILE;
-        } catch (MdnsPacket.ParseException e) {
-            LOGGER.e(e.getMessage(), e);
-            return e.code;
+            throw new MdnsPacket.ParseException(MdnsResponseErrorCode.ERROR_END_OF_FILE,
+                    "Reached the end of the mDNS response unexpectedly.", e);
         }
+    }
 
+    /**
+     * Builds mDNS responses for the desired service type from a packet.
+     *
+     * The class does not check the responses for completeness; the caller should do that.
+     *
+     * @param mdnsPacket The packet to read from.
+     * @param interfaceIndex the network interface index (or {@link
+     *     MdnsSocket#INTERFACE_INDEX_UNSPECIFIED} if not known) at which the packet was received.
+     * @param network the network at which the packet was received, or null if it is unknown.
+     */
+    public List<MdnsResponse> buildResponses(@NonNull MdnsPacket mdnsPacket, int interfaceIndex,
+            @Nullable Network network) {
         final ArrayList<MdnsRecord> records = new ArrayList<>(
                 mdnsPacket.questions.size() + mdnsPacket.answers.size()
                         + mdnsPacket.authorityRecords.size() + mdnsPacket.additionalRecords.size());
@@ -139,6 +134,8 @@
         records.addAll(mdnsPacket.authorityRecords);
         records.addAll(mdnsPacket.additionalRecords);
 
+        final ArrayList<MdnsResponse> responses = new ArrayList<>();
+
         // The response records are structured in a hierarchy, where some records reference
         // others, as follows:
         //
@@ -224,8 +221,7 @@
                 }
             }
         }
-
-        return SUCCESS;
+        return responses;
     }
 
     private static void assignInetRecord(MdnsResponse response, MdnsInetAddressRecord inetRecord) {
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
index e33686f..707a57e 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -21,7 +21,6 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.net.Network;
-import android.os.SystemClock;
 import android.text.TextUtils;
 import android.util.ArrayMap;
 import android.util.Pair;
@@ -54,6 +53,7 @@
     private final String serviceType;
     private final String[] serviceTypeLabels;
     private final MdnsSocketClientBase socketClient;
+    private final MdnsResponseDecoder responseDecoder;
     private final ScheduledExecutorService executor;
     private final Object lock = new Object();
     private final ArrayMap<MdnsServiceBrowserListener, MdnsSearchOptions> listeners =
@@ -64,6 +64,8 @@
     private final boolean allowSearchOptionsToRemoveExpiredService =
             MdnsConfigs.allowSearchOptionsToRemoveExpiredService();
 
+    private final MdnsResponseDecoder.Clock clock;
+
     @Nullable private MdnsSearchOptions searchOptions;
 
     // The session ID increases when startSendAndReceive() is called where we schedule a
@@ -85,10 +87,21 @@
             @NonNull String serviceType,
             @NonNull MdnsSocketClientBase socketClient,
             @NonNull ScheduledExecutorService executor) {
+        this(serviceType, socketClient, executor, new MdnsResponseDecoder.Clock());
+    }
+
+    @VisibleForTesting
+    public MdnsServiceTypeClient(
+            @NonNull String serviceType,
+            @NonNull MdnsSocketClientBase socketClient,
+            @NonNull ScheduledExecutorService executor,
+            @NonNull MdnsResponseDecoder.Clock clock) {
         this.serviceType = serviceType;
         this.socketClient = socketClient;
         this.executor = executor;
-        serviceTypeLabels = TextUtils.split(serviceType, "\\.");
+        this.serviceTypeLabels = TextUtils.split(serviceType, "\\.");
+        this.responseDecoder = new MdnsResponseDecoder(clock, serviceTypeLabels);
+        this.clock = clock;
     }
 
     private static MdnsServiceInfo buildMdnsServiceInfoFromResponse(
@@ -198,24 +211,32 @@
         return serviceTypeLabels;
     }
 
-    public synchronized void processResponse(@NonNull MdnsResponse response) {
-        if (shouldRemoveServiceAfterTtlExpires()) {
-            // Because {@link QueryTask} and {@link processResponse} are running in different
-            // threads. We need to synchronize {@link lock} to protect
-            // {@link instanceNameToResponse} won’t be modified at the same time.
-            synchronized (lock) {
+    /**
+     * Process an incoming response packet.
+     */
+    public synchronized void processResponse(@NonNull MdnsPacket packet, int interfaceIndex,
+            Network network) {
+        final List<MdnsResponse> responses = responseDecoder.buildResponses(packet, interfaceIndex,
+                network);
+        for (MdnsResponse response : responses) {
+            if (shouldRemoveServiceAfterTtlExpires()) {
+                // Because {@link QueryTask} and {@link processResponse} are running in different
+                // threads. We need to synchronize {@link lock} to protect
+                // {@link instanceNameToResponse} won’t be modified at the same time.
+                synchronized (lock) {
+                    if (response.isGoodbye()) {
+                        onGoodbyeReceived(response.getServiceInstanceName());
+                    } else {
+                        onResponseReceived(response);
+                    }
+                }
+            } else {
                 if (response.isGoodbye()) {
                     onGoodbyeReceived(response.getServiceInstanceName());
                 } else {
                     onResponseReceived(response);
                 }
             }
-        } else {
-            if (response.isGoodbye()) {
-                onGoodbyeReceived(response.getServiceInstanceName());
-            } else {
-                onResponseReceived(response);
-            }
         }
     }
 
@@ -481,7 +502,7 @@
                         if (existingResponse.hasServiceRecord()
                                 && existingResponse
                                 .getServiceRecord()
-                                .getRemainingTTL(SystemClock.elapsedRealtime())
+                                .getRemainingTTL(clock.elapsedRealtime())
                                 == 0) {
                             iter.remove();
                             for (int i = 0; i < listeners.size(); i++) {
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClient.java
index 907687e..c03e6aa 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClient.java
@@ -16,8 +16,6 @@
 
 package com.android.server.connectivity.mdns;
 
-import static com.android.server.connectivity.mdns.MdnsSocketClientBase.Callback;
-
 import android.Manifest.permission;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -35,7 +33,6 @@
 import java.net.DatagramPacket;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
-import java.util.LinkedList;
 import java.util.List;
 import java.util.Queue;
 import java.util.Timer;
@@ -73,7 +70,6 @@
     private final Context context;
     private final byte[] multicastReceiverBuffer = new byte[RECEIVER_BUFFER_SIZE];
     @Nullable private final byte[] unicastReceiverBuffer;
-    private final MdnsResponseDecoder responseDecoder;
     private final MulticastLock multicastLock;
     private final boolean useSeparateSocketForUnicast =
             MdnsConfigs.useSeparateSocketToSendUnicastQuery();
@@ -110,7 +106,6 @@
     public MdnsSocketClient(@NonNull Context context, @NonNull MulticastLock multicastLock) {
         this.context = context;
         this.multicastLock = multicastLock;
-        responseDecoder = new MdnsResponseDecoder(new MdnsResponseDecoder.Clock(), null);
         if (useSeparateSocketForUnicast) {
             unicastReceiverBuffer = new byte[RECEIVER_BUFFER_SIZE];
         } else {
@@ -421,37 +416,27 @@
             int interfaceIndex, @Nullable Network network) {
         int packetNumber = ++receivedPacketNumber;
 
-        List<MdnsResponse> responses = new LinkedList<>();
-        int errorCode = responseDecoder.decode(packet, responses, interfaceIndex, network);
-        if (errorCode == MdnsResponseDecoder.SUCCESS) {
-            if (responseType.equals(MULTICAST_TYPE)) {
-                receivedMulticastResponse = true;
-                if (cannotReceiveMulticastResponse.getAndSet(false)) {
-                    // If we are already in the bad state, receiving a multicast response means
-                    // we are recovered.
-                    LOGGER.e(
-                            "Recovered from the state where the phone can't receive any multicast"
-                                    + " response");
-                }
-            } else {
-                receivedUnicastResponse = true;
-            }
-            for (MdnsResponse response : responses) {
-                String serviceInstanceName = response.getServiceInstanceName();
-                LOGGER.log("mDNS %s response received: %s at ifIndex %d", responseType,
-                        serviceInstanceName, interfaceIndex);
-                if (callback != null) {
-                    callback.onResponseReceived(response);
-                }
-            }
-        } else if (errorCode != MdnsResponseErrorCode.ERROR_NOT_RESPONSE_MESSAGE) {
+        final MdnsPacket response;
+        try {
+            response = MdnsResponseDecoder.parseResponse(packet.getData(), packet.getLength());
+        } catch (MdnsPacket.ParseException e) {
             LOGGER.w(String.format("Error while decoding %s packet (%d): %d",
-                    responseType, packetNumber, errorCode));
+                    responseType, packetNumber, e.code));
             if (callback != null) {
-                callback.onFailedToParseMdnsResponse(packetNumber, errorCode);
+                callback.onFailedToParseMdnsResponse(packetNumber, e.code);
             }
+            return e.code;
         }
-        return errorCode;
+
+        if (response == null) {
+            return MdnsResponseErrorCode.ERROR_NOT_RESPONSE_MESSAGE;
+        }
+
+        if (callback != null) {
+            callback.onResponseReceived(response, interfaceIndex, network);
+        }
+
+        return MdnsResponseErrorCode.SUCCESS;
     }
 
     @VisibleForTesting
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java
index 23504a0..796dc83 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java
@@ -72,7 +72,8 @@
     /*** Callback for mdns response  */
     interface Callback {
         /*** Receive a mdns response */
-        void onResponseReceived(@NonNull MdnsResponse response);
+        void onResponseReceived(@NonNull MdnsPacket packet, int interfaceIndex,
+                @Nullable Network network);
 
         /*** Parse a mdns response failed */
         void onFailedToParseMdnsResponse(int receivedPacketNumber, int errorCode);
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
index 83e7696..e6b8326 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
@@ -23,6 +23,7 @@
 import static org.mockito.Mockito.when;
 
 import android.annotation.NonNull;
+import android.net.Network;
 import android.text.TextUtils;
 
 import com.android.testutils.DevSdkIgnoreRule;
@@ -35,7 +36,10 @@
 import org.mockito.MockitoAnnotations;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.List;
 
 /** Tests for {@link MdnsDiscoveryManager}. */
 @RunWith(DevSdkIgnoreRunner.class)
@@ -111,25 +115,38 @@
         discoveryManager.registerListener(
                 SERVICE_TYPE_2, mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
 
-        MdnsResponse responseForServiceTypeOne = createMockResponse(SERVICE_TYPE_1);
-        discoveryManager.onResponseReceived(responseForServiceTypeOne);
-        verify(mockServiceTypeClientOne).processResponse(responseForServiceTypeOne);
+        MdnsPacket responseForServiceTypeOne = createMdnsPacket(SERVICE_TYPE_1);
+        final int ifIndex = 1;
+        final Network network = mock(Network.class);
+        discoveryManager.onResponseReceived(responseForServiceTypeOne, ifIndex, network);
+        verify(mockServiceTypeClientOne).processResponse(responseForServiceTypeOne, ifIndex,
+                network);
 
-        MdnsResponse responseForServiceTypeTwo = createMockResponse(SERVICE_TYPE_2);
-        discoveryManager.onResponseReceived(responseForServiceTypeTwo);
-        verify(mockServiceTypeClientTwo).processResponse(responseForServiceTypeTwo);
+        MdnsPacket responseForServiceTypeTwo = createMdnsPacket(SERVICE_TYPE_2);
+        discoveryManager.onResponseReceived(responseForServiceTypeTwo, ifIndex, network);
+        verify(mockServiceTypeClientTwo).processResponse(responseForServiceTypeTwo, ifIndex,
+                network);
 
-        MdnsResponse responseForSubtype = createMockResponse("subtype._sub._googlecast._tcp.local");
-        discoveryManager.onResponseReceived(responseForSubtype);
-        verify(mockServiceTypeClientOne).processResponse(responseForSubtype);
+        MdnsPacket responseForSubtype = createMdnsPacket("subtype._sub._googlecast._tcp.local");
+        discoveryManager.onResponseReceived(responseForSubtype, ifIndex, network);
+        verify(mockServiceTypeClientOne).processResponse(responseForSubtype, ifIndex, network);
     }
 
-    private MdnsResponse createMockResponse(String serviceType) {
-        MdnsPointerRecord mockPointerRecord = mock(MdnsPointerRecord.class);
-        MdnsResponse mockResponse = mock(MdnsResponse.class);
-        when(mockResponse.getPointerRecords())
-                .thenReturn(Collections.singletonList(mockPointerRecord));
-        when(mockPointerRecord.getName()).thenReturn(TextUtils.split(serviceType, "\\."));
-        return mockResponse;
+    private MdnsPacket createMdnsPacket(String serviceType) {
+        final String[] type = TextUtils.split(serviceType, "\\.");
+        final ArrayList<String> name = new ArrayList<>(type.length + 1);
+        name.add("TestName");
+        name.addAll(Arrays.asList(type));
+        return new MdnsPacket(0 /* flags */,
+                Collections.emptyList() /* questions */,
+                List.of(new MdnsPointerRecord(
+                        type,
+                        0L /* receiptTimeMillis */,
+                        false /* cacheFlush */,
+                        120000 /* ttlMillis */,
+                        name.toArray(new String[0])
+                        )) /* answers */,
+                Collections.emptyList() /* authorityRecords */,
+                Collections.emptyList() /* additionalRecords */);
     }
 }
\ No newline at end of file
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
index 9d42a65..1e322e4 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
@@ -19,9 +19,9 @@
 import static com.android.server.connectivity.mdns.MdnsSocketProvider.SocketCallback;
 import static com.android.server.connectivity.mdns.MulticastPacketReader.PacketHandler;
 
-import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.timeout;
@@ -146,16 +146,23 @@
         // Send the data and verify the received records.
         final PacketHandler handler = handlerCaptor.getValue();
         handler.handlePacket(data, data.length, null /* src */);
-        final ArgumentCaptor<MdnsResponse> responseCaptor =
-                ArgumentCaptor.forClass(MdnsResponse.class);
-        verify(mCallback).onResponseReceived(responseCaptor.capture());
-        final MdnsResponse response = responseCaptor.getValue();
-        assertTrue(response.hasPointerRecords());
-        assertArrayEquals("_testtype._tcp.local".split("\\."),
-                response.getPointerRecords().get(0).getName());
-        assertTrue(response.hasServiceRecord());
-        assertEquals("testservice", response.getServiceRecord().getServiceInstanceName());
-        assertEquals("Android.local".split("\\."),
-                response.getServiceRecord().getServiceHost());
+        final ArgumentCaptor<MdnsPacket> responseCaptor =
+                ArgumentCaptor.forClass(MdnsPacket.class);
+        verify(mCallback).onResponseReceived(responseCaptor.capture(), anyInt(), any());
+        final MdnsPacket response = responseCaptor.getValue();
+        assertEquals(0, response.questions.size());
+        assertEquals(0, response.additionalRecords.size());
+        assertEquals(0, response.authorityRecords.size());
+
+        final String[] serviceName = "testservice._testtype._tcp.local".split("\\.");
+        assertEquals(List.of(
+                new MdnsPointerRecord("_testtype._tcp.local".split("\\."),
+                        0L /* receiptTimeMillis */, false /* cacheFlush */, 4500000 /* ttlMillis */,
+                        serviceName),
+                new MdnsServiceRecord(serviceName, 0L /* receiptTimeMillis */,
+                        false /* cacheFlush */, 4500000 /* ttlMillis */, 0 /* servicePriority */,
+                        0 /* serviceWeight */, 31234 /* servicePort */,
+                        new String[] { "Android", "local" } /* serviceHost */)
+        ), response.answers);
     }
 }
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsResponseDecoderTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsResponseDecoderTests.java
index 4cae447..aaef048 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsResponseDecoderTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsResponseDecoderTests.java
@@ -43,7 +43,6 @@
 import java.net.Inet4Address;
 import java.net.Inet6Address;
 import java.net.InetSocketAddress;
-import java.util.LinkedList;
 import java.util.List;
 
 @RunWith(DevSdkIgnoreRunner.class)
@@ -154,35 +153,22 @@
     private static final String[] MATTER_SERVICE_TYPE =
             new String[] {MATTER_SERVICE_NAME, "_tcp", "local"};
 
-    private final List<MdnsResponse> responses = new LinkedList<>();
+    private List<MdnsResponse> responses;
 
     private final Clock mClock = mock(Clock.class);
 
     @Before
-    public void setUp() {
+    public void setUp() throws Exception {
         MdnsResponseDecoder decoder = new MdnsResponseDecoder(mClock, CAST_SERVICE_TYPE);
         assertNotNull(data);
-        DatagramPacket packet = new DatagramPacket(data, data.length);
-        packet.setSocketAddress(
-                new InetSocketAddress(MdnsConstants.getMdnsIPv4Address(), MdnsConstants.MDNS_PORT));
-        responses.clear();
-        int errorCode = decoder.decode(
-                packet, responses, MdnsSocket.INTERFACE_INDEX_UNSPECIFIED, mock(Network.class));
-        assertEquals(MdnsResponseDecoder.SUCCESS, errorCode);
+        responses = decode(decoder, data);
         assertEquals(1, responses.size());
     }
 
     @Test
-    public void testDecodeWithNullServiceType() {
+    public void testDecodeWithNullServiceType() throws Exception {
         MdnsResponseDecoder decoder = new MdnsResponseDecoder(mClock, null);
-        assertNotNull(data);
-        DatagramPacket packet = new DatagramPacket(data, data.length);
-        packet.setSocketAddress(
-                new InetSocketAddress(MdnsConstants.getMdnsIPv4Address(), MdnsConstants.MDNS_PORT));
-        responses.clear();
-        int errorCode = decoder.decode(
-                packet, responses, MdnsSocket.INTERFACE_INDEX_UNSPECIFIED, mock(Network.class));
-        assertEquals(MdnsResponseDecoder.SUCCESS, errorCode);
+        responses = decode(decoder, data);
         assertEquals(2, responses.size());
     }
 
@@ -235,15 +221,9 @@
     public void testDecodeIPv6AnswerPacket() throws IOException {
         MdnsResponseDecoder decoder = new MdnsResponseDecoder(mClock, CAST_SERVICE_TYPE);
         assertNotNull(data6);
-        DatagramPacket packet = new DatagramPacket(data6, data6.length);
-        packet.setSocketAddress(
-                new InetSocketAddress(MdnsConstants.getMdnsIPv6Address(), MdnsConstants.MDNS_PORT));
 
-        responses.clear();
-        int errorCode = decoder.decode(
-                packet, responses, MdnsSocket.INTERFACE_INDEX_UNSPECIFIED, mock(Network.class));
-        assertEquals(MdnsResponseDecoder.SUCCESS, errorCode);
-
+        responses = decode(decoder, data6);
+        assertEquals(1, responses.size());
         MdnsResponse response = responses.get(0);
         assertTrue(response.isComplete());
 
@@ -283,39 +263,33 @@
     }
 
     @Test
-    public void decode_withInterfaceIndex_populatesInterfaceIndex() {
+    public void decode_withInterfaceIndex_populatesInterfaceIndex() throws Exception {
         MdnsResponseDecoder decoder = new MdnsResponseDecoder(mClock, CAST_SERVICE_TYPE);
         assertNotNull(data6);
         DatagramPacket packet = new DatagramPacket(data6, data6.length);
         packet.setSocketAddress(
                 new InetSocketAddress(MdnsConstants.getMdnsIPv6Address(), MdnsConstants.MDNS_PORT));
 
-        responses.clear();
+        final MdnsPacket parsedPacket = MdnsResponseDecoder.parseResponse(data6, data6.length);
+        assertNotNull(parsedPacket);
+
         final Network network = mock(Network.class);
-        int errorCode = decoder.decode(
-                packet, responses, /* interfaceIndex= */ 10, network);
-        assertEquals(errorCode, MdnsResponseDecoder.SUCCESS);
+        responses = decoder.buildResponses(parsedPacket,
+                /* interfaceIndex= */ 10, network /* expireOnExit= */);
+
         assertEquals(responses.size(), 1);
         assertEquals(responses.get(0).getInterfaceIndex(), 10);
         assertEquals(network, responses.get(0).getNetwork());
     }
 
     @Test
-    public void decode_singleHostname_multipleSrvRecords_flagEnabled_multipleCompleteResponses() {
+    public void decode_singleHostname_multipleSrvRecords_flagEnabled_multipleCompleteResponses()
+            throws Exception {
         //MdnsScannerConfigsFlagsImpl.allowMultipleSrvRecordsPerHost.override(true);
         MdnsResponseDecoder decoder = new MdnsResponseDecoder(mClock, MATTER_SERVICE_TYPE);
         assertNotNull(matterDuplicateHostname);
 
-        DatagramPacket packet =
-                new DatagramPacket(matterDuplicateHostname, matterDuplicateHostname.length);
-
-        packet.setSocketAddress(
-                new InetSocketAddress(MdnsConstants.getMdnsIPv6Address(), MdnsConstants.MDNS_PORT));
-
-        responses.clear();
-        int errorCode = decoder.decode(
-                packet, responses, /* interfaceIndex= */ 0, mock(Network.class));
-        assertEquals(MdnsResponseDecoder.SUCCESS, errorCode);
+        responses = decode(decoder, matterDuplicateHostname);
 
         // This should emit two records:
         assertEquals(2, responses.size());
@@ -336,21 +310,13 @@
 
     @Test
     @Ignore("MdnsConfigs is not configurable currently.")
-    public void decode_singleHostname_multipleSrvRecords_flagDisabled_singleCompleteResponse() {
+    public void decode_singleHostname_multipleSrvRecords_flagDisabled_singleCompleteResponse()
+            throws Exception {
         //MdnsScannerConfigsFlagsImpl.allowMultipleSrvRecordsPerHost.override(false);
         MdnsResponseDecoder decoder = new MdnsResponseDecoder(mClock, MATTER_SERVICE_TYPE);
         assertNotNull(matterDuplicateHostname);
 
-        DatagramPacket packet =
-                new DatagramPacket(matterDuplicateHostname, matterDuplicateHostname.length);
-
-        packet.setSocketAddress(
-                new InetSocketAddress(MdnsConstants.getMdnsIPv6Address(), MdnsConstants.MDNS_PORT));
-
-        responses.clear();
-        int errorCode = decoder.decode(
-                packet, responses, /* interfaceIndex= */ 0, mock(Network.class));
-        assertEquals(MdnsResponseDecoder.SUCCESS, errorCode);
+        responses = decode(decoder, matterDuplicateHostname);
 
         // This should emit only two records:
         assertEquals(2, responses.size());
@@ -359,4 +325,14 @@
         assertTrue(responses.get(0).isComplete());
         assertFalse(responses.get(1).isComplete());
     }
+
+
+    private List<MdnsResponse> decode(MdnsResponseDecoder decoder, byte[] data)
+            throws MdnsPacket.ParseException {
+        final MdnsPacket parsedPacket = MdnsResponseDecoder.parseResponse(data, data.length);
+        assertNotNull(parsedPacket);
+
+        return decoder.buildResponses(parsedPacket,
+                MdnsSocket.INTERFACE_INDEX_UNSPECIFIED, mock(Network.class));
+    }
 }
\ No newline at end of file
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
index a45ca68..d266b3d 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
@@ -24,13 +24,10 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.inOrder;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -60,8 +57,7 @@
 
 import java.io.IOException;
 import java.net.DatagramPacket;
-import java.net.Inet4Address;
-import java.net.Inet6Address;
+import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -85,6 +81,9 @@
     private static final InetSocketAddress IPV6_ADDRESS = new InetSocketAddress(
             MdnsConstants.getMdnsIPv6Address(), MdnsConstants.MDNS_PORT);
 
+    private static final long TEST_TTL = 120000L;
+    private static final long TEST_ELAPSED_REALTIME = 123L;
+
     @Mock
     private MdnsServiceBrowserListener mockListenerOne;
     @Mock
@@ -95,6 +94,8 @@
     private MdnsMultinetworkSocketClient mockSocketClient;
     @Mock
     private Network mockNetwork;
+    @Mock
+    private MdnsResponseDecoder.Clock mockDecoderClock;
     @Captor
     private ArgumentCaptor<MdnsServiceInfo> serviceInfoCaptor;
 
@@ -111,6 +112,7 @@
     @SuppressWarnings("DoNotMock")
     public void setUp() throws IOException {
         MockitoAnnotations.initMocks(this);
+        doReturn(TEST_ELAPSED_REALTIME).when(mockDecoderClock).elapsedRealtime();
 
         expectedIPv4Packets = new DatagramPacket[16];
         expectedIPv6Packets = new DatagramPacket[16];
@@ -160,7 +162,8 @@
                 .thenReturn(expectedIPv6Packets[15]);
 
         client =
-                new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor) {
+                new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
+                        mockDecoderClock) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -424,6 +427,7 @@
         assertEquals(port, serviceInfo.getPort());
         assertEquals(subTypes, serviceInfo.getSubtypes());
         for (String key : attributes.keySet()) {
+            assertTrue(attributes.containsKey(key));
             assertEquals(attributes.get(key), serviceInfo.getAttributeByKey(key));
         }
         assertEquals(interfaceIndex, serviceInfo.getInterfaceIndex());
@@ -434,22 +438,19 @@
     public void processResponse_incompleteResponse() {
         client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
 
-        MdnsResponse response = mock(MdnsResponse.class);
-        when(response.getServiceInstanceName()).thenReturn("service-instance-1");
-        doReturn(INTERFACE_INDEX).when(response).getInterfaceIndex();
-        doReturn(mockNetwork).when(response).getNetwork();
-        when(response.isComplete()).thenReturn(false);
-
-        client.processResponse(response);
+        client.processResponse(createResponse(
+                "service-instance-1", null /* host */, 0 /* port */,
+                SERVICE_TYPE_LABELS,
+                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
         verify(mockListenerOne).onServiceNameDiscovered(serviceInfoCaptor.capture());
         verifyServiceInfo(serviceInfoCaptor.getAllValues().get(0),
                 "service-instance-1",
                 SERVICE_TYPE_LABELS,
-                null /* ipv4Address */,
-                null /* ipv6Address */,
-                0 /* port */,
-                List.of() /* subTypes */,
-                Collections.singletonMap("key", null) /* attributes */,
+                /* ipv4Address= */ null,
+                /* ipv6Address= */ null,
+                /* port= */ 0,
+                /* subTypes= */ List.of(),
+                Collections.emptyMap(),
                 INTERFACE_INDEX,
                 mockNetwork);
 
@@ -463,28 +464,17 @@
         client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
 
         // Process the initial response.
-        MdnsResponse initialResponse =
-                createResponse(
-                        "service-instance-1",
-                        ipV4Address,
-                        5353,
-                        /* subtype= */ "ABCDE",
-                        Collections.emptyMap(),
-                        /* interfaceIndex= */ 20,
-                        mockNetwork);
-        client.processResponse(initialResponse);
+        client.processResponse(createResponse(
+                "service-instance-1", ipV4Address, 5353,
+                /* subtype= */ "ABCDE",
+                Collections.emptyMap(), TEST_TTL), /* interfaceIndex= */ 20, mockNetwork);
 
         // Process a second response with a different port and updated text attributes.
-        MdnsResponse secondResponse =
-                createResponse(
-                        "service-instance-1",
-                        ipV4Address,
-                        5354,
-                        /* subtype= */ "ABCDE",
-                        Collections.singletonMap("key", "value"),
-                        /* interfaceIndex= */ 20,
-                        mockNetwork);
-        client.processResponse(secondResponse);
+        client.processResponse(createResponse(
+                "service-instance-1", ipV4Address, 5354,
+                /* subtype= */ "ABCDE",
+                Collections.singletonMap("key", "value"), TEST_TTL),
+                /* interfaceIndex= */ 20, mockNetwork);
 
         // Verify onServiceNameDiscovered was called once for the initial response.
         verify(mockListenerOne).onServiceNameDiscovered(serviceInfoCaptor.capture());
@@ -529,31 +519,17 @@
         client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
 
         // Process the initial response.
-        MdnsResponse initialResponse =
-                createResponse(
-                        "service-instance-1",
-                        ipV6Address,
-                        5353,
-                        /* subtype= */ "ABCDE",
-                        Collections.emptyMap(),
-                        /* interfaceIndex= */ 20,
-                        mockNetwork);
-        client.processResponse(initialResponse);
+        client.processResponse(createResponse(
+                "service-instance-1", ipV6Address, 5353,
+                /* subtype= */ "ABCDE",
+                Collections.emptyMap(), TEST_TTL), /* interfaceIndex= */ 20, mockNetwork);
 
         // Process a second response with a different port and updated text attributes.
-        MdnsResponse secondResponse =
-                createResponse(
-                        "service-instance-1",
-                        ipV6Address,
-                        5354,
-                        /* subtype= */ "ABCDE",
-                        Collections.singletonMap("key", "value"),
-                        /* interfaceIndex= */ 20,
-                        mockNetwork);
-        client.processResponse(secondResponse);
-
-        System.out.println("secondResponses ip"
-                + secondResponse.getInet6AddressRecord().getInet6Address().getHostAddress());
+        client.processResponse(createResponse(
+                "service-instance-1", ipV6Address, 5354,
+                /* subtype= */ "ABCDE",
+                Collections.singletonMap("key", "value"), TEST_TTL),
+                /* interfaceIndex= */ 20, mockNetwork);
 
         // Verify onServiceNameDiscovered was called once for the initial response.
         verify(mockListenerOne).onServiceNameDiscovered(serviceInfoCaptor.capture());
@@ -619,29 +595,25 @@
         final String serviceName = "service-instance-1";
         final String ipV6Address = "2000:3333::da6c:63ff:fe7c:7483";
         // Process the initial response.
-        final MdnsResponse initialResponse =
-                createResponse(
-                        serviceName,
-                        ipV6Address,
-                        5353 /* port */,
-                        /* subtype= */ "ABCDE",
-                        Collections.emptyMap(),
-                        INTERFACE_INDEX,
-                        mockNetwork);
-        client.processResponse(initialResponse);
-        MdnsResponse response = mock(MdnsResponse.class);
-        doReturn("goodbye-service").when(response).getServiceInstanceName();
-        doReturn(INTERFACE_INDEX).when(response).getInterfaceIndex();
-        doReturn(mockNetwork).when(response).getNetwork();
-        doReturn(true).when(response).isGoodbye();
-        client.processResponse(response);
+        client.processResponse(createResponse(
+                serviceName, ipV6Address, 5353,
+                SERVICE_TYPE_LABELS,
+                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
+
+        client.processResponse(createResponse(
+                "goodbye-service", ipV6Address, 5353,
+                SERVICE_TYPE_LABELS,
+                Collections.emptyMap(), /* ptrTtlMillis= */ 0L), INTERFACE_INDEX, mockNetwork);
+
         // Verify removed callback won't be called if the service is not existed.
         verifyServiceRemovedNoCallback(mockListenerOne);
         verifyServiceRemovedNoCallback(mockListenerTwo);
 
         // Verify removed callback would be called.
-        doReturn(serviceName).when(response).getServiceInstanceName();
-        client.processResponse(response);
+        client.processResponse(createResponse(
+                serviceName, ipV6Address, 5353,
+                SERVICE_TYPE_LABELS,
+                Collections.emptyMap(), 0L), INTERFACE_INDEX, mockNetwork);
         verifyServiceRemovedCallback(
                 mockListenerOne, serviceName, SERVICE_TYPE_LABELS, INTERFACE_INDEX, mockNetwork);
         verifyServiceRemovedCallback(
@@ -651,16 +623,10 @@
     @Test
     public void reportExistingServiceToNewlyRegisteredListeners() throws Exception {
         // Process the initial response.
-        MdnsResponse initialResponse =
-                createResponse(
-                        "service-instance-1",
-                        "192.168.1.1",
-                        5353,
-                        /* subtype= */ "ABCDE",
-                        Collections.emptyMap(),
-                        INTERFACE_INDEX,
-                        mockNetwork);
-        client.processResponse(initialResponse);
+        client.processResponse(createResponse(
+                "service-instance-1", "192.168.1.1", 5353,
+                /* subtype= */ "ABCDE",
+                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
 
         client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
 
@@ -687,10 +653,10 @@
         assertNull(existingServiceInfo.getAttributeByKey("key"));
 
         // Process a goodbye message for the existing response.
-        MdnsResponse goodByeResponse = mock(MdnsResponse.class);
-        when(goodByeResponse.getServiceInstanceName()).thenReturn("service-instance-1");
-        when(goodByeResponse.isGoodbye()).thenReturn(true);
-        client.processResponse(goodByeResponse);
+        client.processResponse(createResponse(
+                "service-instance-1", "192.168.1.1", 5353,
+                SERVICE_TYPE_LABELS,
+                Collections.emptyMap(), /* ptrTtlMillis= */ 0L), INTERFACE_INDEX, mockNetwork);
 
         client.startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
 
@@ -709,17 +675,15 @@
         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
 
         // Process the initial response.
-        MdnsResponse initialResponse =
-                createMockResponse(
-                        serviceInstanceName, "192.168.1.1", 5353, List.of("ABCDE"),
-                        Map.of(), INTERFACE_INDEX, mockNetwork);
-        client.processResponse(initialResponse);
+        client.processResponse(createResponse(
+                serviceInstanceName, "192.168.1.1", 5353, /* subtype= */ "ABCDE",
+                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
 
         // Clear the scheduled runnable.
         currentThreadExecutor.getAndClearLastScheduledRunnable();
 
         // Simulate the case where the response is after TTL.
-        when(initialResponse.getServiceRecord().getRemainingTTL(anyLong())).thenReturn((long) 0);
+        doReturn(TEST_ELAPSED_REALTIME + TEST_TTL + 1L).when(mockDecoderClock).elapsedRealtime();
         firstMdnsTask.run();
 
         // Verify removed callback was not called.
@@ -733,7 +697,8 @@
         //MdnsConfigsFlagsImpl.allowSearchOptionsToRemoveExpiredService.override(true);
         final String serviceInstanceName = "service-instance-1";
         client =
-                new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor) {
+                new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
+                        mockDecoderClock) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -743,24 +708,22 @@
         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
 
         // Process the initial response.
-        MdnsResponse initialResponse =
-                createMockResponse(
-                        serviceInstanceName, "192.168.1.1", 5353, List.of("ABCDE"),
-                        Map.of(), INTERFACE_INDEX, mockNetwork);
-        client.processResponse(initialResponse);
+        client.processResponse(createResponse(
+                serviceInstanceName, "192.168.1.1", 5353, /* subtype= */ "ABCDE",
+                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
 
         // Clear the scheduled runnable.
         currentThreadExecutor.getAndClearLastScheduledRunnable();
 
         // Simulate the case where the response is under TTL.
-        when(initialResponse.getServiceRecord().getRemainingTTL(anyLong())).thenReturn((long) 1000);
+        doReturn(TEST_ELAPSED_REALTIME + TEST_TTL - 1L).when(mockDecoderClock).elapsedRealtime();
         firstMdnsTask.run();
 
         // Verify removed callback was not called.
         verifyServiceRemovedNoCallback(mockListenerOne);
 
         // Simulate the case where the response is after TTL.
-        when(initialResponse.getServiceRecord().getRemainingTTL(anyLong())).thenReturn((long) 0);
+        doReturn(TEST_ELAPSED_REALTIME + TEST_TTL + 1L).when(mockDecoderClock).elapsedRealtime();
         firstMdnsTask.run();
 
         // Verify removed callback was called.
@@ -773,7 +736,8 @@
             throws Exception {
         final String serviceInstanceName = "service-instance-1";
         client =
-                new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor) {
+                new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
+                        mockDecoderClock) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -783,17 +747,15 @@
         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
 
         // Process the initial response.
-        MdnsResponse initialResponse =
-                createMockResponse(
-                        serviceInstanceName, "192.168.1.1", 5353, List.of("ABCDE"),
-                        Map.of(), INTERFACE_INDEX, mockNetwork);
-        client.processResponse(initialResponse);
+        client.processResponse(createResponse(
+                serviceInstanceName, "192.168.1.1", 5353, /* subtype= */ "ABCDE",
+                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
 
         // Clear the scheduled runnable.
         currentThreadExecutor.getAndClearLastScheduledRunnable();
 
         // Simulate the case where the response is after TTL.
-        when(initialResponse.getServiceRecord().getRemainingTTL(anyLong())).thenReturn((long) 0);
+        doReturn(TEST_ELAPSED_REALTIME + TEST_TTL + 1L).when(mockDecoderClock).elapsedRealtime();
         firstMdnsTask.run();
 
         // Verify removed callback was not called.
@@ -807,7 +769,8 @@
         //MdnsConfigsFlagsImpl.removeServiceAfterTtlExpires.override(true);
         final String serviceInstanceName = "service-instance-1";
         client =
-                new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor) {
+                new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
+                        mockDecoderClock) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -817,17 +780,15 @@
         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
 
         // Process the initial response.
-        MdnsResponse initialResponse =
-                createMockResponse(
-                        serviceInstanceName, "192.168.1.1", 5353, List.of("ABCDE"),
-                        Map.of(), INTERFACE_INDEX, mockNetwork);
-        client.processResponse(initialResponse);
+        client.processResponse(createResponse(
+                serviceInstanceName, "192.168.1.1", 5353, /* subtype= */ "ABCDE",
+                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
 
         // Clear the scheduled runnable.
         currentThreadExecutor.getAndClearLastScheduledRunnable();
 
         // Simulate the case where the response is after TTL.
-        when(initialResponse.getServiceRecord().getRemainingTTL(anyLong())).thenReturn((long) 0);
+        doReturn(TEST_ELAPSED_REALTIME + TEST_TTL + 1L).when(mockDecoderClock).elapsedRealtime();
         firstMdnsTask.run();
 
         // Verify removed callback was called.
@@ -844,46 +805,26 @@
         InOrder inOrder = inOrder(mockListenerOne);
 
         // Process the initial response which is incomplete.
-        final MdnsResponse initialResponse =
-                createResponse(
-                        serviceName,
-                        null,
-                        5353,
-                        "ABCDE" /* subtype */,
-                        Collections.emptyMap(),
-                        INTERFACE_INDEX,
-                        mockNetwork);
-        client.processResponse(initialResponse);
+        final String subtype = "ABCDE";
+        client.processResponse(createResponse(
+                serviceName, null, 5353, subtype,
+                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
 
         // Process a second response which has ip address to make response become complete.
-        final MdnsResponse secondResponse =
-                createResponse(
-                        serviceName,
-                        ipV4Address,
-                        5353,
-                        "ABCDE" /* subtype */,
-                        Collections.emptyMap(),
-                        INTERFACE_INDEX,
-                        mockNetwork);
-        client.processResponse(secondResponse);
+        client.processResponse(createResponse(
+                serviceName, ipV4Address, 5353, subtype,
+                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
 
         // Process a third response with a different ip address, port and updated text attributes.
-        final MdnsResponse thirdResponse =
-                createResponse(
-                        serviceName,
-                        ipV6Address,
-                        5354,
-                        "ABCDE" /* subtype */,
-                        Collections.singletonMap("key", "value"),
-                        INTERFACE_INDEX,
-                        mockNetwork);
-        client.processResponse(thirdResponse);
+        client.processResponse(createResponse(
+                serviceName, ipV6Address, 5354, subtype,
+                Collections.singletonMap("key", "value"), TEST_TTL), INTERFACE_INDEX, mockNetwork);
 
-        // Process the last response which is goodbye message.
-        final MdnsResponse lastResponse = mock(MdnsResponse.class);
-        doReturn(serviceName).when(lastResponse).getServiceInstanceName();
-        doReturn(true).when(lastResponse).isGoodbye();
-        client.processResponse(lastResponse);
+        // Process the last response which is goodbye message (with the main type, not subtype).
+        client.processResponse(createResponse(
+                serviceName, ipV6Address, 5354, SERVICE_TYPE_LABELS,
+                Collections.singletonMap("key", "value"), /* ptrTtlMillis= */ 0L),
+                INTERFACE_INDEX, mockNetwork);
 
         // Verify onServiceNameDiscovered was first called for the initial response.
         inOrder.verify(mockListenerOne).onServiceNameDiscovered(serviceInfoCaptor.capture());
@@ -893,7 +834,7 @@
                 null /* ipv4Address */,
                 null /* ipv6Address */,
                 5353 /* port */,
-                Collections.singletonList("ABCDE") /* subTypes */,
+                Collections.singletonList(subtype) /* subTypes */,
                 Collections.singletonMap("key", null) /* attributes */,
                 INTERFACE_INDEX,
                 mockNetwork);
@@ -906,7 +847,7 @@
                 ipV4Address /* ipv4Address */,
                 null /* ipv6Address */,
                 5353 /* port */,
-                Collections.singletonList("ABCDE") /* subTypes */,
+                Collections.singletonList(subtype) /* subTypes */,
                 Collections.singletonMap("key", null) /* attributes */,
                 INTERFACE_INDEX,
                 mockNetwork);
@@ -919,7 +860,7 @@
                 ipV4Address /* ipv4Address */,
                 ipV6Address /* ipv6Address */,
                 5354 /* port */,
-                Collections.singletonList("ABCDE") /* subTypes */,
+                Collections.singletonList(subtype) /* subTypes */,
                 Collections.singletonMap("key", "value") /* attributes */,
                 INTERFACE_INDEX,
                 mockNetwork);
@@ -1029,106 +970,68 @@
         }
     }
 
-    // Creates a mock mDNS response.
-    private MdnsResponse createMockResponse(
-            @NonNull String serviceInstanceName,
-            @NonNull String host,
-            int port,
-            @NonNull List<String> subtypes,
-            @NonNull Map<String, String> textAttributes,
-            int interfaceIndex,
-            Network network)
-            throws Exception {
-        String[] hostName = new String[]{"hostname"};
-        MdnsServiceRecord serviceRecord = mock(MdnsServiceRecord.class);
-        when(serviceRecord.getServiceHost()).thenReturn(hostName);
-        when(serviceRecord.getServicePort()).thenReturn(port);
-
-        MdnsResponse response = spy(new MdnsResponse(0, interfaceIndex, network));
-
-        MdnsInetAddressRecord inetAddressRecord = mock(MdnsInetAddressRecord.class);
-        if (host.contains(":")) {
-            when(inetAddressRecord.getInet6Address())
-                    .thenReturn((Inet6Address) Inet6Address.getByName(host));
-            response.setInet6AddressRecord(inetAddressRecord);
-        } else {
-            when(inetAddressRecord.getInet4Address())
-                    .thenReturn((Inet4Address) Inet4Address.getByName(host));
-            response.setInet4AddressRecord(inetAddressRecord);
-        }
-
-        MdnsTextRecord textRecord = mock(MdnsTextRecord.class);
-        List<String> textStrings = new ArrayList<>();
-        List<TextEntry> textEntries = new ArrayList<>();
-        for (Map.Entry<String, String> kv : textAttributes.entrySet()) {
-            textStrings.add(kv.getKey() + "=" + kv.getValue());
-            textEntries.add(new TextEntry(kv.getKey(), kv.getValue().getBytes(UTF_8)));
-        }
-        when(textRecord.getStrings()).thenReturn(textStrings);
-        when(textRecord.getEntries()).thenReturn(textEntries);
-
-        response.setServiceRecord(serviceRecord);
-        response.setTextRecord(textRecord);
-
-        doReturn(false).when(response).isGoodbye();
-        doReturn(true).when(response).isComplete();
-        doReturn(serviceInstanceName).when(response).getServiceInstanceName();
-        doReturn(new ArrayList<>(subtypes)).when(response).getSubtypes();
-        return response;
-    }
-
-    // Creates a mDNS response.
-    private MdnsResponse createResponse(
+    private MdnsPacket createResponse(
             @NonNull String serviceInstanceName,
             @Nullable String host,
             int port,
             @NonNull String subtype,
             @NonNull Map<String, String> textAttributes,
-            int interfaceIndex,
-            Network network)
+            long ptrTtlMillis)
             throws Exception {
-        MdnsResponse response = new MdnsResponse(0, interfaceIndex, network);
+        final ArrayList<String> type = new ArrayList<>();
+        type.add(subtype);
+        type.add(MdnsConstants.SUBTYPE_LABEL);
+        type.addAll(Arrays.asList(SERVICE_TYPE_LABELS));
+        return createResponse(serviceInstanceName, host, port, type.toArray(new String[0]),
+                textAttributes, ptrTtlMillis);
+    }
+
+    // Creates a mDNS response.
+    private MdnsPacket createResponse(
+            @NonNull String serviceInstanceName,
+            @Nullable String host,
+            int port,
+            @NonNull String[] type,
+            @NonNull Map<String, String> textAttributes,
+            long ptrTtlMillis) {
+
+        final ArrayList<MdnsRecord> answerRecords = new ArrayList<>();
 
         // Set PTR record
+        final ArrayList<String> serviceNameList = new ArrayList<>();
+        serviceNameList.add(serviceInstanceName);
+        serviceNameList.addAll(Arrays.asList(type));
+        final String[] serviceName = serviceNameList.toArray(new String[0]);
         final MdnsPointerRecord pointerRecord = new MdnsPointerRecord(
-                new String[]{subtype, MdnsConstants.SUBTYPE_LABEL, "test"} /* name */,
-                0L /* receiptTimeMillis */,
+                type,
+                TEST_ELAPSED_REALTIME /* receiptTimeMillis */,
                 false /* cacheFlush */,
-                120000L /* ttlMillis */,
-                new String[]{serviceInstanceName});
-        response.addPointerRecord(pointerRecord);
+                ptrTtlMillis,
+                serviceName);
+        answerRecords.add(pointerRecord);
 
         // Set SRV record.
         final MdnsServiceRecord serviceRecord = new MdnsServiceRecord(
-                new String[] {"service"} /* name */,
-                0L /* receiptTimeMillis */,
+                serviceName,
+                TEST_ELAPSED_REALTIME /* receiptTimeMillis */,
                 false /* cacheFlush */,
-                120000L /* ttlMillis */,
+                TEST_TTL,
                 0 /* servicePriority */,
                 0 /* serviceWeight */,
                 port,
                 new String[]{"hostname"});
-        response.setServiceRecord(serviceRecord);
+        answerRecords.add(serviceRecord);
 
         // Set A/AAAA record.
         if (host != null) {
-            if (InetAddresses.parseNumericAddress(host) instanceof Inet6Address) {
-                final MdnsInetAddressRecord inetAddressRecord = new MdnsInetAddressRecord(
-                        new String[] {"address"} /* name */,
-                        0L /* receiptTimeMillis */,
-                        false /* cacheFlush */,
-                        120000L /* ttlMillis */,
-                        Inet6Address.getByName(host));
-                response.setInet6AddressRecord(inetAddressRecord);
-            } else {
-                final MdnsInetAddressRecord inetAddressRecord = new MdnsInetAddressRecord(
-                        new String[] {"address"} /* name */,
-                        0L /* receiptTimeMillis */,
-                        false /* cacheFlush */,
-                        120000L /* ttlMillis */,
-                        Inet4Address.getByName(host));
-                response.setInet4AddressRecord(inetAddressRecord);
-            }
+            final InetAddress addr = InetAddresses.parseNumericAddress(host);
+            final MdnsInetAddressRecord inetAddressRecord = new MdnsInetAddressRecord(
+                    new String[] {"hostname"} /* name */,
+                    TEST_ELAPSED_REALTIME /* receiptTimeMillis */,
+                    false /* cacheFlush */,
+                    TEST_TTL,
+                    addr);
+            answerRecords.add(inetAddressRecord);
         }
 
         // Set TXT record.
@@ -1137,12 +1040,18 @@
             textEntries.add(new TextEntry(kv.getKey(), kv.getValue().getBytes(UTF_8)));
         }
         final MdnsTextRecord textRecord = new MdnsTextRecord(
-                new String[] {"text"} /* name */,
-                0L /* receiptTimeMillis */,
+                serviceName,
+                TEST_ELAPSED_REALTIME /* receiptTimeMillis */,
                 false /* cacheFlush */,
-                120000L /* ttlMillis */,
+                TEST_TTL,
                 textEntries);
-        response.setTextRecord(textRecord);
-        return response;
+        answerRecords.add(textRecord);
+        return new MdnsPacket(
+                0 /* flags */,
+                Collections.emptyList() /* questions */,
+                answerRecords,
+                Collections.emptyList() /* authorityRecords */,
+                Collections.emptyList() /* additionalRecords */
+        );
     }
 }
\ No newline at end of file
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketClientTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketClientTests.java
index 1d61cd3..9048686 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketClientTests.java
@@ -18,13 +18,11 @@
 
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
-import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.never;
@@ -48,7 +46,6 @@
 import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
 import org.mockito.ArgumentMatchers;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
@@ -373,7 +370,7 @@
         mdnsClient.startDiscovery();
 
         verify(mockCallback, timeout(TIMEOUT).atLeast(1))
-                .onResponseReceived(any(MdnsResponse.class));
+                .onResponseReceived(any(MdnsPacket.class), anyInt(), any());
     }
 
     @Test
@@ -382,7 +379,7 @@
         mdnsClient.startDiscovery();
 
         verify(mockCallback, timeout(TIMEOUT).atLeastOnce())
-                .onResponseReceived(any(MdnsResponse.class));
+                .onResponseReceived(any(MdnsPacket.class), anyInt(), any());
 
         mdnsClient.stopDiscovery();
     }
@@ -514,7 +511,7 @@
         mdnsClient.startDiscovery();
 
         verify(mockCallback, timeout(TIMEOUT).atLeastOnce())
-                .onResponseReceived(argThat(response -> response.getInterfaceIndex() == 21));
+                .onResponseReceived(any(), eq(21), any());
     }
 
     @Test
@@ -536,11 +533,7 @@
         mdnsClient.setCallback(mockCallback);
         mdnsClient.startDiscovery();
 
-        ArgumentCaptor<MdnsResponse> mdnsResponseCaptor =
-                ArgumentCaptor.forClass(MdnsResponse.class);
         verify(mockMulticastSocket, never()).getInterfaceIndex();
-        verify(mockCallback, timeout(TIMEOUT).atLeast(1))
-                .onResponseReceived(mdnsResponseCaptor.capture());
-        assertEquals(-1, mdnsResponseCaptor.getValue().getInterfaceIndex());
+        verify(mockCallback, timeout(TIMEOUT).atLeast(1)).onResponseReceived(any(), eq(-1), any());
     }
 }
\ No newline at end of file