From 653a8e6abe38af8ba554c3915aed08722b04b5ca Mon Sep 17 00:00:00 2001
From: Yan Yan <evitayan@google.com>
Date: Tue, 31 Jan 2023 01:38:17 +0000
Subject: [PATCH] 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
---
 core/java/android/net/vcn/VcnConfig.java      |  7 ++++
 .../android/server/VcnManagementService.java  |  7 ++--
 .../java/android/net/vcn/VcnConfigTest.java   | 32 +++++++++++++++++++
 3 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/core/java/android/net/vcn/VcnConfig.java b/core/java/android/net/vcn/VcnConfig.java
index 6f9c9dd918d1..a27e9230d473 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 @@ public final class VcnConfig implements Parcelable {
     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 @@ public final class VcnConfig implements Parcelable {
                                 + 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 78d4708e70a2..e8c85ce68f22 100644
--- a/services/core/java/com/android/server/VcnManagementService.java
+++ b/services/core/java/com/android/server/VcnManagementService.java
@@ -1074,9 +1074,10 @@ public class VcnManagementService extends IVcnManagementService.Stub {
                             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 b313c9fc6c28..73a0a6183cb6 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 @@ public class VcnConfigTest {
         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);