Merge "Fix tests in TrafficStatsTest that are affected by adb over network"
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/PacketReflector.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/PacketReflector.java
index a4a2956..124c2c3 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/PacketReflector.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/PacketReflector.java
@@ -16,6 +16,11 @@
 
 package com.android.cts.net.hostside;
 
+import static android.system.OsConstants.ICMP6_ECHO_REPLY;
+import static android.system.OsConstants.ICMP6_ECHO_REQUEST;
+import static android.system.OsConstants.ICMP_ECHO;
+import static android.system.OsConstants.ICMP_ECHOREPLY;
+
 import android.system.ErrnoException;
 import android.system.Os;
 import android.util.Log;
@@ -47,8 +52,6 @@
 
     private static final byte ICMP_ECHO = 8;
     private static final byte ICMP_ECHOREPLY = 0;
-    private static final byte ICMPV6_ECHO_REQUEST = (byte) 128;
-    private static final byte ICMPV6_ECHO_REPLY = (byte) 129;
 
     private static String TAG = "PacketReflector";
 
@@ -125,7 +128,7 @@
 
         byte type = buf[hdrLen];
         if (!(version == 4 && type == ICMP_ECHO) &&
-            !(version == 6 && type == ICMPV6_ECHO_REQUEST)) {
+            !(version == 6 && type == (byte) ICMP6_ECHO_REQUEST)) {
             return;
         }
 
@@ -145,10 +148,18 @@
             return;
         }
 
+        byte replyType = buf[hdrLen];
+        if ((type == ICMP_ECHO && replyType != ICMP_ECHOREPLY)
+                || (type == (byte) ICMP6_ECHO_REQUEST && replyType != (byte) ICMP6_ECHO_REPLY)) {
+            Log.i(TAG, "Received unexpected ICMP reply: original " + type
+                    + ", reply " + replyType);
+            return;
+        }
+
         // Compare the response we got with the original packet.
         // The only thing that should have changed are addresses, type and checksum.
         // Overwrite them with the received bytes and see if the packet is otherwise identical.
-        request[hdrLen] = buf[hdrLen];          // Type.
+        request[hdrLen] = buf[hdrLen];          // Type
         request[hdrLen + 2] = buf[hdrLen + 2];  // Checksum byte 1.
         request[hdrLen + 3] = buf[hdrLen + 3];  // Checksum byte 2.
 
diff --git a/tests/cts/net/src/android/net/cts/DnsResolverTest.java b/tests/cts/net/src/android/net/cts/DnsResolverTest.java
index f6cc768..0ff6cd8 100644
--- a/tests/cts/net/src/android/net/cts/DnsResolverTest.java
+++ b/tests/cts/net/src/android/net/cts/DnsResolverTest.java
@@ -370,7 +370,6 @@
 
         @Override
         public void onQueryException(@NonNull ErrnoException e) {
-            if (mCancelSignal.isCanceled() && e.errno == EBADF) return;
             fail(mMsg + e.getMessage());
         }
     }
@@ -380,7 +379,7 @@
         final String msg = "Test cancel query " + dname;
         // Start a DNS query and the cancel it immediately. Use VerifyCancelCallback to expect
         // that the query is cancelled before it succeeds. If it is not cancelled before it
-        // succeeds, retry the until it is.
+        // succeeds, retry the test until it is.
         for (Network network : getTestableNetworks()) {
             boolean retry = false;
             int round = 0;
@@ -426,7 +425,7 @@
         final String msg = "Test cancel raw Query " + byteArrayToHexString(blob);
         // Start a DNS query and the cancel it immediately. Use VerifyCancelCallback to expect
         // that the query is cancelled before it succeeds. If it is not cancelled before it
-        // succeeds, retry the until it is.
+        // succeeds, retry the test until it is.
         for (Network network : getTestableNetworks()) {
             boolean retry = false;
             int round = 0;
@@ -470,4 +469,108 @@
             }
         }
     }
+
+    /**
+     * A query callback for InetAddress that ensures that the query is
+     * cancelled and that onAnswer is never called. If the query succeeds
+     * before it is cancelled, needRetry will return true so the
+     * test can retry.
+     */
+    class VerifyCancelInetAddressCallback extends DnsResolver.InetAddressAnswerCallback {
+        private static final int CANCEL_TIMEOUT = 3_000;
+
+        private final CountDownLatch mLatch = new CountDownLatch(1);
+        private final String mMsg;
+        private final List<InetAddress> mAnswers;
+        private final CancellationSignal mCancelSignal;
+
+        VerifyCancelInetAddressCallback(@NonNull String msg, @Nullable CancellationSignal cancel) {
+            this.mMsg = msg;
+            this.mCancelSignal = cancel;
+            mAnswers = new ArrayList<>();
+        }
+
+        public boolean waitForAnswer() throws InterruptedException {
+            return mLatch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS);
+        }
+
+        public boolean needRetry() throws InterruptedException {
+            return mLatch.await(CANCEL_TIMEOUT, TimeUnit.MILLISECONDS);
+        }
+
+        public boolean isAnswerEmpty() {
+            return mAnswers.isEmpty();
+        }
+
+        @Override
+        public void onAnswer(@NonNull List<InetAddress> answerList) {
+            if (mCancelSignal != null && mCancelSignal.isCanceled()) {
+                fail(mMsg + " should not have returned any answers");
+            }
+            mAnswers.clear();
+            mAnswers.addAll(answerList);
+            mLatch.countDown();
+        }
+
+        @Override
+        public void onParseException(@NonNull ParseException e) {
+            fail(mMsg + e.getMessage());
+        }
+
+        @Override
+        public void onQueryException(@NonNull ErrnoException e) {
+            fail(mMsg + e.getMessage());
+        }
+    }
+
+    public void testQueryForInetAddress() {
+        final String dname = "www.google.com";
+        final String msg = "Test query for InetAddress " + dname;
+        for (Network network : getTestableNetworks()) {
+            final VerifyCancelInetAddressCallback callback =
+                    new VerifyCancelInetAddressCallback(msg, null);
+            mDns.query(network, dname, FLAG_NO_CACHE_LOOKUP,
+                    mExecutor, null, callback);
+            try {
+                assertTrue(msg + " but no answer after " + TIMEOUT_MS + "ms.",
+                        callback.waitForAnswer());
+                assertTrue(msg + " returned 0 results", !callback.isAnswerEmpty());
+            } catch (InterruptedException e) {
+                fail(msg + " Waiting for DNS lookup was interrupted");
+            }
+        }
+    }
+
+    public void testQueryCancelForInetAddress() throws ErrnoException {
+        final String dname = "www.google.com";
+        final String msg = "Test cancel query for InetAddress " + dname;
+        // Start a DNS query and the cancel it immediately. Use VerifyCancelCallback to expect
+        // that the query is cancelled before it succeeds. If it is not cancelled before it
+        // succeeds, retry the test until it is.
+        for (Network network : getTestableNetworks()) {
+            boolean retry = false;
+            int round = 0;
+            do {
+                if (++round > CANCEL_RETRY_TIMES) {
+                    fail(msg + " cancel failed " + CANCEL_RETRY_TIMES + " times");
+                }
+                final CountDownLatch latch = new CountDownLatch(1);
+                final CancellationSignal cancelSignal = new CancellationSignal();
+                final VerifyCancelInetAddressCallback callback =
+                        new VerifyCancelInetAddressCallback(msg, cancelSignal);
+                mDns.query(network, dname, FLAG_EMPTY, mExecutor, cancelSignal, callback);
+                mExecutor.execute(() -> {
+                    cancelSignal.cancel();
+                    latch.countDown();
+                });
+                try {
+                    retry = callback.needRetry();
+                    assertTrue(msg + " query was not cancelled",
+                        latch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
+                } catch (InterruptedException e) {
+                    fail(msg + "Waiting for DNS lookup was interrupted");
+                }
+            } while (retry);
+        }
+    }
 }
diff --git a/tests/cts/net/src/android/net/ipv6/cts/PingTest.java b/tests/cts/net/src/android/net/ipv6/cts/PingTest.java
index c23ad30..146fd83 100644
--- a/tests/cts/net/src/android/net/ipv6/cts/PingTest.java
+++ b/tests/cts/net/src/android/net/ipv6/cts/PingTest.java
@@ -61,7 +61,7 @@
 
     /** The beginning of an ICMPv6 echo request: type, code, and uninitialized checksum. */
     private static final byte[] PING_HEADER = new byte[] {
-        (byte) 0x80, (byte) 0x00, (byte) 0x00, (byte) 0x00
+        (byte) ICMP6_ECHO_REQUEST, (byte) 0x00, (byte) 0x00, (byte) 0x00
     };
 
     /**
@@ -135,7 +135,7 @@
         byte[] response = new byte[bytesRead];
         responseBuffer.flip();
         responseBuffer.get(response, 0, bytesRead);
-        assertEquals((byte) 0x81, response[0]);
+        assertEquals((byte) ICMP6_ECHO_REPLY, response[0]);
 
         // Find out what ICMP ID was used in the packet that was sent.
         int id = ((InetSocketAddress) Os.getsockname(s)).getPort();