Merge "Update CS so that per-app OEM APIs can be tested" into sc-dev
diff --git a/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java b/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java
index 36a2f10..4a05c9f 100644
--- a/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java
+++ b/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java
@@ -786,6 +786,7 @@
         dumpIpv6ForwardingRules(pw);
         dumpIpv4ForwardingRules(pw);
         pw.decreaseIndent();
+        pw.println();
 
         pw.println("Device map:");
         pw.increaseIndent();
@@ -872,24 +873,33 @@
         }
     }
 
-    private String ipv4RuleToString(long now, Tether4Key key, Tether4Value value) {
-        final String private4, public4, dst4;
+    private String ipv4RuleToString(long now, boolean downstream,
+            Tether4Key key, Tether4Value value) {
+        final String src4, public4, dst4;
+        final int publicPort;
         try {
-            private4 = InetAddress.getByAddress(key.src4).getHostAddress();
-            dst4 = InetAddress.getByAddress(key.dst4).getHostAddress();
-            public4 = InetAddress.getByAddress(value.src46).getHostAddress();
+            src4 = InetAddress.getByAddress(key.src4).getHostAddress();
+            if (downstream) {
+                public4 = InetAddress.getByAddress(key.dst4).getHostAddress();
+                publicPort = key.dstPort;
+            } else {
+                public4 = InetAddress.getByAddress(value.src46).getHostAddress();
+                publicPort = value.srcPort;
+            }
+            dst4 = InetAddress.getByAddress(value.dst46).getHostAddress();
         } catch (UnknownHostException impossible) {
-            throw new AssertionError("4-byte array not valid IPv4 address!");
+            throw new AssertionError("IP address array not valid IPv4 address!");
         }
-        long ageMs = (now - value.lastUsed) / 1_000_000;
-        return String.format("[%s] %d(%s) %s:%d -> %d(%s) %s:%d -> %s:%d %dms",
-                key.dstMac, key.iif, getIfName(key.iif), private4, key.srcPort,
+
+        final long ageMs = (now - value.lastUsed) / 1_000_000;
+        return String.format("[%s] %d(%s) %s:%d -> %d(%s) %s:%d -> %s:%d [%s] %dms",
+                key.dstMac, key.iif, getIfName(key.iif), src4, key.srcPort,
                 value.oif, getIfName(value.oif),
-                public4, value.srcPort, dst4, key.dstPort, ageMs);
+                public4, publicPort, dst4, value.dstPort, value.ethDstMac, ageMs);
     }
 
-    private void dumpIpv4ForwardingRuleMap(long now, BpfMap<Tether4Key, Tether4Value> map,
-            IndentingPrintWriter pw) throws ErrnoException {
+    private void dumpIpv4ForwardingRuleMap(long now, boolean downstream,
+            BpfMap<Tether4Key, Tether4Value> map, IndentingPrintWriter pw) throws ErrnoException {
         if (map == null) {
             pw.println("No IPv4 support");
             return;
@@ -898,7 +908,7 @@
             pw.println("No rules");
             return;
         }
-        map.forEach((k, v) -> pw.println(ipv4RuleToString(now, k, v)));
+        map.forEach((k, v) -> pw.println(ipv4RuleToString(now, downstream, k, v)));
     }
 
     private void dumpIpv4ForwardingRules(IndentingPrintWriter pw) {
@@ -906,14 +916,14 @@
 
         try (BpfMap<Tether4Key, Tether4Value> upstreamMap = mDeps.getBpfUpstream4Map();
                 BpfMap<Tether4Key, Tether4Value> downstreamMap = mDeps.getBpfDownstream4Map()) {
-            pw.println("IPv4 Upstream: [inDstMac] iif(iface) src -> nat -> dst");
+            pw.println("IPv4 Upstream: [inDstMac] iif(iface) src -> nat -> dst [outDstMac] age");
             pw.increaseIndent();
-            dumpIpv4ForwardingRuleMap(now, upstreamMap, pw);
+            dumpIpv4ForwardingRuleMap(now, UPSTREAM, upstreamMap, pw);
             pw.decreaseIndent();
 
-            pw.println("IPv4 Downstream: [inDstMac] iif(iface) src -> nat -> dst");
+            pw.println("IPv4 Downstream: [inDstMac] iif(iface) src -> nat -> dst [outDstMac] age");
             pw.increaseIndent();
-            dumpIpv4ForwardingRuleMap(now, downstreamMap, pw);
+            dumpIpv4ForwardingRuleMap(now, DOWNSTREAM, downstreamMap, pw);
             pw.decreaseIndent();
         } catch (ErrnoException e) {
             pw.println("Error dumping IPv4 map: " + e);
diff --git a/tests/common/java/android/net/NetworkProviderTest.kt b/tests/common/java/android/net/NetworkProviderTest.kt
index 7424157..97d3c5a 100644
--- a/tests/common/java/android/net/NetworkProviderTest.kt
+++ b/tests/common/java/android/net/NetworkProviderTest.kt
@@ -18,6 +18,7 @@
 
 import android.app.Instrumentation
 import android.content.Context
+import android.net.NetworkCapabilities.NET_CAPABILITY_TRUSTED
 import android.net.NetworkCapabilities.TRANSPORT_TEST
 import android.net.NetworkProviderTest.TestNetworkCallback.CallbackEntry.OnUnavailable
 import android.net.NetworkProviderTest.TestNetworkProvider.CallbackEntry.OnNetworkRequestWithdrawn
@@ -25,14 +26,18 @@
 import android.os.Build
 import android.os.HandlerThread
 import android.os.Looper
+import android.util.Log
 import androidx.test.InstrumentationRegistry
 import com.android.net.module.util.ArrayTrackRecord
 import com.android.testutils.CompatUtil
+import com.android.testutils.DevSdkIgnoreRule
+import com.android.testutils.DevSdkIgnoreRule.IgnoreAfter
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
 import com.android.testutils.DevSdkIgnoreRunner
 import com.android.testutils.isDevSdkInRange
 import org.junit.After
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mockito.doReturn
@@ -41,6 +46,7 @@
 import java.util.UUID
 import kotlin.test.assertEquals
 import kotlin.test.assertNotEquals
+import kotlin.test.fail
 
 private const val DEFAULT_TIMEOUT_MS = 5000L
 private val instrumentation: Instrumentation
@@ -51,6 +57,8 @@
 @RunWith(DevSdkIgnoreRunner::class)
 @IgnoreUpTo(Build.VERSION_CODES.Q)
 class NetworkProviderTest {
+    @Rule @JvmField
+    val mIgnoreRule = DevSdkIgnoreRule()
     private val mCm = context.getSystemService(ConnectivityManager::class.java)
     private val mHandlerThread = HandlerThread("${javaClass.simpleName} handler thread")
 
@@ -68,6 +76,7 @@
 
     private class TestNetworkProvider(context: Context, looper: Looper) :
             NetworkProvider(context, looper, PROVIDER_NAME) {
+        private val TAG = this::class.simpleName
         private val seenEvents = ArrayTrackRecord<CallbackEntry>().newReadHead()
 
         sealed class CallbackEntry {
@@ -80,22 +89,30 @@
         }
 
         override fun onNetworkRequested(request: NetworkRequest, score: Int, id: Int) {
+            Log.d(TAG, "onNetworkRequested $request, $score, $id")
             seenEvents.add(OnNetworkRequested(request, score, id))
         }
 
         override fun onNetworkRequestWithdrawn(request: NetworkRequest) {
+            Log.d(TAG, "onNetworkRequestWithdrawn $request")
             seenEvents.add(OnNetworkRequestWithdrawn(request))
         }
 
-        inline fun <reified T : CallbackEntry> expectCallback(
+        inline fun <reified T : CallbackEntry> eventuallyExpectCallbackThat(
             crossinline predicate: (T) -> Boolean
         ) = seenEvents.poll(DEFAULT_TIMEOUT_MS) { it is T && predicate(it) }
+                ?: fail("Did not receive callback after ${DEFAULT_TIMEOUT_MS}ms")
     }
 
     private fun createNetworkProvider(ctx: Context = context): TestNetworkProvider {
         return TestNetworkProvider(ctx, mHandlerThread.looper)
     }
 
+    // In S+ framework, do not run this test, since the provider will no longer receive
+    // onNetworkRequested for every request. Instead, provider needs to
+    // call {@code registerNetworkOffer} with the description of networks they
+    // might have ability to setup, and expects {@link NetworkOfferCallback#onNetworkNeeded}.
+    @IgnoreAfter(Build.VERSION_CODES.R)
     @Test
     fun testOnNetworkRequested() {
         val provider = createNetworkProvider()
@@ -105,13 +122,15 @@
 
         val specifier = CompatUtil.makeTestNetworkSpecifier(
                 UUID.randomUUID().toString())
+        // Test network is not allowed to be trusted.
         val nr: NetworkRequest = NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_TEST)
+                .removeCapability(NET_CAPABILITY_TRUSTED)
                 .setNetworkSpecifier(specifier)
                 .build()
         val cb = ConnectivityManager.NetworkCallback()
         mCm.requestNetwork(nr, cb)
-        provider.expectCallback<OnNetworkRequested>() { callback ->
+        provider.eventuallyExpectCallbackThat<OnNetworkRequested>() { callback ->
             callback.request.getNetworkSpecifier() == specifier &&
             callback.request.hasTransport(TRANSPORT_TEST)
         }
@@ -131,22 +150,24 @@
         val config = NetworkAgentConfig.Builder().build()
         val agent = object : NetworkAgent(context, mHandlerThread.looper, "TestAgent", nc, lp,
                 initialScore, config, provider) {}
+        agent.register()
+        agent.markConnected()
 
-        provider.expectCallback<OnNetworkRequested>() { callback ->
+        provider.eventuallyExpectCallbackThat<OnNetworkRequested>() { callback ->
             callback.request.getNetworkSpecifier() == specifier &&
             callback.score == initialScore &&
             callback.id == agent.providerId
         }
 
         agent.sendNetworkScore(updatedScore)
-        provider.expectCallback<OnNetworkRequested>() { callback ->
+        provider.eventuallyExpectCallbackThat<OnNetworkRequested>() { callback ->
             callback.request.getNetworkSpecifier() == specifier &&
             callback.score == updatedScore &&
             callback.id == agent.providerId
         }
 
         mCm.unregisterNetworkCallback(cb)
-        provider.expectCallback<OnNetworkRequestWithdrawn>() { callback ->
+        provider.eventuallyExpectCallbackThat<OnNetworkRequestWithdrawn>() { callback ->
             callback.request.getNetworkSpecifier() == specifier &&
             callback.request.hasTransport(TRANSPORT_TEST)
         }
diff --git a/tests/cts/net/Android.bp b/tests/cts/net/Android.bp
index 25596a9..85942b0 100644
--- a/tests/cts/net/Android.bp
+++ b/tests/cts/net/Android.bp
@@ -41,6 +41,7 @@
     srcs: [
         "src/**/*.java",
         "src/**/*.kt",
+        ":ike-aes-xcbc",
     ],
     jarjar_rules: "jarjar-rules-shared.txt",
     static_libs: [
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index defb3ff..acec9a4 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -650,6 +650,18 @@
         mCm.getBackgroundDataSetting();
     }
 
+    private NetworkRequest makeDefaultRequest() {
+        // Make a request that is similar to the way framework tracks the system
+        // default network.
+        return new NetworkRequest.Builder()
+                .clearCapabilities()
+                .addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
+                .addCapability(NetworkCapabilities.NET_CAPABILITY_TRUSTED)
+                .addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_VPN)
+                .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
+                .build();
+    }
+
     private NetworkRequest makeWifiNetworkRequest() {
         return new NetworkRequest.Builder()
                 .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
@@ -688,12 +700,14 @@
 
         final TestNetworkCallback systemDefaultCallback = new TestNetworkCallback();
         final TestNetworkCallback perUidCallback = new TestNetworkCallback();
+        final TestNetworkCallback bestMatchingCallback = new TestNetworkCallback();
         final Handler h = new Handler(Looper.getMainLooper());
         if (TestUtils.shouldTestSApis()) {
             runWithShellPermissionIdentity(() -> {
                 mCmShim.registerSystemDefaultNetworkCallback(systemDefaultCallback, h);
                 mCmShim.registerDefaultNetworkCallbackForUid(Process.myUid(), perUidCallback, h);
             }, NETWORK_SETTINGS);
+            mCm.registerBestMatchingNetworkCallback(makeDefaultRequest(), bestMatchingCallback, h);
         }
 
         Network wifiNetwork = null;
@@ -719,6 +733,10 @@
                 assertNotNull("Did not receive onAvailable on per-UID default network callback",
                         perUidNetwork);
                 assertEquals(defaultNetwork, perUidNetwork);
+                final Network bestMatchingNetwork = bestMatchingCallback.waitForAvailable();
+                assertNotNull("Did not receive onAvailable on best matching network callback",
+                        bestMatchingNetwork);
+                assertEquals(defaultNetwork, bestMatchingNetwork);
             }
 
         } catch (InterruptedException e) {
@@ -729,6 +747,7 @@
             if (TestUtils.shouldTestSApis()) {
                 mCm.unregisterNetworkCallback(systemDefaultCallback);
                 mCm.unregisterNetworkCallback(perUidCallback);
+                mCm.unregisterNetworkCallback(bestMatchingCallback);
             }
         }
     }
@@ -2015,6 +2034,7 @@
         assertThrows(SecurityException.class, () -> mCm.factoryReset());
     }
 
+    @AppModeFull(reason = "Cannot get WifiManager in instant app mode")
     @Test
     public void testFactoryReset() throws Exception {
         assumeTrue(TestUtils.shouldTestSApis());
diff --git a/tests/cts/net/src/android/net/cts/Ikev2VpnTest.java b/tests/cts/net/src/android/net/cts/Ikev2VpnTest.java
index c6d8d65..6e9f0cd 100644
--- a/tests/cts/net/src/android/net/cts/Ikev2VpnTest.java
+++ b/tests/cts/net/src/android/net/cts/Ikev2VpnTest.java
@@ -39,13 +39,11 @@
 import android.net.ConnectivityManager;
 import android.net.Ikev2VpnProfile;
 import android.net.IpSecAlgorithm;
-import android.net.LinkAddress;
 import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.NetworkRequest;
 import android.net.ProxyInfo;
 import android.net.TestNetworkInterface;
-import android.net.TestNetworkManager;
 import android.net.VpnManager;
 import android.net.cts.util.CtsNetUtils;
 import android.os.Build;
@@ -439,41 +437,37 @@
         assertEquals(vpnNetwork, cb.lastLostNetwork);
     }
 
-    private void doTestStartStopVpnProfile(boolean testIpv6) throws Exception {
-        // Non-final; these variables ensure we clean up properly after our test if we have
-        // allocated test network resources
-        final TestNetworkManager tnm = sContext.getSystemService(TestNetworkManager.class);
-        TestNetworkInterface testIface = null;
-        TestNetworkCallback tunNetworkCallback = null;
+    private class VerifyStartStopVpnProfileTest implements TestNetworkRunnable.Test {
+        private final boolean mTestIpv6Only;
 
-        try {
-            // Build underlying test network
-            testIface = tnm.createTunInterface(
-                    new LinkAddress[] {
-                            new LinkAddress(LOCAL_OUTER_4, IP4_PREFIX_LEN),
-                            new LinkAddress(LOCAL_OUTER_6, IP6_PREFIX_LEN)});
+        /**
+         * Constructs the test
+         *
+         * @param testIpv6Only if true, builds a IPv6-only test; otherwise builds a IPv4-only test
+         */
+        VerifyStartStopVpnProfileTest(boolean testIpv6Only) {
+            mTestIpv6Only = testIpv6Only;
+        }
 
-            // Hold on to this callback to ensure network does not get reaped.
-            tunNetworkCallback = mCtsNetUtils.setupAndGetTestNetwork(
-                    testIface.getInterfaceName());
+        @Override
+        public void runTest(TestNetworkInterface testIface, TestNetworkCallback tunNetworkCallback)
+                throws Exception {
             final IkeTunUtils tunUtils = new IkeTunUtils(testIface.getFileDescriptor());
 
-            checkStartStopVpnProfileBuildsNetworks(tunUtils, testIpv6);
-        } finally {
-            // Make sure to stop the VPN profile. This is safe to call multiple times.
+            checkStartStopVpnProfileBuildsNetworks(tunUtils, mTestIpv6Only);
+        }
+
+        @Override
+        public void cleanupTest() {
             sVpnMgr.stopProvisionedVpnProfile();
+        }
 
-            if (testIface != null) {
-                testIface.getFileDescriptor().close();
-            }
-
-            if (tunNetworkCallback != null) {
-                sCM.unregisterNetworkCallback(tunNetworkCallback);
-            }
-
-            final Network testNetwork = tunNetworkCallback.currentNetwork;
-            if (testNetwork != null) {
-                tnm.teardownTestNetwork(testNetwork);
+        @Override
+        public InetAddress[] getTestNetworkAddresses() {
+            if (mTestIpv6Only) {
+                return new InetAddress[] {LOCAL_OUTER_6};
+            } else {
+                return new InetAddress[] {LOCAL_OUTER_4};
             }
         }
     }
@@ -483,9 +477,8 @@
         assumeTrue(mCtsNetUtils.hasIpsecTunnelsFeature());
 
         // Requires shell permission to update appops.
-        runWithShellPermissionIdentity(() -> {
-            doTestStartStopVpnProfile(false);
-        });
+        runWithShellPermissionIdentity(
+                new TestNetworkRunnable(new VerifyStartStopVpnProfileTest(false)));
     }
 
     @Test
@@ -493,9 +486,8 @@
         assumeTrue(mCtsNetUtils.hasIpsecTunnelsFeature());
 
         // Requires shell permission to update appops.
-        runWithShellPermissionIdentity(() -> {
-            doTestStartStopVpnProfile(true);
-        });
+        runWithShellPermissionIdentity(
+                new TestNetworkRunnable(new VerifyStartStopVpnProfileTest(true)));
     }
 
     private static class CertificateAndKey {
diff --git a/tests/cts/net/src/android/net/cts/IpSecAlgorithmImplTest.java b/tests/cts/net/src/android/net/cts/IpSecAlgorithmImplTest.java
new file mode 100644
index 0000000..2f29273
--- /dev/null
+++ b/tests/cts/net/src/android/net/cts/IpSecAlgorithmImplTest.java
@@ -0,0 +1,306 @@
+/*
+ * Copyright (C) 2020 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 android.net.cts;
+
+import static android.net.IpSecAlgorithm.AUTH_AES_CMAC;
+import static android.net.IpSecAlgorithm.AUTH_AES_XCBC;
+import static android.net.IpSecAlgorithm.AUTH_CRYPT_CHACHA20_POLY1305;
+import static android.net.IpSecAlgorithm.CRYPT_AES_CTR;
+import static android.net.cts.PacketUtils.AES_CMAC;
+import static android.net.cts.PacketUtils.AES_CMAC_ICV_LEN;
+import static android.net.cts.PacketUtils.AES_CMAC_KEY_LEN;
+import static android.net.cts.PacketUtils.AES_CTR;
+import static android.net.cts.PacketUtils.AES_CTR_BLK_SIZE;
+import static android.net.cts.PacketUtils.AES_CTR_IV_LEN;
+import static android.net.cts.PacketUtils.AES_CTR_KEY_LEN_20;
+import static android.net.cts.PacketUtils.AES_CTR_KEY_LEN_28;
+import static android.net.cts.PacketUtils.AES_CTR_KEY_LEN_36;
+import static android.net.cts.PacketUtils.AES_CTR_SALT_LEN;
+import static android.net.cts.PacketUtils.AES_XCBC;
+import static android.net.cts.PacketUtils.AES_XCBC_ICV_LEN;
+import static android.net.cts.PacketUtils.AES_XCBC_KEY_LEN;
+import static android.net.cts.PacketUtils.CHACHA20_POLY1305;
+import static android.net.cts.PacketUtils.CHACHA20_POLY1305_BLK_SIZE;
+import static android.net.cts.PacketUtils.CHACHA20_POLY1305_ICV_LEN;
+import static android.net.cts.PacketUtils.CHACHA20_POLY1305_IV_LEN;
+import static android.net.cts.PacketUtils.CHACHA20_POLY1305_KEY_LEN;
+import static android.net.cts.PacketUtils.CHACHA20_POLY1305_SALT_LEN;
+import static android.net.cts.PacketUtils.ESP_HDRLEN;
+import static android.net.cts.PacketUtils.IP6_HDRLEN;
+import static android.net.cts.PacketUtils.getIpHeader;
+import static android.net.cts.util.CtsNetUtils.TestNetworkCallback;
+
+import static com.android.compatibility.common.util.SystemUtil.runWithShellPermissionIdentity;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assume.assumeTrue;
+
+import android.net.IpSecAlgorithm;
+import android.net.IpSecManager;
+import android.net.IpSecTransform;
+import android.net.Network;
+import android.net.TestNetworkInterface;
+import android.net.cts.PacketUtils.BytePayload;
+import android.net.cts.PacketUtils.EspAeadCipher;
+import android.net.cts.PacketUtils.EspAuth;
+import android.net.cts.PacketUtils.EspAuthNull;
+import android.net.cts.PacketUtils.EspCipher;
+import android.net.cts.PacketUtils.EspCipherNull;
+import android.net.cts.PacketUtils.EspCryptCipher;
+import android.net.cts.PacketUtils.EspHeader;
+import android.net.cts.PacketUtils.IpHeader;
+import android.net.cts.PacketUtils.UdpHeader;
+import android.platform.test.annotations.AppModeFull;
+
+import androidx.test.InstrumentationRegistry;
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.net.DatagramPacket;
+import java.net.DatagramSocket;
+import java.net.InetAddress;
+import java.util.Arrays;
+
+@RunWith(AndroidJUnit4.class)
+@AppModeFull(reason = "Socket cannot bind in instant app mode")
+public class IpSecAlgorithmImplTest extends IpSecBaseTest {
+    private static final InetAddress LOCAL_ADDRESS =
+            InetAddress.parseNumericAddress("2001:db8:1::1");
+    private static final InetAddress REMOTE_ADDRESS =
+            InetAddress.parseNumericAddress("2001:db8:1::2");
+
+    private static final int REMOTE_PORT = 12345;
+    private static final IpSecManager sIpSecManager =
+            InstrumentationRegistry.getContext().getSystemService(IpSecManager.class);
+
+    private static class CheckCryptoImplTest implements TestNetworkRunnable.Test {
+        private final IpSecAlgorithm mIpsecEncryptAlgo;
+        private final IpSecAlgorithm mIpsecAuthAlgo;
+        private final IpSecAlgorithm mIpsecAeadAlgo;
+        private final EspCipher mEspCipher;
+        private final EspAuth mEspAuth;
+
+        CheckCryptoImplTest(
+                IpSecAlgorithm ipsecEncryptAlgo,
+                IpSecAlgorithm ipsecAuthAlgo,
+                IpSecAlgorithm ipsecAeadAlgo,
+                EspCipher espCipher,
+                EspAuth espAuth) {
+            mIpsecEncryptAlgo = ipsecEncryptAlgo;
+            mIpsecAuthAlgo = ipsecAuthAlgo;
+            mIpsecAeadAlgo = ipsecAeadAlgo;
+            mEspCipher = espCipher;
+            mEspAuth = espAuth;
+        }
+
+        private static byte[] buildTransportModeEspPayload(
+                int srcPort, int dstPort, int spi, EspCipher espCipher, EspAuth espAuth)
+                throws Exception {
+            final UdpHeader udpPayload =
+                    new UdpHeader(srcPort, dstPort, new BytePayload(TEST_DATA));
+            final IpHeader preEspIpHeader =
+                    getIpHeader(
+                            udpPayload.getProtocolId(), LOCAL_ADDRESS, REMOTE_ADDRESS, udpPayload);
+
+            final PacketUtils.EspHeader espPayload =
+                    new EspHeader(
+                            udpPayload.getProtocolId(),
+                            spi,
+                            1 /* sequence number */,
+                            udpPayload.getPacketBytes(preEspIpHeader),
+                            espCipher,
+                            espAuth);
+            return espPayload.getPacketBytes(preEspIpHeader);
+        }
+
+        @Override
+        public void runTest(TestNetworkInterface testIface, TestNetworkCallback tunNetworkCallback)
+                throws Exception {
+            final TunUtils tunUtils = new TunUtils(testIface.getFileDescriptor());
+            tunNetworkCallback.waitForAvailable();
+            final Network testNetwork = tunNetworkCallback.currentNetwork;
+
+            final IpSecTransform.Builder transformBuilder =
+                    new IpSecTransform.Builder(InstrumentationRegistry.getContext());
+            if (mIpsecAeadAlgo != null) {
+                transformBuilder.setAuthenticatedEncryption(mIpsecAeadAlgo);
+            } else {
+                if (mIpsecEncryptAlgo != null) {
+                    transformBuilder.setEncryption(mIpsecEncryptAlgo);
+                }
+                if (mIpsecAuthAlgo != null) {
+                    transformBuilder.setAuthentication(mIpsecAuthAlgo);
+                }
+            }
+
+            try (IpSecManager.SecurityParameterIndex outSpi =
+                            sIpSecManager.allocateSecurityParameterIndex(REMOTE_ADDRESS);
+                    IpSecManager.SecurityParameterIndex inSpi =
+                            sIpSecManager.allocateSecurityParameterIndex(LOCAL_ADDRESS);
+                    IpSecTransform outTransform =
+                            transformBuilder.buildTransportModeTransform(LOCAL_ADDRESS, outSpi);
+                    IpSecTransform inTransform =
+                            transformBuilder.buildTransportModeTransform(REMOTE_ADDRESS, inSpi);
+                    // Bind localSocket to a random available port.
+                    DatagramSocket localSocket = new DatagramSocket(0)) {
+                sIpSecManager.applyTransportModeTransform(
+                        localSocket, IpSecManager.DIRECTION_IN, inTransform);
+                sIpSecManager.applyTransportModeTransform(
+                        localSocket, IpSecManager.DIRECTION_OUT, outTransform);
+
+                // Send ESP packet
+                final DatagramPacket outPacket =
+                        new DatagramPacket(
+                                TEST_DATA, 0, TEST_DATA.length, REMOTE_ADDRESS, REMOTE_PORT);
+                testNetwork.bindSocket(localSocket);
+                localSocket.send(outPacket);
+                final byte[] outEspPacket =
+                        tunUtils.awaitEspPacket(outSpi.getSpi(), false /* useEncap */);
+
+                // Remove transform for good hygiene
+                sIpSecManager.removeTransportModeTransforms(localSocket);
+
+                // Get the kernel-generated ESP payload
+                final byte[] outEspPayload = new byte[outEspPacket.length - IP6_HDRLEN];
+                System.arraycopy(outEspPacket, IP6_HDRLEN, outEspPayload, 0, outEspPayload.length);
+
+                // Get the IV of the kernel-generated ESP payload
+                final byte[] iv =
+                        Arrays.copyOfRange(
+                                outEspPayload, ESP_HDRLEN, ESP_HDRLEN + mEspCipher.ivLen);
+
+                // Build ESP payload using the kernel-generated IV and the user space crypto
+                // implementations
+                mEspCipher.updateIv(iv);
+                final byte[] expectedEspPayload =
+                        buildTransportModeEspPayload(
+                                localSocket.getLocalPort(),
+                                REMOTE_PORT,
+                                outSpi.getSpi(),
+                                mEspCipher,
+                                mEspAuth);
+
+                // Compare user-space-generated and kernel-generated ESP payload
+                assertArrayEquals(expectedEspPayload, outEspPayload);
+            }
+        }
+
+        @Override
+        public void cleanupTest() {
+            // Do nothing
+        }
+
+        @Override
+        public InetAddress[] getTestNetworkAddresses() {
+            return new InetAddress[] {LOCAL_ADDRESS};
+        }
+    }
+
+    private void checkAesCtr(int keyLen) throws Exception {
+        final byte[] cryptKey = getKeyBytes(keyLen);
+
+        final IpSecAlgorithm ipsecEncryptAlgo =
+                new IpSecAlgorithm(IpSecAlgorithm.CRYPT_AES_CTR, cryptKey);
+        final EspCipher espCipher =
+                new EspCryptCipher(
+                        AES_CTR, AES_CTR_BLK_SIZE, cryptKey, AES_CTR_IV_LEN, AES_CTR_SALT_LEN);
+
+        runWithShellPermissionIdentity(new TestNetworkRunnable(new CheckCryptoImplTest(
+                ipsecEncryptAlgo, null /* ipsecAuthAlgo */, null /* ipsecAeadAlgo */,
+                espCipher, EspAuthNull.getInstance())));
+    }
+
+    @Test
+    public void testAesCtr160() throws Exception {
+        assumeTrue(hasIpSecAlgorithm(CRYPT_AES_CTR));
+
+        checkAesCtr(AES_CTR_KEY_LEN_20);
+    }
+
+    @Test
+    public void testAesCtr224() throws Exception {
+        assumeTrue(hasIpSecAlgorithm(CRYPT_AES_CTR));
+
+        checkAesCtr(AES_CTR_KEY_LEN_28);
+    }
+
+    @Test
+    public void testAesCtr288() throws Exception {
+        assumeTrue(hasIpSecAlgorithm(CRYPT_AES_CTR));
+
+        checkAesCtr(AES_CTR_KEY_LEN_36);
+    }
+
+    @Test
+    public void testAesXcbc() throws Exception {
+        assumeTrue(hasIpSecAlgorithm(AUTH_AES_XCBC));
+
+        final byte[] authKey = getKeyBytes(AES_XCBC_KEY_LEN);
+        final IpSecAlgorithm ipsecAuthAlgo =
+                new IpSecAlgorithm(IpSecAlgorithm.AUTH_AES_XCBC, authKey, AES_XCBC_ICV_LEN * 8);
+        final EspAuth espAuth = new EspAuth(AES_XCBC, authKey, AES_XCBC_ICV_LEN);
+
+        runWithShellPermissionIdentity(new TestNetworkRunnable(new CheckCryptoImplTest(
+                null /* ipsecEncryptAlgo */, ipsecAuthAlgo, null /* ipsecAeadAlgo */,
+                EspCipherNull.getInstance(), espAuth)));
+    }
+
+    @Test
+    public void testAesCmac() throws Exception {
+        assumeTrue(hasIpSecAlgorithm(AUTH_AES_CMAC));
+
+        final byte[] authKey = getKeyBytes(AES_CMAC_KEY_LEN);
+        final IpSecAlgorithm ipsecAuthAlgo =
+                new IpSecAlgorithm(IpSecAlgorithm.AUTH_AES_CMAC, authKey, AES_CMAC_ICV_LEN * 8);
+        final EspAuth espAuth = new EspAuth(AES_CMAC, authKey, AES_CMAC_ICV_LEN);
+
+        runWithShellPermissionIdentity(new TestNetworkRunnable(new CheckCryptoImplTest(
+                null /* ipsecEncryptAlgo */, ipsecAuthAlgo, null /* ipsecAeadAlgo */,
+                EspCipherNull.getInstance(), espAuth)));
+    }
+
+    @Test
+    public void testChaCha20Poly1305() throws Exception {
+        assumeTrue(hasIpSecAlgorithm(AUTH_CRYPT_CHACHA20_POLY1305));
+
+        final byte[] cryptKey = getKeyBytes(CHACHA20_POLY1305_KEY_LEN);
+        final IpSecAlgorithm ipsecAeadAlgo =
+                new IpSecAlgorithm(
+                        IpSecAlgorithm.AUTH_CRYPT_CHACHA20_POLY1305,
+                        cryptKey,
+                        CHACHA20_POLY1305_ICV_LEN * 8);
+        final EspAeadCipher espAead =
+                new EspAeadCipher(
+                        CHACHA20_POLY1305,
+                        CHACHA20_POLY1305_BLK_SIZE,
+                        cryptKey,
+                        CHACHA20_POLY1305_IV_LEN,
+                        CHACHA20_POLY1305_ICV_LEN,
+                        CHACHA20_POLY1305_SALT_LEN);
+
+        runWithShellPermissionIdentity(
+                new TestNetworkRunnable(
+                        new CheckCryptoImplTest(
+                                null /* ipsecEncryptAlgo */,
+                                null /* ipsecAuthAlgo */,
+                                ipsecAeadAlgo,
+                                espAead,
+                                EspAuthNull.getInstance())));
+    }
+}
diff --git a/tests/cts/net/src/android/net/cts/IpSecManagerTest.java b/tests/cts/net/src/android/net/cts/IpSecManagerTest.java
index e7e1d67..5c95aa3 100644
--- a/tests/cts/net/src/android/net/cts/IpSecManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/IpSecManagerTest.java
@@ -33,7 +33,7 @@
 import static android.net.cts.PacketUtils.AES_CMAC_KEY_LEN;
 import static android.net.cts.PacketUtils.AES_CTR_BLK_SIZE;
 import static android.net.cts.PacketUtils.AES_CTR_IV_LEN;
-import static android.net.cts.PacketUtils.AES_CTR_KEY_LEN;
+import static android.net.cts.PacketUtils.AES_CTR_KEY_LEN_20;
 import static android.net.cts.PacketUtils.AES_GCM_BLK_SIZE;
 import static android.net.cts.PacketUtils.AES_GCM_IV_LEN;
 import static android.net.cts.PacketUtils.AES_XCBC_ICV_LEN;
@@ -74,7 +74,6 @@
 
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
-import com.android.testutils.SkipPresubmit;
 
 import org.junit.Rule;
 import org.junit.Test;
@@ -771,7 +770,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testAesCbcHmacMd5Tcp6() throws Exception {
         IpSecAlgorithm crypt = new IpSecAlgorithm(IpSecAlgorithm.CRYPT_AES_CBC, CRYPT_KEY);
         IpSecAlgorithm auth = new IpSecAlgorithm(IpSecAlgorithm.AUTH_HMAC_MD5, getKey(128), 96);
@@ -804,7 +802,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testAesCbcHmacSha1Tcp6() throws Exception {
         IpSecAlgorithm crypt = new IpSecAlgorithm(IpSecAlgorithm.CRYPT_AES_CBC, CRYPT_KEY);
         IpSecAlgorithm auth = new IpSecAlgorithm(IpSecAlgorithm.AUTH_HMAC_SHA1, getKey(160), 96);
@@ -837,7 +834,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testAesCbcHmacSha256Tcp6() throws Exception {
         IpSecAlgorithm crypt = new IpSecAlgorithm(IpSecAlgorithm.CRYPT_AES_CBC, CRYPT_KEY);
         IpSecAlgorithm auth = new IpSecAlgorithm(IpSecAlgorithm.AUTH_HMAC_SHA256, getKey(256), 128);
@@ -870,7 +866,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testAesCbcHmacSha384Tcp6() throws Exception {
         IpSecAlgorithm crypt = new IpSecAlgorithm(IpSecAlgorithm.CRYPT_AES_CBC, CRYPT_KEY);
         IpSecAlgorithm auth = new IpSecAlgorithm(IpSecAlgorithm.AUTH_HMAC_SHA384, getKey(384), 192);
@@ -903,7 +898,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testAesCbcHmacSha512Tcp6() throws Exception {
         IpSecAlgorithm crypt = new IpSecAlgorithm(IpSecAlgorithm.CRYPT_AES_CBC, CRYPT_KEY);
         IpSecAlgorithm auth = new IpSecAlgorithm(IpSecAlgorithm.AUTH_HMAC_SHA512, getKey(512), 256);
@@ -928,7 +922,7 @@
     }
 
     private static IpSecAlgorithm buildCryptAesCtr() throws Exception {
-        return new IpSecAlgorithm(CRYPT_AES_CTR, getKeyBytes(AES_CTR_KEY_LEN));
+        return new IpSecAlgorithm(CRYPT_AES_CTR, getKeyBytes(AES_CTR_KEY_LEN_20));
     }
 
     private static IpSecAlgorithm buildAuthHmacSha512() throws Exception {
@@ -947,7 +941,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testAesCtrHmacSha512Tcp6() throws Exception {
         assumeTrue(hasIpSecAlgorithm(CRYPT_AES_CTR));
 
@@ -1002,7 +995,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testAesCbcAesXCbcTcp6() throws Exception {
         assumeTrue(hasIpSecAlgorithm(AUTH_AES_XCBC));
 
@@ -1043,7 +1035,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testAesCbcAesCmacTcp6() throws Exception {
         assumeTrue(hasIpSecAlgorithm(AUTH_AES_CMAC));
 
@@ -1082,7 +1073,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testAesGcm64Tcp6() throws Exception {
         IpSecAlgorithm authCrypt =
                 new IpSecAlgorithm(IpSecAlgorithm.AUTH_CRYPT_AES_GCM, AEAD_KEY, 64);
@@ -1115,7 +1105,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testAesGcm96Tcp6() throws Exception {
         IpSecAlgorithm authCrypt =
                 new IpSecAlgorithm(IpSecAlgorithm.AUTH_CRYPT_AES_GCM, AEAD_KEY, 96);
@@ -1148,7 +1137,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testAesGcm128Tcp6() throws Exception {
         IpSecAlgorithm authCrypt =
                 new IpSecAlgorithm(IpSecAlgorithm.AUTH_CRYPT_AES_GCM, AEAD_KEY, 128);
@@ -1187,7 +1175,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testChaCha20Poly1305Tcp6() throws Exception {
         assumeTrue(hasIpSecAlgorithm(AUTH_CRYPT_CHACHA20_POLY1305));
 
@@ -1463,7 +1450,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testCryptTcp6() throws Exception {
         IpSecAlgorithm crypt = new IpSecAlgorithm(IpSecAlgorithm.CRYPT_AES_CBC, CRYPT_KEY);
         checkTransform(IPPROTO_TCP, IPV6_LOOPBACK, crypt, null, null, false, 1, false);
@@ -1471,7 +1457,6 @@
     }
 
     @Test
-    @SkipPresubmit(reason = "b/186608065 - kernel 5.10 regression in TrafficStats with ipsec")
     public void testAuthTcp6() throws Exception {
         IpSecAlgorithm auth = new IpSecAlgorithm(IpSecAlgorithm.AUTH_HMAC_SHA256, getKey(256), 128);
         checkTransform(IPPROTO_TCP, IPV6_LOOPBACK, null, auth, null, false, 1, false);
diff --git a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
index 8e2b310..c505cef 100644
--- a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
@@ -66,10 +66,12 @@
 import android.net.cts.NetworkAgentTest.TestableNetworkAgent.CallbackEntry.OnStopSocketKeepalive
 import android.net.cts.NetworkAgentTest.TestableNetworkAgent.CallbackEntry.OnValidationStatus
 import android.os.Build
+import android.os.Handler
 import android.os.HandlerThread
 import android.os.Looper
 import android.os.Message
 import android.os.SystemClock
+import android.telephony.TelephonyManager
 import android.util.DebugUtils.valueToString
 import androidx.test.InstrumentationRegistry
 import com.android.compatibility.common.util.SystemUtil.runWithShellPermissionIdentity
@@ -268,10 +270,9 @@
             history.add(OnSignalStrengthThresholdsUpdated(thresholds))
         }
 
-        fun expectEmptySignalStrengths() {
+        fun expectSignalStrengths(thresholds: IntArray? = intArrayOf()) {
             expectCallback<OnSignalStrengthThresholdsUpdated>().let {
-                // intArrayOf() without arguments makes an empty array
-                assertArrayEquals(intArrayOf(), it.thresholds)
+                assertArrayEquals(thresholds, it.thresholds)
             }
         }
 
@@ -291,7 +292,7 @@
         // a NetworkAgent whose network does not require validation (which test networks do
         // not, since they lack the INTERNET capability). It always contains the default argument
         // for the URI.
-        fun expectNoInternetValidationStatus() = expectCallback<OnValidationStatus>().let {
+        fun expectValidationBypassedStatus() = expectCallback<OnValidationStatus>().let {
             assertEquals(it.status, VALID_NETWORK)
             // The returned Uri is parsed from the empty string, which means it's an
             // instance of the (private) Uri.StringUri. There are no real good ways
@@ -331,9 +332,30 @@
         callbacksToCleanUp.add(callback)
     }
 
+    private fun registerBestMatchingNetworkCallback(
+        request: NetworkRequest,
+        callback: TestableNetworkCallback,
+        handler: Handler
+    ) {
+        mCM!!.registerBestMatchingNetworkCallback(request, callback, handler)
+        callbacksToCleanUp.add(callback)
+    }
+
+    private fun makeTestNetworkRequest(specifier: String? = null): NetworkRequest {
+        return NetworkRequest.Builder()
+                .clearCapabilities()
+                .addTransportType(TRANSPORT_TEST)
+                .also {
+                    if (specifier != null) {
+                        it.setNetworkSpecifier(CompatUtil.makeEthernetNetworkSpecifier(specifier))
+                    }
+                }
+                .build()
+    }
+
     private fun createNetworkAgent(
         context: Context = realContext,
-        name: String? = null,
+        specifier: String? = null,
         initialNc: NetworkCapabilities? = null,
         initialLp: LinkProperties? = null,
         initialConfig: NetworkAgentConfig? = null
@@ -348,8 +370,8 @@
             if (SdkLevel.isAtLeastS()) {
                 addCapability(NET_CAPABILITY_NOT_VCN_MANAGED)
             }
-            if (null != name) {
-                setNetworkSpecifier(CompatUtil.makeEthernetNetworkSpecifier(name))
+            if (null != specifier) {
+                setNetworkSpecifier(CompatUtil.makeEthernetNetworkSpecifier(specifier))
             }
         }
         val lp = initialLp ?: LinkProperties().apply {
@@ -362,19 +384,24 @@
         }
     }
 
-    private fun createConnectedNetworkAgent(context: Context = realContext, name: String? = null):
-            Pair<TestableNetworkAgent, TestableNetworkCallback> {
-        val request: NetworkRequest = NetworkRequest.Builder()
-                .clearCapabilities()
-                .addTransportType(TRANSPORT_TEST)
-                .build()
+    private fun createConnectedNetworkAgent(
+        context: Context = realContext,
+        specifier: String? = UUID.randomUUID().toString(),
+        initialConfig: NetworkAgentConfig? = null,
+        expectedInitSignalStrengthThresholds: IntArray? = intArrayOf()
+    ): Pair<TestableNetworkAgent, TestableNetworkCallback> {
         val callback = TestableNetworkCallback(timeoutMs = DEFAULT_TIMEOUT_MS)
-        requestNetwork(request, callback)
-        val agent = createNetworkAgent(context, name)
+        // Ensure this NetworkAgent is never unneeded by filing a request with its specifier.
+        requestNetwork(makeTestNetworkRequest(specifier = specifier), callback)
+        val agent = createNetworkAgent(context, specifier, initialConfig = initialConfig)
         agent.setTeardownDelayMillis(0)
+        // Connect the agent and verify initial status callbacks.
         agent.register()
         agent.markConnected()
         agent.expectCallback<OnNetworkCreated>()
+        agent.expectSignalStrengths(expectedInitSignalStrengthThresholds)
+        agent.expectValidationBypassedStatus()
+        callback.expectAvailableThenValidatedCallbacks(agent.network!!)
         return agent to callback
     }
 
@@ -383,14 +410,52 @@
     }
 
     @Test
+    fun testSetSubtypeNameAndExtraInfoByAgentConfig() {
+        val subtypeLTE = TelephonyManager.NETWORK_TYPE_LTE
+        val subtypeNameLTE = "LTE"
+        val legacyExtraInfo = "mylegacyExtraInfo"
+        val config = NetworkAgentConfig.Builder()
+                .setLegacySubType(subtypeLTE)
+                .setLegacySubTypeName(subtypeNameLTE)
+                .setLegacyExtraInfo(legacyExtraInfo).build()
+        val (agent, callback) = createConnectedNetworkAgent(initialConfig = config)
+        val networkInfo = mCM.getNetworkInfo(agent.network)
+        assertEquals(subtypeLTE, networkInfo.getSubtype())
+        assertEquals(subtypeNameLTE, networkInfo.getSubtypeName())
+        assertEquals(legacyExtraInfo, config.getLegacyExtraInfo())
+    }
+
+    @Test
+    fun testSetLegacySubtypeInNetworkAgent() {
+        val subtypeLTE = TelephonyManager.NETWORK_TYPE_LTE
+        val subtypeUMTS = TelephonyManager.NETWORK_TYPE_UMTS
+        val subtypeNameLTE = "LTE"
+        val subtypeNameUMTS = "UMTS"
+        val config = NetworkAgentConfig.Builder()
+                .setLegacySubType(subtypeLTE)
+                .setLegacySubTypeName(subtypeNameLTE).build()
+        val (agent, callback) = createConnectedNetworkAgent(initialConfig = config)
+            agent.setLegacySubtype(subtypeUMTS, subtypeNameUMTS)
+
+            // There is no callback when networkInfo changes,
+            // so use the NetworkCapabilities callback to ensure
+            // that networkInfo is ready for verification.
+            val nc = NetworkCapabilities(agent.nc)
+            nc.addCapability(NET_CAPABILITY_NOT_METERED)
+            agent.sendNetworkCapabilities(nc)
+            callback.expectCapabilitiesThat(agent.network) {
+                it.hasCapability(NET_CAPABILITY_NOT_METERED)
+            }
+            val networkInfo = mCM.getNetworkInfo(agent.network)
+            assertEquals(subtypeUMTS, networkInfo.getSubtype())
+            assertEquals(subtypeNameUMTS, networkInfo.getSubtypeName())
+    }
+
+    @Test
     fun testConnectAndUnregister() {
         val (agent, callback) = createConnectedNetworkAgent()
-        callback.expectAvailableThenValidatedCallbacks(agent.network)
-        agent.expectEmptySignalStrengths()
-        agent.expectNoInternetValidationStatus()
-
         unregister(agent)
-        callback.expectCallback<Lost>(agent.network)
+        callback.expectCallback<Lost>(agent.network!!)
         assertFailsWith<IllegalStateException>("Must not be able to register an agent twice") {
             agent.register()
         }
@@ -398,11 +463,8 @@
 
     @Test
     fun testOnBandwidthUpdateRequested() {
-        val (agent, callback) = createConnectedNetworkAgent()
-        callback.expectAvailableThenValidatedCallbacks(agent.network)
-        agent.expectEmptySignalStrengths()
-        agent.expectNoInternetValidationStatus()
-        mCM.requestBandwidthUpdate(agent.network)
+        val (agent, _) = createConnectedNetworkAgent()
+        mCM.requestBandwidthUpdate(agent.network!!)
         agent.expectCallback<OnBandwidthUpdateRequested>()
         unregister(agent)
     }
@@ -420,13 +482,8 @@
                 registerNetworkCallback(request, it)
             }
         }
-        createConnectedNetworkAgent().let { (agent, callback) ->
-            callback.expectAvailableThenValidatedCallbacks(agent.network)
-            agent.expectCallback<OnSignalStrengthThresholdsUpdated>().let {
-                assertArrayEquals(it.thresholds, thresholds)
-            }
-            agent.expectNoInternetValidationStatus()
-
+        createConnectedNetworkAgent(expectedInitSignalStrengthThresholds = thresholds).let {
+            (agent, callback) ->
             // Send signal strength and check that the callbacks are called appropriately.
             val nc = NetworkCapabilities(agent.nc)
             nc.setSignalStrength(20)
@@ -435,21 +492,21 @@
 
             nc.setSignalStrength(40)
             agent.sendNetworkCapabilities(nc)
-            callbacks[0].expectAvailableCallbacks(agent.network)
+            callbacks[0].expectAvailableCallbacks(agent.network!!)
             callbacks[1].assertNoCallback(NO_CALLBACK_TIMEOUT)
             callbacks[2].assertNoCallback(NO_CALLBACK_TIMEOUT)
 
             nc.setSignalStrength(80)
             agent.sendNetworkCapabilities(nc)
-            callbacks[0].expectCapabilitiesThat(agent.network) { it.signalStrength == 80 }
-            callbacks[1].expectAvailableCallbacks(agent.network)
-            callbacks[2].expectAvailableCallbacks(agent.network)
+            callbacks[0].expectCapabilitiesThat(agent.network!!) { it.signalStrength == 80 }
+            callbacks[1].expectAvailableCallbacks(agent.network!!)
+            callbacks[2].expectAvailableCallbacks(agent.network!!)
 
             nc.setSignalStrength(55)
             agent.sendNetworkCapabilities(nc)
-            callbacks[0].expectCapabilitiesThat(agent.network) { it.signalStrength == 55 }
-            callbacks[1].expectCapabilitiesThat(agent.network) { it.signalStrength == 55 }
-            callbacks[2].expectCallback<Lost>(agent.network)
+            callbacks[0].expectCapabilitiesThat(agent.network!!) { it.signalStrength == 55 }
+            callbacks[1].expectCapabilitiesThat(agent.network!!) { it.signalStrength == 55 }
+            callbacks[2].expectCallback<Lost>(agent.network!!)
         }
         callbacks.forEach {
             mCM.unregisterNetworkCallback(it)
@@ -498,20 +555,17 @@
 
     @Test
     fun testSendUpdates(): Unit = createConnectedNetworkAgent().let { (agent, callback) ->
-        callback.expectAvailableThenValidatedCallbacks(agent.network)
-        agent.expectEmptySignalStrengths()
-        agent.expectNoInternetValidationStatus()
         val ifaceName = "adhocIface"
         val lp = LinkProperties(agent.lp)
         lp.setInterfaceName(ifaceName)
         agent.sendLinkProperties(lp)
-        callback.expectLinkPropertiesThat(agent.network) {
+        callback.expectLinkPropertiesThat(agent.network!!) {
             it.getInterfaceName() == ifaceName
         }
         val nc = NetworkCapabilities(agent.nc)
         nc.addCapability(NET_CAPABILITY_NOT_METERED)
         agent.sendNetworkCapabilities(nc)
-        callback.expectCapabilitiesThat(agent.network) {
+        callback.expectCapabilitiesThat(agent.network!!) {
             it.hasCapability(NET_CAPABILITY_NOT_METERED)
         }
     }
@@ -520,56 +574,32 @@
     fun testSendScore() {
         // This test will create two networks and check that the one with the stronger
         // score wins out for a request that matches them both.
-        // First create requests to make sure both networks are kept up, using the
-        // specifier so they are specific to each network
-        val name1 = UUID.randomUUID().toString()
-        val name2 = UUID.randomUUID().toString()
-        val request1 = NetworkRequest.Builder()
-                .clearCapabilities()
-                .addTransportType(TRANSPORT_TEST)
-                .setNetworkSpecifier(CompatUtil.makeEthernetNetworkSpecifier(name1))
-                .build()
-        val request2 = NetworkRequest.Builder()
-                .clearCapabilities()
-                .addTransportType(TRANSPORT_TEST)
-                .setNetworkSpecifier(CompatUtil.makeEthernetNetworkSpecifier(name2))
-                .build()
-        val callback1 = TestableNetworkCallback(timeoutMs = DEFAULT_TIMEOUT_MS)
-        val callback2 = TestableNetworkCallback(timeoutMs = DEFAULT_TIMEOUT_MS)
-        requestNetwork(request1, callback1)
-        requestNetwork(request2, callback2)
 
-        // Then file the interesting request
-        val request = NetworkRequest.Builder()
-                .clearCapabilities()
-                .addTransportType(TRANSPORT_TEST)
-                .build()
+        // File the interesting request
         val callback = TestableNetworkCallback(timeoutMs = DEFAULT_TIMEOUT_MS)
-        requestNetwork(request, callback)
+        requestNetwork(makeTestNetworkRequest(), callback)
 
-        // Connect the first Network
-        createConnectedNetworkAgent(name = name1).let { (agent1, _) ->
-            callback.expectAvailableThenValidatedCallbacks(agent1.network)
-            // If using the int ranking, agent1 must be upgraded to a better score so that there is
-            // no ambiguity when agent2 connects that agent1 is still better. If using policy
-            // ranking, this is not necessary.
-            agent1.sendNetworkScore(NetworkScore.Builder().setLegacyInt(BETTER_NETWORK_SCORE)
-                    .build())
-            // Connect the second agent
-            createConnectedNetworkAgent(name = name2).let { (agent2, _) ->
-                agent2.markConnected()
-                // The callback should not see anything yet. With int ranking, agent1 was upgraded
-                // to a stronger score beforehand. With policy ranking, agent1 is preferred by
-                // virtue of already satisfying the request.
-                callback.assertNoCallback(NO_CALLBACK_TIMEOUT)
-                // Now downgrade the score and expect the callback now prefers agent2
-                agent1.sendNetworkScore(NetworkScore.Builder()
-                        .setLegacyInt(WORSE_NETWORK_SCORE)
-                        .setExiting(true)
-                        .build())
-                callback.expectCallback<Available>(agent2.network)
-            }
-        }
+        // Connect the first Network, with an unused callback that kept the network up.
+        val (agent1, _) = createConnectedNetworkAgent()
+        callback.expectAvailableThenValidatedCallbacks(agent1.network!!)
+        // If using the int ranking, agent1 must be upgraded to a better score so that there is
+        // no ambiguity when agent2 connects that agent1 is still better. If using policy
+        // ranking, this is not necessary.
+        agent1.sendNetworkScore(NetworkScore.Builder().setLegacyInt(BETTER_NETWORK_SCORE)
+                .build())
+
+        // Connect the second agent.
+        val (agent2, _) = createConnectedNetworkAgent()
+        // The callback should not see anything yet. With int ranking, agent1 was upgraded
+        // to a stronger score beforehand. With policy ranking, agent1 is preferred by
+        // virtue of already satisfying the request.
+        callback.assertNoCallback(NO_CALLBACK_TIMEOUT)
+        // Now downgrade the score and expect the callback now prefers agent2
+        agent1.sendNetworkScore(NetworkScore.Builder()
+                .setLegacyInt(WORSE_NETWORK_SCORE)
+                .setExiting(true)
+                .build())
+        callback.expectCallback<Available>(agent2.network!!)
 
         // tearDown() will unregister the requests and agents
     }
@@ -610,7 +640,7 @@
         callback.expectAvailableThenValidatedCallbacks(agent.network!!)
 
         // Check that the default network's transport is propagated to the VPN.
-        var vpnNc = mCM.getNetworkCapabilities(agent.network)
+        var vpnNc = mCM.getNetworkCapabilities(agent.network!!)
         assertNotNull(vpnNc)
         assertEquals(VpnManager.TYPE_VPN_SERVICE,
                 (vpnNc.transportInfo as VpnTransportInfo).type)
@@ -642,7 +672,7 @@
         // This is not very accurate because the test does not control the capabilities of the
         // underlying networks, and because not congested, not roaming, and not suspended are the
         // default anyway. It's still useful as an extra check though.
-        vpnNc = mCM.getNetworkCapabilities(agent.network)
+        vpnNc = mCM.getNetworkCapabilities(agent.network!!)
         for (cap in listOf(NET_CAPABILITY_NOT_CONGESTED,
                 NET_CAPABILITY_NOT_ROAMING,
                 NET_CAPABILITY_NOT_SUSPENDED)) {
@@ -653,7 +683,7 @@
         }
 
         unregister(agent)
-        callback.expectCallback<Lost>(agent.network)
+        callback.expectCallback<Lost>(agent.network!!)
     }
 
     private fun unregister(agent: TestableNetworkAgent) {
@@ -741,43 +771,24 @@
     fun testTemporarilyUnmeteredCapability() {
         // This test will create a networks with/without NET_CAPABILITY_TEMPORARILY_NOT_METERED
         // and check that the callback reflects the capability changes.
-        // First create a request to make sure the network is kept up
-        val request1 = NetworkRequest.Builder()
-                .clearCapabilities()
-                .addTransportType(TRANSPORT_TEST)
-                .build()
-        val callback1 = TestableNetworkCallback(timeoutMs = DEFAULT_TIMEOUT_MS).also {
-            registerNetworkCallback(request1, it)
-        }
-        requestNetwork(request1, callback1)
-
-        // Then file the interesting request
-        val request = NetworkRequest.Builder()
-                .clearCapabilities()
-                .addTransportType(TRANSPORT_TEST)
-                .build()
-        val callback = TestableNetworkCallback(timeoutMs = DEFAULT_TIMEOUT_MS)
-        requestNetwork(request, callback)
 
         // Connect the network
-        createConnectedNetworkAgent().let { (agent, _) ->
-            callback.expectAvailableThenValidatedCallbacks(agent.network)
+        val (agent, callback) = createConnectedNetworkAgent()
 
-            // Send TEMP_NOT_METERED and check that the callback is called appropriately.
-            val nc1 = NetworkCapabilities(agent.nc)
-                    .addCapability(NET_CAPABILITY_TEMPORARILY_NOT_METERED)
-            agent.sendNetworkCapabilities(nc1)
-            callback.expectCapabilitiesThat(agent.network) {
-                it.hasCapability(NET_CAPABILITY_TEMPORARILY_NOT_METERED)
-            }
+        // Send TEMP_NOT_METERED and check that the callback is called appropriately.
+        val nc1 = NetworkCapabilities(agent.nc)
+                .addCapability(NET_CAPABILITY_TEMPORARILY_NOT_METERED)
+        agent.sendNetworkCapabilities(nc1)
+        callback.expectCapabilitiesThat(agent.network!!) {
+            it.hasCapability(NET_CAPABILITY_TEMPORARILY_NOT_METERED)
+        }
 
-            // Remove TEMP_NOT_METERED and check that the callback is called appropriately.
-            val nc2 = NetworkCapabilities(agent.nc)
-                    .removeCapability(NET_CAPABILITY_TEMPORARILY_NOT_METERED)
-            agent.sendNetworkCapabilities(nc2)
-            callback.expectCapabilitiesThat(agent.network) {
-                !it.hasCapability(NET_CAPABILITY_TEMPORARILY_NOT_METERED)
-            }
+        // Remove TEMP_NOT_METERED and check that the callback is called appropriately.
+        val nc2 = NetworkCapabilities(agent.nc)
+                .removeCapability(NET_CAPABILITY_TEMPORARILY_NOT_METERED)
+        agent.sendNetworkCapabilities(nc2)
+        callback.expectCapabilitiesThat(agent.network!!) {
+            !it.hasCapability(NET_CAPABILITY_TEMPORARILY_NOT_METERED)
         }
 
         // tearDown() will unregister the requests and agents
@@ -790,88 +801,137 @@
         // score wins out for a request that matches them both. And the weaker agent will
         // be disconnected after customized linger duration.
 
-        // Connect the first Network
-        val name1 = UUID.randomUUID().toString()
-        val name2 = UUID.randomUUID().toString()
-        val (agent1, callback) = createConnectedNetworkAgent(name = name1)
-        callback.expectAvailableThenValidatedCallbacks(agent1.network!!)
-        // Downgrade agent1 to a worse score so that there is no ambiguity when
-        // agent2 connects.
-        agent1.sendNetworkScore(NetworkScore.Builder().setLegacyInt(WORSE_NETWORK_SCORE)
+        // Request the first Network, with a request that could moved to agentStronger in order to
+        // make agentWeaker linger later.
+        val specifierWeaker = UUID.randomUUID().toString()
+        val specifierStronger = UUID.randomUUID().toString()
+        val commonCallback = TestableNetworkCallback(timeoutMs = DEFAULT_TIMEOUT_MS)
+        requestNetwork(makeTestNetworkRequest(), commonCallback)
+        val agentWeaker = createNetworkAgent(specifier = specifierWeaker)
+        agentWeaker.register()
+        agentWeaker.markConnected()
+        commonCallback.expectAvailableThenValidatedCallbacks(agentWeaker.network!!)
+        // Downgrade agentWeaker to a worse score so that there is no ambiguity when
+        // agentStronger connects.
+        agentWeaker.sendNetworkScore(NetworkScore.Builder().setLegacyInt(WORSE_NETWORK_SCORE)
                 .setExiting(true).build())
 
         // Verify invalid linger duration cannot be set.
         assertFailsWith<IllegalArgumentException> {
-            agent1.setLingerDuration(Duration.ofMillis(-1))
+            agentWeaker.setLingerDuration(Duration.ofMillis(-1))
         }
-        assertFailsWith<IllegalArgumentException> { agent1.setLingerDuration(Duration.ZERO) }
+        assertFailsWith<IllegalArgumentException> { agentWeaker.setLingerDuration(Duration.ZERO) }
         assertFailsWith<IllegalArgumentException> {
-            agent1.setLingerDuration(Duration.ofMillis(Integer.MIN_VALUE.toLong()))
+            agentWeaker.setLingerDuration(Duration.ofMillis(Integer.MIN_VALUE.toLong()))
         }
         assertFailsWith<IllegalArgumentException> {
-            agent1.setLingerDuration(Duration.ofMillis(Integer.MAX_VALUE.toLong() + 1))
+            agentWeaker.setLingerDuration(Duration.ofMillis(Integer.MAX_VALUE.toLong() + 1))
         }
         assertFailsWith<IllegalArgumentException> {
-            agent1.setLingerDuration(Duration.ofMillis(
+            agentWeaker.setLingerDuration(Duration.ofMillis(
                     NetworkAgent.MIN_LINGER_TIMER_MS.toLong() - 1))
         }
         // Verify valid linger timer can be set, but it should not take effect since the network
         // is still needed.
-        agent1.setLingerDuration(Duration.ofMillis(Integer.MAX_VALUE.toLong()))
-        callback.assertNoCallback(NO_CALLBACK_TIMEOUT)
+        agentWeaker.setLingerDuration(Duration.ofMillis(Integer.MAX_VALUE.toLong()))
+        commonCallback.assertNoCallback(NO_CALLBACK_TIMEOUT)
         // Set to the value we want to verify the functionality.
-        agent1.setLingerDuration(Duration.ofMillis(NetworkAgent.MIN_LINGER_TIMER_MS.toLong()))
-        // Make a listener which can observe agent1 lost later.
+        agentWeaker.setLingerDuration(Duration.ofMillis(NetworkAgent.MIN_LINGER_TIMER_MS.toLong()))
+        // Make a listener which can observe agentWeaker lost later.
         val callbackWeaker = TestableNetworkCallback(timeoutMs = DEFAULT_TIMEOUT_MS)
         registerNetworkCallback(NetworkRequest.Builder()
                 .clearCapabilities()
                 .addTransportType(TRANSPORT_TEST)
-                .setNetworkSpecifier(CompatUtil.makeEthernetNetworkSpecifier(name1))
+                .setNetworkSpecifier(CompatUtil.makeEthernetNetworkSpecifier(specifierWeaker))
                 .build(), callbackWeaker)
-        callbackWeaker.expectAvailableCallbacks(agent1.network!!)
+        callbackWeaker.expectAvailableCallbacks(agentWeaker.network!!)
 
-        // Connect the second agent with a score better than agent1. Verify the callback for
-        // agent1 sees the linger expiry while the callback for both sees the winner.
+        // Connect the agentStronger with a score better than agentWeaker. Verify the callback for
+        // agentWeaker sees the linger expiry while the callback for both sees the winner.
         // Record linger start timestamp prior to send score to prevent possible race, the actual
         // timestamp should be slightly late than this since the service handles update
         // network score asynchronously.
         val lingerStart = SystemClock.elapsedRealtime()
-        val agent2 = createNetworkAgent(name = name2)
-        agent2.register()
-        agent2.markConnected()
-        callback.expectAvailableCallbacks(agent2.network!!)
-        callbackWeaker.expectCallback<Losing>(agent1.network!!)
+        val agentStronger = createNetworkAgent(specifier = specifierStronger)
+        agentStronger.register()
+        agentStronger.markConnected()
+        commonCallback.expectAvailableCallbacks(agentStronger.network!!)
+        callbackWeaker.expectCallback<Losing>(agentWeaker.network!!)
         val expectedRemainingLingerDuration = lingerStart +
                 NetworkAgent.MIN_LINGER_TIMER_MS.toLong() - SystemClock.elapsedRealtime()
         // If the available callback is too late. The remaining duration will be reduced.
         assertTrue(expectedRemainingLingerDuration > 0,
                 "expected remaining linger duration is $expectedRemainingLingerDuration")
         callbackWeaker.assertNoCallback(expectedRemainingLingerDuration)
-        callbackWeaker.expectCallback<Lost>(agent1.network!!)
+        callbackWeaker.expectCallback<Lost>(agentWeaker.network!!)
     }
 
     @Test
     @IgnoreUpTo(Build.VERSION_CODES.R)
     fun testSetSubscriberId() {
-        val name = "TEST-AGENT"
         val imsi = UUID.randomUUID().toString()
         val config = NetworkAgentConfig.Builder().setSubscriberId(imsi).build()
 
-        val request: NetworkRequest = NetworkRequest.Builder()
-                .clearCapabilities()
-                .addTransportType(TRANSPORT_TEST)
-                .setNetworkSpecifier(CompatUtil.makeEthernetNetworkSpecifier(name))
-                .build()
-        val callback = TestableNetworkCallback(timeoutMs = DEFAULT_TIMEOUT_MS)
-        requestNetwork(request, callback)
-
-        val agent = createNetworkAgent(name = name, initialConfig = config)
-        agent.register()
-        agent.markConnected()
-        callback.expectAvailableThenValidatedCallbacks(agent.network!!)
+        val (agent, _) = createConnectedNetworkAgent(initialConfig = config)
         val snapshots = runWithShellPermissionIdentity(ThrowingSupplier {
                 mCM!!.allNetworkStateSnapshots }, NETWORK_SETTINGS)
         val testNetworkSnapshot = snapshots.findLast { it.network == agent.network }
         assertEquals(imsi, testNetworkSnapshot!!.subscriberId)
     }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.R)
+    // TODO: Refactor helper functions to util class and move this test case to
+    //  {@link android.net.cts.ConnectivityManagerTest}.
+    fun testRegisterBestMatchingNetworkCallback() {
+        // Register best matching network callback with additional condition that will be
+        // exercised later. This assumes the test network agent has NOT_VCN_MANAGED in it and
+        // does not have NET_CAPABILITY_TEMPORARILY_NOT_METERED.
+        val bestMatchingCb = TestableNetworkCallback(timeoutMs = DEFAULT_TIMEOUT_MS)
+        registerBestMatchingNetworkCallback(NetworkRequest.Builder()
+                .clearCapabilities()
+                .addTransportType(TRANSPORT_TEST)
+                .addCapability(NET_CAPABILITY_NOT_VCN_MANAGED)
+                .build(), bestMatchingCb, mHandlerThread.threadHandler)
+
+        val (agent1, _) = createConnectedNetworkAgent(specifier = "AGENT-1")
+        bestMatchingCb.expectAvailableThenValidatedCallbacks(agent1.network!!)
+        // Make agent1 worse so when agent2 shows up, the callback will see that.
+        agent1.sendNetworkScore(NetworkScore.Builder().setExiting(true).build())
+        bestMatchingCb.assertNoCallback(NO_CALLBACK_TIMEOUT)
+
+        val (agent2, _) = createConnectedNetworkAgent(specifier = "AGENT-2")
+        bestMatchingCb.expectAvailableDoubleValidatedCallbacks(agent2.network!!)
+
+        // Change something on agent1 to trigger capabilities changed, since the callback
+        // only cares about the best network, verify it received nothing from agent1.
+        val ncAgent1 = agent1.nc
+        ncAgent1.addCapability(NET_CAPABILITY_TEMPORARILY_NOT_METERED)
+        agent1.sendNetworkCapabilities(ncAgent1)
+        bestMatchingCb.assertNoCallback(NO_CALLBACK_TIMEOUT)
+
+        // Make agent1 the best network again, verify the callback now tracks agent1.
+        agent1.sendNetworkScore(NetworkScore.Builder()
+                .setExiting(false).setTransportPrimary(true).build())
+        bestMatchingCb.expectAvailableCallbacks(agent1.network!!)
+
+        // Make agent1 temporary vcn managed, which will not satisfying the request.
+        // Verify the callback switch from/to the other network accordingly.
+        ncAgent1.removeCapability(NET_CAPABILITY_NOT_VCN_MANAGED)
+        agent1.sendNetworkCapabilities(ncAgent1)
+        bestMatchingCb.expectAvailableCallbacks(agent2.network!!)
+        ncAgent1.addCapability(NET_CAPABILITY_NOT_VCN_MANAGED)
+        agent1.sendNetworkCapabilities(ncAgent1)
+        bestMatchingCb.expectAvailableDoubleValidatedCallbacks(agent1.network!!)
+
+        // Verify the callback doesn't care about agent2 disconnect.
+        agent2.unregister()
+        agentsToCleanUp.remove(agent2)
+        bestMatchingCb.assertNoCallback()
+        agent1.unregister()
+        agentsToCleanUp.remove(agent1)
+        bestMatchingCb.expectCallback<Lost>(agent1.network!!)
+
+        // tearDown() will unregister the requests and agents
+    }
 }
diff --git a/tests/cts/net/src/android/net/cts/PacketUtils.java b/tests/cts/net/src/android/net/cts/PacketUtils.java
index 7e622f6..4d924d1 100644
--- a/tests/cts/net/src/android/net/cts/PacketUtils.java
+++ b/tests/cts/net/src/android/net/cts/PacketUtils.java
@@ -19,6 +19,10 @@
 import static android.system.OsConstants.IPPROTO_IPV6;
 import static android.system.OsConstants.IPPROTO_UDP;
 
+import android.util.ArraySet;
+
+import com.android.internal.net.ipsec.ike.crypto.AesXCbcImpl;
+
 import java.net.Inet4Address;
 import java.net.Inet6Address;
 import java.net.InetAddress;
@@ -27,6 +31,7 @@
 import java.security.GeneralSecurityException;
 import java.security.SecureRandom;
 import java.util.Arrays;
+import java.util.Set;
 
 import javax.crypto.Cipher;
 import javax.crypto.Mac;
@@ -43,7 +48,9 @@
     static final int UDP_HDRLEN = 8;
     static final int TCP_HDRLEN = 20;
     static final int TCP_HDRLEN_WITH_TIMESTAMP_OPT = TCP_HDRLEN + 12;
+    static final int ESP_HDRLEN = 8;
     static final int ESP_BLK_SIZE = 4; // ESP has to be 4-byte aligned
+    static final int ESP_TRAILER_LEN = 2;
 
     // Not defined in OsConstants
     static final int IPPROTO_IPV4 = 4;
@@ -52,19 +59,25 @@
     // Encryption parameters
     static final int AES_CBC_IV_LEN = 16;
     static final int AES_CBC_BLK_SIZE = 16;
+    static final int AES_CTR_SALT_LEN = 4;
 
-    static final int AES_CTR_KEY_LEN = 20;
+    static final int AES_CTR_KEY_LEN_20 = 20;
+    static final int AES_CTR_KEY_LEN_28 = 28;
+    static final int AES_CTR_KEY_LEN_36 = 36;
     static final int AES_CTR_BLK_SIZE = ESP_BLK_SIZE;
     static final int AES_CTR_IV_LEN = 8;
 
     // AEAD parameters
     static final int AES_GCM_IV_LEN = 8;
     static final int AES_GCM_BLK_SIZE = 4;
+    static final int CHACHA20_POLY1305_KEY_LEN = 36;
     static final int CHACHA20_POLY1305_BLK_SIZE = ESP_BLK_SIZE;
     static final int CHACHA20_POLY1305_IV_LEN = 8;
+    static final int CHACHA20_POLY1305_SALT_LEN = 4;
     static final int CHACHA20_POLY1305_ICV_LEN = 16;
 
     // Authentication parameters
+    static final int HMAC_SHA256_ICV_LEN = 16;
     static final int HMAC_SHA512_KEY_LEN = 64;
     static final int HMAC_SHA512_ICV_LEN = 32;
     static final int AES_XCBC_KEY_LEN = 16;
@@ -72,10 +85,25 @@
     static final int AES_CMAC_KEY_LEN = 16;
     static final int AES_CMAC_ICV_LEN = 12;
 
+    // Block counter field should be 32 bits and starts from value one as per RFC 3686
+    static final byte[] AES_CTR_INITIAL_COUNTER = new byte[] {0x00, 0x00, 0x00, 0x01};
+
     // Encryption algorithms
     static final String AES = "AES";
     static final String AES_CBC = "AES/CBC/NoPadding";
+    static final String AES_CTR = "AES/CTR/NoPadding";
+
+    // AEAD algorithms
+    static final String CHACHA20_POLY1305 = "ChaCha20/Poly1305/NoPadding";
+
+    // Authentication algorithms
+    static final String HMAC_MD5 = "HmacMD5";
+    static final String HMAC_SHA1 = "HmacSHA1";
     static final String HMAC_SHA_256 = "HmacSHA256";
+    static final String HMAC_SHA_384 = "HmacSHA384";
+    static final String HMAC_SHA_512 = "HmacSHA512";
+    static final String AES_CMAC = "AESCMAC";
+    static final String AES_XCBC = "AesXCbc";
 
     public interface Payload {
         byte[] getPacketBytes(IpHeader header) throws Exception;
@@ -328,8 +356,9 @@
         public final int nextHeader;
         public final int spi;
         public final int seqNum;
-        public final byte[] key;
         public final byte[] payload;
+        public final EspCipher cipher;
+        public final EspAuth auth;
 
         /**
          * Generic constructor for ESP headers.
@@ -340,11 +369,48 @@
          * calculated using the pre-encryption IP header
          */
         public EspHeader(int nextHeader, int spi, int seqNum, byte[] key, byte[] payload) {
+            this(nextHeader, spi, seqNum, payload, getDefaultCipher(key), getDefaultAuth(key));
+        }
+
+        /**
+         * Generic constructor for ESP headers that allows configuring encryption and authentication
+         * algortihms.
+         *
+         * <p>For Tunnel mode, payload will be a full IP header + attached payloads
+         *
+         * <p>For Transport mode, payload will be only the attached payloads, but with the checksum
+         * calculated using the pre-encryption IP header
+         */
+        public EspHeader(
+                int nextHeader,
+                int spi,
+                int seqNum,
+                byte[] payload,
+                EspCipher cipher,
+                EspAuth auth) {
             this.nextHeader = nextHeader;
             this.spi = spi;
             this.seqNum = seqNum;
-            this.key = key;
             this.payload = payload;
+            this.cipher = cipher;
+            this.auth = auth;
+
+            if (cipher instanceof EspCipherNull && auth instanceof EspAuthNull) {
+                throw new IllegalArgumentException("No algorithm is provided");
+            }
+
+            if (cipher instanceof EspAeadCipher && !(auth instanceof EspAuthNull)) {
+                throw new IllegalArgumentException(
+                        "AEAD is provided with an authentication" + " algorithm.");
+            }
+        }
+
+        private static EspCipher getDefaultCipher(byte[] key) {
+            return new EspCryptCipher(AES_CBC, AES_CBC_BLK_SIZE, key, AES_CBC_IV_LEN);
+        }
+
+        private static EspAuth getDefaultAuth(byte[] key) {
+            return new EspAuth(HMAC_SHA_256, key, HMAC_SHA256_ICV_LEN);
         }
 
         public int getProtocolId() {
@@ -352,9 +418,10 @@
         }
 
         public short length() {
-            // ALWAYS uses AES-CBC, HMAC-SHA256 (128b trunc len)
-            return (short)
-                    calculateEspPacketSize(payload.length, AES_CBC_IV_LEN, AES_CBC_BLK_SIZE, 128);
+            final int icvLen =
+                    cipher instanceof EspAeadCipher ? ((EspAeadCipher) cipher).icvLen : auth.icvLen;
+            return calculateEspPacketSize(
+                    payload.length, cipher.ivLen, cipher.blockSize, icvLen * 8);
         }
 
         public byte[] getPacketBytes(IpHeader header) throws Exception {
@@ -368,58 +435,12 @@
             ByteBuffer espPayloadBuffer = ByteBuffer.allocate(DATA_BUFFER_LEN);
             espPayloadBuffer.putInt(spi);
             espPayloadBuffer.putInt(seqNum);
-            espPayloadBuffer.put(getCiphertext(key));
 
-            espPayloadBuffer.put(getIcv(getByteArrayFromBuffer(espPayloadBuffer)), 0, 16);
+            espPayloadBuffer.put(cipher.getCipherText(nextHeader, payload, spi, seqNum));
+            espPayloadBuffer.put(auth.getIcv(getByteArrayFromBuffer(espPayloadBuffer)));
+
             resultBuffer.put(getByteArrayFromBuffer(espPayloadBuffer));
         }
-
-        private byte[] getIcv(byte[] authenticatedSection) throws GeneralSecurityException {
-            Mac sha256HMAC = Mac.getInstance(HMAC_SHA_256);
-            SecretKeySpec authKey = new SecretKeySpec(key, HMAC_SHA_256);
-            sha256HMAC.init(authKey);
-
-            return sha256HMAC.doFinal(authenticatedSection);
-        }
-
-        /**
-         * Encrypts and builds ciphertext block. Includes the IV, Padding and Next-Header blocks
-         *
-         * <p>The ciphertext does NOT include the SPI/Sequence numbers, or the ICV.
-         */
-        private byte[] getCiphertext(byte[] key) throws GeneralSecurityException {
-            int paddedLen = calculateEspEncryptedLength(payload.length, AES_CBC_BLK_SIZE);
-            ByteBuffer paddedPayload = ByteBuffer.allocate(paddedLen);
-            paddedPayload.put(payload);
-
-            // Add padding - consecutive integers from 0x01
-            int pad = 1;
-            while (paddedPayload.position() < paddedPayload.limit()) {
-                paddedPayload.put((byte) pad++);
-            }
-
-            paddedPayload.position(paddedPayload.limit() - 2);
-            paddedPayload.put((byte) (paddedLen - 2 - payload.length)); // Pad length
-            paddedPayload.put((byte) nextHeader);
-
-            // Generate Initialization Vector
-            byte[] iv = new byte[AES_CBC_IV_LEN];
-            new SecureRandom().nextBytes(iv);
-            IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
-            SecretKeySpec secretKeySpec = new SecretKeySpec(key, AES);
-
-            // Encrypt payload
-            Cipher cipher = Cipher.getInstance(AES_CBC);
-            cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, ivParameterSpec);
-            byte[] encrypted = cipher.doFinal(getByteArrayFromBuffer(paddedPayload));
-
-            // Build ciphertext
-            ByteBuffer cipherText = ByteBuffer.allocate(AES_CBC_IV_LEN + encrypted.length);
-            cipherText.put(iv);
-            cipherText.put(encrypted);
-
-            return getByteArrayFromBuffer(cipherText);
-        }
     }
 
     private static int addAndWrapForChecksum(int currentChecksum, int value) {
@@ -436,15 +457,14 @@
         return (short) ((~val) & 0xffff);
     }
 
-    public static int calculateEspPacketSize(
+    public static short calculateEspPacketSize(
             int payloadLen, int cryptIvLength, int cryptBlockSize, int authTruncLen) {
-        final int ESP_HDRLEN = 4 + 4; // SPI + Seq#
         final int ICV_LEN = authTruncLen / 8; // Auth trailer; based on truncation length
-        payloadLen += cryptIvLength; // Initialization Vector
 
         // Align to block size of encryption algorithm
         payloadLen = calculateEspEncryptedLength(payloadLen, cryptBlockSize);
-        return payloadLen + ESP_HDRLEN + ICV_LEN;
+        payloadLen += cryptIvLength; // Initialization Vector
+        return (short) (payloadLen + ESP_HDRLEN + ICV_LEN);
     }
 
     private static int calculateEspEncryptedLength(int payloadLen, int cryptBlockSize) {
@@ -475,6 +495,237 @@
         }
     }
 
+    public abstract static class EspCipher {
+        protected static final int SALT_LEN_UNUSED = 0;
+
+        public final String algoName;
+        public final int blockSize;
+        public final byte[] key;
+        public final int ivLen;
+        public final int saltLen;
+        protected byte[] mIv;
+
+        public EspCipher(String algoName, int blockSize, byte[] key, int ivLen, int saltLen) {
+            this.algoName = algoName;
+            this.blockSize = blockSize;
+            this.key = key;
+            this.ivLen = ivLen;
+            this.saltLen = saltLen;
+            this.mIv = getIv(ivLen);
+        }
+
+        public void updateIv(byte[] iv) {
+            this.mIv = iv;
+        }
+
+        public static byte[] getPaddedPayload(int nextHeader, byte[] payload, int blockSize) {
+            final int paddedLen = calculateEspEncryptedLength(payload.length, blockSize);
+            final ByteBuffer paddedPayload = ByteBuffer.allocate(paddedLen);
+            paddedPayload.put(payload);
+
+            // Add padding - consecutive integers from 0x01
+            byte pad = 1;
+            while (paddedPayload.position() < paddedPayload.limit() - ESP_TRAILER_LEN) {
+                paddedPayload.put((byte) pad++);
+            }
+
+            // Add padding length and next header
+            paddedPayload.put((byte) (paddedLen - ESP_TRAILER_LEN - payload.length));
+            paddedPayload.put((byte) nextHeader);
+
+            return getByteArrayFromBuffer(paddedPayload);
+        }
+
+        private static byte[] getIv(int ivLen) {
+            final byte[] iv = new byte[ivLen];
+            new SecureRandom().nextBytes(iv);
+            return iv;
+        }
+
+        public abstract byte[] getCipherText(int nextHeader, byte[] payload, int spi, int seqNum)
+                throws GeneralSecurityException;
+    }
+
+    public static final class EspCipherNull extends EspCipher {
+        private static final String CRYPT_NULL = "CRYPT_NULL";
+        private static final int IV_LEN_UNUSED = 0;
+        private static final byte[] KEY_UNUSED = new byte[0];
+
+        private static final EspCipherNull sInstance = new EspCipherNull();
+
+        private EspCipherNull() {
+            super(CRYPT_NULL, ESP_BLK_SIZE, KEY_UNUSED, IV_LEN_UNUSED, SALT_LEN_UNUSED);
+        }
+
+        public static EspCipherNull getInstance() {
+            return sInstance;
+        }
+
+        @Override
+        public byte[] getCipherText(int nextHeader, byte[] payload, int spi, int seqNum)
+                throws GeneralSecurityException {
+            return getPaddedPayload(nextHeader, payload, blockSize);
+        }
+    }
+
+    public static final class EspCryptCipher extends EspCipher {
+        public EspCryptCipher(String algoName, int blockSize, byte[] key, int ivLen) {
+            this(algoName, blockSize, key, ivLen, SALT_LEN_UNUSED);
+        }
+
+        public EspCryptCipher(String algoName, int blockSize, byte[] key, int ivLen, int saltLen) {
+            super(algoName, blockSize, key, ivLen, saltLen);
+        }
+
+        @Override
+        public byte[] getCipherText(int nextHeader, byte[] payload, int spi, int seqNum)
+                throws GeneralSecurityException {
+            final IvParameterSpec ivParameterSpec;
+            final SecretKeySpec secretKeySpec;
+
+            if (AES_CBC.equals(algoName)) {
+                ivParameterSpec = new IvParameterSpec(mIv);
+                secretKeySpec = new SecretKeySpec(key, algoName);
+            } else if (AES_CTR.equals(algoName)) {
+                // Provided key consists of encryption/decryption key plus 4-byte salt. Salt is used
+                // with ESP payload IV and initial block counter value to build IvParameterSpec.
+                final byte[] secretKey = Arrays.copyOfRange(key, 0, key.length - saltLen);
+                final byte[] salt = Arrays.copyOfRange(key, secretKey.length, key.length);
+                secretKeySpec = new SecretKeySpec(secretKey, algoName);
+
+                final ByteBuffer ivParameterBuffer =
+                        ByteBuffer.allocate(mIv.length + saltLen + AES_CTR_INITIAL_COUNTER.length);
+                ivParameterBuffer.put(salt);
+                ivParameterBuffer.put(mIv);
+                ivParameterBuffer.put(AES_CTR_INITIAL_COUNTER);
+                ivParameterSpec = new IvParameterSpec(ivParameterBuffer.array());
+            } else {
+                throw new IllegalArgumentException("Invalid algorithm " + algoName);
+            }
+
+            // Encrypt payload
+            final Cipher cipher = Cipher.getInstance(algoName);
+            cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, ivParameterSpec);
+            final byte[] encrypted =
+                    cipher.doFinal(getPaddedPayload(nextHeader, payload, blockSize));
+
+            // Build ciphertext
+            final ByteBuffer cipherText = ByteBuffer.allocate(mIv.length + encrypted.length);
+            cipherText.put(mIv);
+            cipherText.put(encrypted);
+
+            return getByteArrayFromBuffer(cipherText);
+        }
+    }
+
+    public static final class EspAeadCipher extends EspCipher {
+        public final int icvLen;
+
+        public EspAeadCipher(
+                String algoName, int blockSize, byte[] key, int ivLen, int icvLen, int saltLen) {
+            super(algoName, blockSize, key, ivLen, saltLen);
+            this.icvLen = icvLen;
+        }
+
+        @Override
+        public byte[] getCipherText(int nextHeader, byte[] payload, int spi, int seqNum)
+                throws GeneralSecurityException {
+            // Provided key consists of encryption/decryption key plus salt. Salt is used
+            // with ESP payload IV to build IvParameterSpec.
+            final byte[] secretKey = Arrays.copyOfRange(key, 0, key.length - saltLen);
+            final byte[] salt = Arrays.copyOfRange(key, secretKey.length, key.length);
+
+            final SecretKeySpec secretKeySpec = new SecretKeySpec(secretKey, algoName);
+
+            final ByteBuffer ivParameterBuffer = ByteBuffer.allocate(saltLen + mIv.length);
+            ivParameterBuffer.put(salt);
+            ivParameterBuffer.put(mIv);
+            final IvParameterSpec ivParameterSpec = new IvParameterSpec(ivParameterBuffer.array());
+
+            final ByteBuffer aadBuffer = ByteBuffer.allocate(ESP_HDRLEN);
+            aadBuffer.putInt(spi);
+            aadBuffer.putInt(seqNum);
+
+            // Encrypt payload
+            final Cipher cipher = Cipher.getInstance(algoName);
+            cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, ivParameterSpec);
+            cipher.updateAAD(aadBuffer.array());
+            final byte[] encryptedTextAndIcv =
+                    cipher.doFinal(getPaddedPayload(nextHeader, payload, blockSize));
+
+            // Build ciphertext
+            final ByteBuffer cipherText =
+                    ByteBuffer.allocate(mIv.length + encryptedTextAndIcv.length);
+            cipherText.put(mIv);
+            cipherText.put(encryptedTextAndIcv);
+
+            return getByteArrayFromBuffer(cipherText);
+        }
+    }
+
+    public static class EspAuth {
+        public final String algoName;
+        public final byte[] key;
+        public final int icvLen;
+
+        private static final Set<String> JCE_SUPPORTED_MACS = new ArraySet<>();
+
+        static {
+            JCE_SUPPORTED_MACS.add(HMAC_MD5);
+            JCE_SUPPORTED_MACS.add(HMAC_SHA1);
+            JCE_SUPPORTED_MACS.add(HMAC_SHA_256);
+            JCE_SUPPORTED_MACS.add(HMAC_SHA_384);
+            JCE_SUPPORTED_MACS.add(HMAC_SHA_512);
+            JCE_SUPPORTED_MACS.add(AES_CMAC);
+        }
+
+        public EspAuth(String algoName, byte[] key, int icvLen) {
+            this.algoName = algoName;
+            this.key = key;
+            this.icvLen = icvLen;
+        }
+
+        public byte[] getIcv(byte[] authenticatedSection) throws GeneralSecurityException {
+            if (AES_XCBC.equals(algoName)) {
+                final Cipher aesCipher = Cipher.getInstance(AES_CBC);
+                return new AesXCbcImpl().mac(key, authenticatedSection, true /* needTruncation */);
+            } else if (JCE_SUPPORTED_MACS.contains(algoName)) {
+                final Mac mac = Mac.getInstance(algoName);
+                final SecretKeySpec authKey = new SecretKeySpec(key, algoName);
+                mac.init(authKey);
+
+                final ByteBuffer buffer = ByteBuffer.wrap(mac.doFinal(authenticatedSection));
+                final byte[] icv = new byte[icvLen];
+                buffer.get(icv);
+                return icv;
+            } else {
+                throw new IllegalArgumentException("Invalid algorithm: " + algoName);
+            }
+        }
+    }
+
+    public static final class EspAuthNull extends EspAuth {
+        private static final String AUTH_NULL = "AUTH_NULL";
+        private static final int ICV_LEN_UNUSED = 0;
+        private static final byte[] KEY_UNUSED = new byte[0];
+        private static final byte[] ICV_EMPTY = new byte[0];
+
+        private static final EspAuthNull sInstance = new EspAuthNull();
+
+        private EspAuthNull() {
+            super(AUTH_NULL, KEY_UNUSED, ICV_LEN_UNUSED);
+        }
+
+        public static EspAuthNull getInstance() {
+            return sInstance;
+        }
+
+        @Override
+        public byte[] getIcv(byte[] authenticatedSection) throws GeneralSecurityException {
+            return ICV_EMPTY;
+        }
+    }
+
     /*
      * Debug printing
      */
diff --git a/tests/cts/net/src/android/net/cts/TestNetworkRunnable.java b/tests/cts/net/src/android/net/cts/TestNetworkRunnable.java
new file mode 100644
index 0000000..0eb5644
--- /dev/null
+++ b/tests/cts/net/src/android/net/cts/TestNetworkRunnable.java
@@ -0,0 +1,132 @@
+/*
+ * Copyright (C) 2020 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 android.net.cts;
+
+import static android.Manifest.permission.MANAGE_TEST_NETWORKS;
+import static android.content.pm.PackageManager.PERMISSION_GRANTED;
+import static android.net.cts.util.CtsNetUtils.TestNetworkCallback;
+
+import static com.android.compatibility.common.util.SystemUtil.runWithShellPermissionIdentity;
+
+import android.content.Context;
+import android.net.ConnectivityManager;
+import android.net.LinkAddress;
+import android.net.Network;
+import android.net.TestNetworkInterface;
+import android.net.TestNetworkManager;
+import android.net.cts.util.CtsNetUtils;
+
+import androidx.test.InstrumentationRegistry;
+
+import com.android.compatibility.common.util.ThrowingRunnable;
+
+import java.net.Inet4Address;
+import java.net.InetAddress;
+
+/** This class supports running a test with a test network. */
+public class TestNetworkRunnable implements ThrowingRunnable {
+    private static final int IP4_PREFIX_LEN = 32;
+    private static final int IP6_PREFIX_LEN = 128;
+
+    private static final InetAddress DEFAULT_ADDRESS_4 =
+            InetAddress.parseNumericAddress("192.2.0.2");
+    private static final InetAddress DEFAULT_ADDRESS_6 =
+            InetAddress.parseNumericAddress("2001:db8:1::2");
+
+    private static final Context sContext = InstrumentationRegistry.getContext();
+    private static final ConnectivityManager sCm =
+            sContext.getSystemService(ConnectivityManager.class);
+
+    private final Test mTest;
+
+    public TestNetworkRunnable(Test test) {
+        mTest = test;
+    }
+
+    private void runTest() throws Exception {
+        final TestNetworkManager tnm = sContext.getSystemService(TestNetworkManager.class);
+
+        // Non-final; these variables ensure we clean up properly after our test if we
+        // have allocated test network resources
+        TestNetworkInterface testIface = null;
+        TestNetworkCallback tunNetworkCallback = null;
+
+        final CtsNetUtils ctsNetUtils = new CtsNetUtils(sContext);
+        final InetAddress[] addresses = mTest.getTestNetworkAddresses();
+        final LinkAddress[] linkAddresses = new LinkAddress[addresses.length];
+        for (int i = 0; i < addresses.length; i++) {
+            InetAddress address = addresses[i];
+            if (address instanceof Inet4Address) {
+                linkAddresses[i] = new LinkAddress(address, IP4_PREFIX_LEN);
+            } else {
+                linkAddresses[i] = new LinkAddress(address, IP6_PREFIX_LEN);
+            }
+        }
+
+        try {
+            // Build underlying test network
+            testIface = tnm.createTunInterface(linkAddresses);
+
+            // Hold on to this callback to ensure network does not get reaped.
+            tunNetworkCallback = ctsNetUtils.setupAndGetTestNetwork(testIface.getInterfaceName());
+
+            mTest.runTest(testIface, tunNetworkCallback);
+        } finally {
+            try {
+                mTest.cleanupTest();
+            } catch (Exception e) {
+                // No action
+            }
+
+            if (testIface != null) {
+                testIface.getFileDescriptor().close();
+            }
+
+            if (tunNetworkCallback != null) {
+                sCm.unregisterNetworkCallback(tunNetworkCallback);
+            }
+
+            final Network testNetwork = tunNetworkCallback.currentNetwork;
+            if (testNetwork != null) {
+                tnm.teardownTestNetwork(testNetwork);
+            }
+        }
+    }
+
+    @Override
+    public void run() throws Exception {
+        if (sContext.checkSelfPermission(MANAGE_TEST_NETWORKS) == PERMISSION_GRANTED) {
+            runTest();
+        } else {
+            runWithShellPermissionIdentity(this::runTest, MANAGE_TEST_NETWORKS);
+        }
+    }
+
+    /** Interface for test caller to configure the test that will be run with a test network */
+    public interface Test {
+        /** Runs the test with a test network */
+        void runTest(TestNetworkInterface testIface, TestNetworkCallback tunNetworkCallback)
+                throws Exception;
+
+        /** Cleans up when the test is finished or interrupted */
+        void cleanupTest();
+
+        /** Returns the IP addresses that will be used by the test network */
+        default InetAddress[] getTestNetworkAddresses() {
+            return new InetAddress[] {DEFAULT_ADDRESS_4, DEFAULT_ADDRESS_6};
+        }
+    }
+}
diff --git a/tests/cts/net/src/android/net/cts/TunUtils.java b/tests/cts/net/src/android/net/cts/TunUtils.java
index 7887385..d8e39b4 100644
--- a/tests/cts/net/src/android/net/cts/TunUtils.java
+++ b/tests/cts/net/src/android/net/cts/TunUtils.java
@@ -147,6 +147,10 @@
         return espPkt; // We've found the packet we're looking for.
     }
 
+    public byte[] awaitEspPacket(int spi, boolean useEncap) throws Exception {
+        return awaitPacket((pkt) -> isEsp(pkt, spi, useEncap));
+    }
+
     private static boolean isSpiEqual(byte[] pkt, int espOffset, int spi) {
         // Check SPI byte by byte.
         return pkt[espOffset] == (byte) ((spi >>> 24) & 0xff)