VCN: Allow restricting TRANSPORT_TEST for CTS

This commit updates VCN to be able to mark a test network
as restricted. Additionally, when VCN is in safe mode, VCN will
delegate to the Android system to decide the restriction policy of
the test network.

This commit allows CTS to verify VCN's ability of restricting
networks based on the transport type.

Bug: 263415068
Test: atest VcnManagerTest (new tests)
Change-Id: I5fe0be4ce445a4d9c20cbef5ee4a2eb55403b3c8
diff --git a/core/java/android/net/vcn/VcnConfig.java b/core/java/android/net/vcn/VcnConfig.java
index 6f9c9dd..a27e923 100644
--- a/core/java/android/net/vcn/VcnConfig.java
+++ b/core/java/android/net/vcn/VcnConfig.java
@@ -16,6 +16,7 @@
 package android.net.vcn;
 
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
+import static android.net.NetworkCapabilities.TRANSPORT_TEST;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 
 import static com.android.internal.annotations.VisibleForTesting.Visibility;
@@ -75,6 +76,7 @@
     static {
         ALLOWED_TRANSPORTS.add(TRANSPORT_WIFI);
         ALLOWED_TRANSPORTS.add(TRANSPORT_CELLULAR);
+        ALLOWED_TRANSPORTS.add(TRANSPORT_TEST);
     }
 
     private static final String PACKAGE_NAME_KEY = "mPackageName";
@@ -155,6 +157,11 @@
                                 + transport
                                 + " which might be from a new version of VcnConfig");
             }
+
+            if (transport == TRANSPORT_TEST && !mIsTestModeProfile) {
+                throw new IllegalArgumentException(
+                        "Found TRANSPORT_TEST in a non-test-mode profile");
+            }
         }
     }
 
diff --git a/services/core/java/com/android/server/VcnManagementService.java b/services/core/java/com/android/server/VcnManagementService.java
index 78d4708..e8c85ce 100644
--- a/services/core/java/com/android/server/VcnManagementService.java
+++ b/services/core/java/com/android/server/VcnManagementService.java
@@ -1074,9 +1074,10 @@
                             subGrp, mLastSnapshot, mConfigs.get(subGrp));
                     for (int restrictedTransport : restrictedTransports) {
                         if (ncCopy.hasTransport(restrictedTransport)) {
-                            if (restrictedTransport == TRANSPORT_CELLULAR) {
-                                // Only make a cell network as restricted when the VCN is in
-                                // active mode.
+                            if (restrictedTransport == TRANSPORT_CELLULAR
+                                    || restrictedTransport == TRANSPORT_TEST) {
+                                // For cell or test network, only mark it as restricted when
+                                // the VCN is in active mode.
                                 isRestricted |= (vcn.getStatus() == VCN_STATUS_CODE_ACTIVE);
                             } else {
                                 isRestricted = true;
diff --git a/tests/vcn/java/android/net/vcn/VcnConfigTest.java b/tests/vcn/java/android/net/vcn/VcnConfigTest.java
index b313c9f..73a0a61 100644
--- a/tests/vcn/java/android/net/vcn/VcnConfigTest.java
+++ b/tests/vcn/java/android/net/vcn/VcnConfigTest.java
@@ -17,6 +17,7 @@
 package android.net.vcn;
 
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
+import static android.net.NetworkCapabilities.TRANSPORT_TEST;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 
 import static org.junit.Assert.assertEquals;
@@ -160,6 +161,37 @@
         assertNotEquals(config, configNotEqual);
     }
 
+    private VcnConfig buildConfigRestrictTransportTest(boolean isTestMode) throws Exception {
+        VcnConfig.Builder builder =
+                new VcnConfig.Builder(mContext)
+                        .setRestrictedUnderlyingNetworkTransports(Set.of(TRANSPORT_TEST));
+        if (isTestMode) {
+            builder.setIsTestModeProfile();
+        }
+
+        for (VcnGatewayConnectionConfig gatewayConnectionConfig : GATEWAY_CONNECTION_CONFIGS) {
+            builder.addGatewayConnectionConfig(gatewayConnectionConfig);
+        }
+
+        return builder.build();
+    }
+
+    @Test
+    public void testRestrictTransportTestInTestModeProfile() throws Exception {
+        final VcnConfig config = buildConfigRestrictTransportTest(true /*  isTestMode */);
+        assertEquals(Set.of(TRANSPORT_TEST), config.getRestrictedUnderlyingNetworkTransports());
+    }
+
+    @Test
+    public void testRestrictTransportTestInNonTestModeProfile() throws Exception {
+        try {
+            buildConfigRestrictTransportTest(false /*  isTestMode */);
+            fail("Expected exception because the config is not a test mode profile");
+        } catch (Exception expected) {
+
+        }
+    }
+
     @Test
     public void testParceling() {
         final VcnConfig config = buildTestConfig(mContext);