Merge "[AWARE] Add UID checking to WifiAwareNetworkSpecifier"
diff --git a/core/java/android/net/NetworkCapabilities.java b/core/java/android/net/NetworkCapabilities.java
index c28294f..1da0d28 100644
--- a/core/java/android/net/NetworkCapabilities.java
+++ b/core/java/android/net/NetworkCapabilities.java
@@ -18,7 +18,9 @@
 
 import android.os.Parcel;
 import android.os.Parcelable;
-import android.util.Log;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.util.BitUtils;
 
 import java.util.Objects;
 
@@ -229,7 +231,8 @@
      * Capabilities that suggest that a network is restricted.
      * {@see #maybeMarkCapabilitiesRestricted}.
      */
-    private static final long RESTRICTED_CAPABILITIES =
+    @VisibleForTesting
+    /* package */ static final long RESTRICTED_CAPABILITIES =
             (1 << NET_CAPABILITY_CBS) |
             (1 << NET_CAPABILITY_DUN) |
             (1 << NET_CAPABILITY_EIMS) |
@@ -240,6 +243,17 @@
             (1 << NET_CAPABILITY_XCAP);
 
     /**
+     * Capabilities that suggest that a network is unrestricted.
+     * {@see #maybeMarkCapabilitiesRestricted}.
+     */
+    @VisibleForTesting
+    /* package */ static final long UNRESTRICTED_CAPABILITIES =
+            (1 << NET_CAPABILITY_INTERNET) |
+            (1 << NET_CAPABILITY_MMS) |
+            (1 << NET_CAPABILITY_SUPL) |
+            (1 << NET_CAPABILITY_WIFI_P2P);
+
+    /**
      * Adds the given capability to this {@code NetworkCapability} instance.
      * Multiple capabilities may be applied sequentially.  Note that when searching
      * for a network to satisfy a request, all capabilities requested must be satisfied.
@@ -279,7 +293,7 @@
      * @hide
      */
     public int[] getCapabilities() {
-        return enumerateBits(mNetworkCapabilities);
+        return BitUtils.unpackBits(mNetworkCapabilities);
     }
 
     /**
@@ -295,19 +309,6 @@
         return ((mNetworkCapabilities & (1 << capability)) != 0);
     }
 
-    private int[] enumerateBits(long val) {
-        int size = Long.bitCount(val);
-        int[] result = new int[size];
-        int index = 0;
-        int resource = 0;
-        while (val > 0) {
-            if ((val & 1) == 1) result[index++] = resource;
-            val = val >> 1;
-            resource++;
-        }
-        return result;
-    }
-
     private void combineNetCapabilities(NetworkCapabilities nc) {
         this.mNetworkCapabilities |= nc.mNetworkCapabilities;
     }
@@ -365,12 +366,16 @@
      * @hide
      */
     public void maybeMarkCapabilitiesRestricted() {
-        // If all the capabilities are typically provided by restricted networks, conclude that this
-        // network is restricted.
-        if ((mNetworkCapabilities & ~(DEFAULT_CAPABILITIES | RESTRICTED_CAPABILITIES)) == 0 &&
-                // Must have at least some restricted capabilities, otherwise a request for an
-                // internet-less network will get marked restricted.
-                (mNetworkCapabilities & RESTRICTED_CAPABILITIES) != 0) {
+        // Verify there aren't any unrestricted capabilities.  If there are we say
+        // the whole thing is unrestricted.
+        final boolean hasUnrestrictedCapabilities =
+                ((mNetworkCapabilities & UNRESTRICTED_CAPABILITIES) != 0);
+
+        // Must have at least some restricted capabilities.
+        final boolean hasRestrictedCapabilities =
+                ((mNetworkCapabilities & RESTRICTED_CAPABILITIES) != 0);
+
+        if (hasRestrictedCapabilities && !hasUnrestrictedCapabilities) {
             removeCapability(NET_CAPABILITY_NOT_RESTRICTED);
         }
     }
@@ -416,6 +421,15 @@
     private static final int MIN_TRANSPORT = TRANSPORT_CELLULAR;
     private static final int MAX_TRANSPORT = TRANSPORT_WIFI_AWARE;
 
+    private static final String[] TRANSPORT_NAMES = {
+        "CELLULAR",
+        "WIFI",
+        "BLUETOOTH",
+        "ETHERNET",
+        "VPN",
+        "WIFI_AWARE"
+    };
+
     /**
      * Adds the given transport type to this {@code NetworkCapability} instance.
      * Multiple transports may be applied sequentially.  Note that when searching
@@ -462,7 +476,7 @@
      * @hide
      */
     public int[] getTransportTypes() {
-        return enumerateBits(mTransportTypes);
+        return BitUtils.unpackBits(mTransportTypes);
     }
 
     /**
@@ -869,18 +883,23 @@
      * @hide
      */
     public static String transportNamesOf(int[] types) {
-        String transports = "";
-        for (int i = 0; i < types.length;) {
-            switch (types[i]) {
-                case TRANSPORT_CELLULAR:    transports += "CELLULAR"; break;
-                case TRANSPORT_WIFI:        transports += "WIFI"; break;
-                case TRANSPORT_BLUETOOTH:   transports += "BLUETOOTH"; break;
-                case TRANSPORT_ETHERNET:    transports += "ETHERNET"; break;
-                case TRANSPORT_VPN:         transports += "VPN"; break;
-                case TRANSPORT_WIFI_AWARE:  transports += "WIFI_AWARE"; break;
-            }
-            if (++i < types.length) transports += "|";
+        if (types == null || types.length == 0) {
+            return "";
         }
-        return transports;
+        StringBuilder transports = new StringBuilder();
+        for (int t : types) {
+            transports.append("|").append(transportNameOf(t));
+        }
+        return transports.substring(1);
+    }
+
+    /**
+     * @hide
+     */
+    public static String transportNameOf(int transport) {
+        if (transport < 0 || TRANSPORT_NAMES.length <= transport) {
+            return "UNKNOWN";
+        }
+        return TRANSPORT_NAMES[transport];
     }
 }
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 96d8907..88bc54d 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -377,11 +377,6 @@
     private static final int EVENT_SET_ACCEPT_UNVALIDATED = 28;
 
     /**
-     * used to specify whether a network should not be penalized when it becomes unvalidated.
-     */
-    private static final int EVENT_SET_AVOID_UNVALIDATED = 35;
-
-    /**
      * used to ask the user to confirm a connection to an unvalidated network.
      * obj  = network
      */
@@ -398,6 +393,16 @@
      */
     private static final int EVENT_REGISTER_NETWORK_LISTENER_WITH_INTENT = 31;
 
+    /**
+     * used to specify whether a network should not be penalized when it becomes unvalidated.
+     */
+    private static final int EVENT_SET_AVOID_UNVALIDATED = 35;
+
+    /**
+     * used to trigger revalidation of a network.
+     */
+    private static final int EVENT_REVALIDATE_NETWORK = 36;
+
     private static String eventName(int what) {
         return sMagicDecoderRing.get(what, Integer.toString(what));
     }
@@ -870,8 +875,8 @@
     }
 
     private void handleMobileDataAlwaysOn() {
-        final boolean enable = (Settings.Global.getInt(
-                mContext.getContentResolver(), Settings.Global.MOBILE_DATA_ALWAYS_ON, 1) == 1);
+        final boolean enable = toBool(Settings.Global.getInt(
+                mContext.getContentResolver(), Settings.Global.MOBILE_DATA_ALWAYS_ON, 1));
         final boolean isEnabled = (mNetworkRequests.get(mDefaultMobileDataRequest) != null);
         if (enable == isEnabled) {
             return;  // Nothing to do.
@@ -2224,7 +2229,7 @@
                 }
                 case NetworkMonitor.EVENT_PROVISIONING_NOTIFICATION: {
                     final int netId = msg.arg2;
-                    final boolean visible = (msg.arg1 != 0);
+                    final boolean visible = toBool(msg.arg1);
                     final NetworkAgentInfo nai;
                     synchronized (mNetworkForNetId) {
                         nai = mNetworkForNetId.get(netId);
@@ -2677,7 +2682,7 @@
     public void setAcceptUnvalidated(Network network, boolean accept, boolean always) {
         enforceConnectivityInternalPermission();
         mHandler.sendMessage(mHandler.obtainMessage(EVENT_SET_ACCEPT_UNVALIDATED,
-                accept ? 1 : 0, always ? 1: 0, network));
+                encodeBool(accept), encodeBool(always), network));
     }
 
     @Override
@@ -2714,7 +2719,7 @@
 
         if (always) {
             nai.asyncChannel.sendMessage(
-                    NetworkAgent.CMD_SAVE_ACCEPT_UNVALIDATED, accept ? 1 : 0);
+                    NetworkAgent.CMD_SAVE_ACCEPT_UNVALIDATED, encodeBool(accept));
         }
 
         if (!accept) {
@@ -2915,7 +2920,8 @@
                     break;
                 }
                 case EVENT_SET_ACCEPT_UNVALIDATED: {
-                    handleSetAcceptUnvalidated((Network) msg.obj, msg.arg1 != 0, msg.arg2 != 0);
+                    Network network = (Network) msg.obj;
+                    handleSetAcceptUnvalidated(network, toBool(msg.arg1), toBool(msg.arg2));
                     break;
                 }
                 case EVENT_SET_AVOID_UNVALIDATED: {
@@ -2949,6 +2955,10 @@
                     }
                     break;
                 }
+                case EVENT_REVALIDATE_NETWORK: {
+                    handleReportNetworkConnectivity((Network) msg.obj, msg.arg1, toBool(msg.arg2));
+                    break;
+                }
             }
         }
     }
@@ -3063,9 +3073,10 @@
     @Override
     public boolean isTetheringSupported() {
         enforceTetherAccessPermission();
-        int defaultVal = (mSystemProperties.get("ro.tether.denied").equals("true") ? 0 : 1);
-        boolean tetherEnabledInSettings = (Settings.Global.getInt(mContext.getContentResolver(),
-                Settings.Global.TETHER_SUPPORTED, defaultVal) != 0)
+        int defaultVal = encodeBool(!mSystemProperties.get("ro.tether.denied").equals("true"));
+        boolean tetherSupported = toBool(Settings.Global.getInt(mContext.getContentResolver(),
+                Settings.Global.TETHER_SUPPORTED, defaultVal));
+        boolean tetherEnabledInSettings = tetherSupported
                 && !mUserManager.hasUserRestriction(UserManager.DISALLOW_CONFIG_TETHERING);
 
         // Elevate to system UID to avoid caller requiring MANAGE_USERS permission.
@@ -3077,8 +3088,7 @@
             Binder.restoreCallingIdentity(token);
         }
 
-        return tetherEnabledInSettings && adminUser &&
-               mTethering.hasTetherableConfiguration();
+        return tetherEnabledInSettings && adminUser && mTethering.hasTetherableConfiguration();
     }
 
     @Override
@@ -3156,8 +3166,14 @@
     public void reportNetworkConnectivity(Network network, boolean hasConnectivity) {
         enforceAccessPermission();
         enforceInternetPermission();
+        final int uid = Binder.getCallingUid();
+        final int connectivityInfo = encodeBool(hasConnectivity);
+        mHandler.sendMessage(
+                mHandler.obtainMessage(EVENT_REVALIDATE_NETWORK, uid, connectivityInfo, network));
+    }
 
-        // TODO: execute this logic on ConnectivityService handler.
+    private void handleReportNetworkConnectivity(
+            Network network, int uid, boolean hasConnectivity) {
         final NetworkAgentInfo nai;
         if (network == null) {
             nai = getDefaultNetwork();
@@ -3172,10 +3188,9 @@
         if (hasConnectivity == nai.lastValidated) {
             return;
         }
-        final int uid = Binder.getCallingUid();
         if (DBG) {
-            log("reportNetworkConnectivity(" + nai.network.netId + ", " + hasConnectivity +
-                    ") by " + uid);
+            int netid = nai.network.netId;
+            log("reportNetworkConnectivity(" + netid + ", " + hasConnectivity + ") by " + uid);
         }
         // Validating a network that has not yet connected could result in a call to
         // rematchNetworkAndRequests() which is not meant to work on such networks.
@@ -3680,17 +3695,9 @@
             existing.shutdown();
         }
 
-        try {
-            if (tracker != null) {
-                mNetd.setFirewallEnabled(true);
-                mNetd.setFirewallInterfaceRule("lo", true);
-                mLockdownTracker = tracker;
-                mLockdownTracker.init();
-            } else {
-                mNetd.setFirewallEnabled(false);
-            }
-        } catch (RemoteException e) {
-            // ignored; NMS lives inside system_server
+        if (tracker != null) {
+            mLockdownTracker = tracker;
+            mLockdownTracker.init();
         }
     }
 
@@ -3886,7 +3893,7 @@
         final long ident = Binder.clearCallingIdentity();
         try {
             final ContentResolver cr = mContext.getContentResolver();
-            Settings.Global.putInt(cr, Settings.Global.AIRPLANE_MODE_ON, enable ? 1 : 0);
+            Settings.Global.putInt(cr, Settings.Global.AIRPLANE_MODE_ON, encodeBool(enable));
             Intent intent = new Intent(Intent.ACTION_AIRPLANE_MODE_CHANGED);
             intent.putExtra("state", enable);
             mContext.sendBroadcastAsUser(intent, UserHandle.ALL);
@@ -5566,4 +5573,12 @@
     private void logNetworkEvent(NetworkAgentInfo nai, int evtype) {
         mMetricsLog.log(new NetworkEvent(nai.network.netId, evtype));
     }
+
+    private static boolean toBool(int encodedBoolean) {
+        return encodedBoolean != 0; // Only 0 means false.
+    }
+
+    private static int encodeBool(boolean b) {
+        return b ? 1 : 0;
+    }
 }
diff --git a/tests/net/java/android/net/NetworkCapabilitiesTest.java b/tests/net/java/android/net/NetworkCapabilitiesTest.java
new file mode 100644
index 0000000..e3b06c8
--- /dev/null
+++ b/tests/net/java/android/net/NetworkCapabilitiesTest.java
@@ -0,0 +1,117 @@
+/*
+ * Copyright (C) 2017 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;
+
+import static android.net.NetworkCapabilities.NET_CAPABILITY_CBS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_EIMS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_METERED;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED;
+import static android.net.NetworkCapabilities.RESTRICTED_CAPABILITIES;
+import static android.net.NetworkCapabilities.UNRESTRICTED_CAPABILITIES;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+
+import android.net.NetworkCapabilities;
+import android.support.test.runner.AndroidJUnit4;
+import android.test.suitebuilder.annotation.SmallTest;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class NetworkCapabilitiesTest {
+    @Test
+    public void testMaybeMarkCapabilitiesRestricted() {
+        // verify EIMS is restricted
+        assertEquals((1 << NET_CAPABILITY_EIMS) & RESTRICTED_CAPABILITIES,
+                (1 << NET_CAPABILITY_EIMS));
+
+        // verify CBS is also restricted
+        assertEquals((1 << NET_CAPABILITY_CBS) & RESTRICTED_CAPABILITIES,
+                (1 << NET_CAPABILITY_CBS));
+
+        // verify default is not restricted
+        assertEquals((1 << NET_CAPABILITY_INTERNET) & RESTRICTED_CAPABILITIES, 0);
+
+        // just to see
+        assertEquals(RESTRICTED_CAPABILITIES & UNRESTRICTED_CAPABILITIES, 0);
+
+        // check that internet does not get restricted
+        NetworkCapabilities netCap = new NetworkCapabilities();
+        netCap.addCapability(NET_CAPABILITY_INTERNET);
+        netCap.maybeMarkCapabilitiesRestricted();
+        assertTrue(netCap.hasCapability(NET_CAPABILITY_NOT_RESTRICTED));
+
+        // metered-ness shouldn't matter
+        netCap = new NetworkCapabilities();
+        netCap.addCapability(NET_CAPABILITY_INTERNET);
+        netCap.addCapability(NET_CAPABILITY_NOT_METERED);
+        netCap.maybeMarkCapabilitiesRestricted();
+        assertTrue(netCap.hasCapability(NET_CAPABILITY_NOT_RESTRICTED));
+        netCap = new NetworkCapabilities();
+        netCap.addCapability(NET_CAPABILITY_INTERNET);
+        netCap.removeCapability(NET_CAPABILITY_NOT_METERED);
+        netCap.maybeMarkCapabilitiesRestricted();
+        assertTrue(netCap.hasCapability(NET_CAPABILITY_NOT_RESTRICTED));
+
+        // add EIMS - bundled with unrestricted means it's unrestricted
+        netCap = new NetworkCapabilities();
+        netCap.addCapability(NET_CAPABILITY_INTERNET);
+        netCap.addCapability(NET_CAPABILITY_EIMS);
+        netCap.addCapability(NET_CAPABILITY_NOT_METERED);
+        netCap.maybeMarkCapabilitiesRestricted();
+        assertTrue(netCap.hasCapability(NET_CAPABILITY_NOT_RESTRICTED));
+        netCap = new NetworkCapabilities();
+        netCap.addCapability(NET_CAPABILITY_INTERNET);
+        netCap.addCapability(NET_CAPABILITY_EIMS);
+        netCap.removeCapability(NET_CAPABILITY_NOT_METERED);
+        netCap.maybeMarkCapabilitiesRestricted();
+        assertTrue(netCap.hasCapability(NET_CAPABILITY_NOT_RESTRICTED));
+
+        // just a restricted cap should be restricted regardless of meteredness
+        netCap = new NetworkCapabilities();
+        netCap.addCapability(NET_CAPABILITY_EIMS);
+        netCap.addCapability(NET_CAPABILITY_NOT_METERED);
+        netCap.maybeMarkCapabilitiesRestricted();
+        assertFalse(netCap.hasCapability(NET_CAPABILITY_NOT_RESTRICTED));
+        netCap = new NetworkCapabilities();
+        netCap.addCapability(NET_CAPABILITY_EIMS);
+        netCap.removeCapability(NET_CAPABILITY_NOT_METERED);
+        netCap.maybeMarkCapabilitiesRestricted();
+        assertFalse(netCap.hasCapability(NET_CAPABILITY_NOT_RESTRICTED));
+
+        // try 2 restricted caps
+        netCap = new NetworkCapabilities();
+        netCap.addCapability(NET_CAPABILITY_CBS);
+        netCap.addCapability(NET_CAPABILITY_EIMS);
+        netCap.addCapability(NET_CAPABILITY_NOT_METERED);
+        netCap.maybeMarkCapabilitiesRestricted();
+        assertFalse(netCap.hasCapability(NET_CAPABILITY_NOT_RESTRICTED));
+        netCap = new NetworkCapabilities();
+        netCap.addCapability(NET_CAPABILITY_CBS);
+        netCap.addCapability(NET_CAPABILITY_EIMS);
+        netCap.removeCapability(NET_CAPABILITY_NOT_METERED);
+        netCap.maybeMarkCapabilitiesRestricted();
+        assertFalse(netCap.hasCapability(NET_CAPABILITY_NOT_RESTRICTED));
+    }
+
+}
diff --git a/tests/net/java/android/net/nsd/NsdManagerTest.java b/tests/net/java/android/net/nsd/NsdManagerTest.java
new file mode 100644
index 0000000..2418450
--- /dev/null
+++ b/tests/net/java/android/net/nsd/NsdManagerTest.java
@@ -0,0 +1,156 @@
+/*
+ * Copyright (C) 2017 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.nsd;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.times;
+
+import android.os.HandlerThread;
+import android.os.Handler;
+import android.os.Looper;
+import android.content.Context;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
+import android.os.Message;
+import android.os.Messenger;
+import com.android.internal.util.AsyncChannel;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class NsdManagerTest {
+
+    @Mock Context mContext;
+    @Mock INsdManager mService;
+    MockServiceHandler mServiceHandler;
+
+    long mTimeoutMs = 100; // non-final so that tests can adjust the value.
+
+    @Before
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+
+        mServiceHandler = spy(MockServiceHandler.create(mContext));
+        when(mService.getMessenger()).thenReturn(new Messenger(mServiceHandler));
+    }
+
+    @Test
+    public void testResolveService() {
+        NsdManager manager = makeManager();
+
+        NsdServiceInfo request = new NsdServiceInfo("a_name", "a_type");
+        NsdServiceInfo reply = new NsdServiceInfo("resolved_name", "resolved_type");
+        NsdManager.ResolveListener listener = mock(NsdManager.ResolveListener.class);
+
+        manager.resolveService(request, listener);
+        int key1 = verifyRequest(NsdManager.RESOLVE_SERVICE);
+        int err = 33;
+        sendResponse(NsdManager.RESOLVE_SERVICE_FAILED, err, key1, null);
+        verify(listener, timeout(mTimeoutMs).times(1)).onResolveFailed(request, err);
+
+        manager.resolveService(request, listener);
+        int key2 = verifyRequest(NsdManager.RESOLVE_SERVICE);
+        sendResponse(NsdManager.RESOLVE_SERVICE_SUCCEEDED, 0, key2, reply);
+        verify(listener, timeout(mTimeoutMs).times(1)).onServiceResolved(reply);
+    }
+
+    @Test
+    public void testParallelResolveService() {
+        NsdManager manager = makeManager();
+
+        NsdServiceInfo request = new NsdServiceInfo("a_name", "a_type");
+        NsdServiceInfo reply = new NsdServiceInfo("resolved_name", "resolved_type");
+
+        NsdManager.ResolveListener listener1 = mock(NsdManager.ResolveListener.class);
+        NsdManager.ResolveListener listener2 = mock(NsdManager.ResolveListener.class);
+
+        manager.resolveService(request, listener1);
+        int key1 = verifyRequest(NsdManager.RESOLVE_SERVICE);
+
+        manager.resolveService(request, listener2);
+        int key2 = verifyRequest(NsdManager.RESOLVE_SERVICE);
+
+        sendResponse(NsdManager.RESOLVE_SERVICE_SUCCEEDED, 0, key2, reply);
+        sendResponse(NsdManager.RESOLVE_SERVICE_SUCCEEDED, 0, key1, reply);
+
+        verify(listener1, timeout(mTimeoutMs).times(1)).onServiceResolved(reply);
+        verify(listener2, timeout(mTimeoutMs).times(1)).onServiceResolved(reply);
+   }
+
+    NsdManager makeManager() {
+        NsdManager manager = new NsdManager(mContext, mService);
+        // Acknowledge first two messages connecting the AsyncChannel.
+        verify(mServiceHandler, timeout(mTimeoutMs).times(2)).handleMessage(any());
+        reset(mServiceHandler);
+        assertNotNull(mServiceHandler.chan);
+        return manager;
+    }
+
+    int verifyRequest(int expectedMessageType) {
+        verify(mServiceHandler, timeout(mTimeoutMs)).handleMessage(any());
+        reset(mServiceHandler);
+        Message received = mServiceHandler.lastMessage;
+        assertEquals(NsdManager.nameOf(expectedMessageType), NsdManager.nameOf(received.what));
+        return received.arg2;
+    }
+
+    void sendResponse(int replyType, int arg, int key, Object obj) {
+        mServiceHandler.chan.sendMessage(replyType, arg, key, obj);
+    }
+
+    // Implements the server side of AsyncChannel connection protocol
+    public static class MockServiceHandler extends Handler {
+        public Context mContext;
+        public AsyncChannel chan;
+        public Message lastMessage;
+
+        MockServiceHandler(Looper looper, Context context) {
+            super(looper);
+            mContext = context;
+        }
+
+        @Override
+        public void handleMessage(Message msg) {
+            lastMessage = obtainMessage();
+            lastMessage.copyFrom(msg);
+            if (msg.what == AsyncChannel.CMD_CHANNEL_FULL_CONNECTION) {
+                chan = new AsyncChannel();
+                chan.connect(mContext, this, msg.replyTo);
+                chan.sendMessage(AsyncChannel.CMD_CHANNEL_FULLY_CONNECTED);
+            }
+
+        }
+
+        public static MockServiceHandler create(Context context) {
+            HandlerThread t = new HandlerThread("mock-service-handler");
+            t.start();
+            return new MockServiceHandler(t.getLooper(), context);
+        }
+    }
+}
diff --git a/tests/net/java/com/android/internal/util/BitUtilsTest.java b/tests/net/java/com/android/internal/util/BitUtilsTest.java
new file mode 100644
index 0000000..0ad8a21
--- /dev/null
+++ b/tests/net/java/com/android/internal/util/BitUtilsTest.java
@@ -0,0 +1,107 @@
+/*
+ * Copyright (C) 2017 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 com.android.internal.util;
+
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
+import java.nio.ByteBuffer;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertEquals;
+import static com.android.internal.util.BitUtils.*;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class BitUtilsTest {
+
+    @Test
+    public void testUnsignedByteWideningConversions() {
+        byte b0 = 0;
+        byte b1 = 1;
+        byte bm1 = -1;
+        assertEquals(0, uint8(b0));
+        assertEquals(1, uint8(b1));
+        assertEquals(127, uint8(Byte.MAX_VALUE));
+        assertEquals(128, uint8(Byte.MIN_VALUE));
+        assertEquals(255, uint8(bm1));
+        assertEquals(255, uint8((byte)255));
+    }
+
+    @Test
+    public void testUnsignedShortWideningConversions() {
+        short s0 = 0;
+        short s1 = 1;
+        short sm1 = -1;
+        assertEquals(0, uint16(s0));
+        assertEquals(1, uint16(s1));
+        assertEquals(32767, uint16(Short.MAX_VALUE));
+        assertEquals(32768, uint16(Short.MIN_VALUE));
+        assertEquals(65535, uint16(sm1));
+        assertEquals(65535, uint16((short)65535));
+    }
+
+    @Test
+    public void testUnsignedIntWideningConversions() {
+        assertEquals(0, uint32(0));
+        assertEquals(1, uint32(1));
+        assertEquals(2147483647L, uint32(Integer.MAX_VALUE));
+        assertEquals(2147483648L, uint32(Integer.MIN_VALUE));
+        assertEquals(4294967295L, uint32(-1));
+        assertEquals(4294967295L, uint32((int)4294967295L));
+    }
+
+    @Test
+    public void testBytesToInt() {
+        assertEquals(0x00000000, bytesToBEInt(bytes(0, 0, 0, 0)));
+        assertEquals(0xffffffff, bytesToBEInt(bytes(255, 255, 255, 255)));
+        assertEquals(0x0a000001, bytesToBEInt(bytes(10, 0, 0, 1)));
+        assertEquals(0x0a000002, bytesToBEInt(bytes(10, 0, 0, 2)));
+        assertEquals(0x0a001fff, bytesToBEInt(bytes(10, 0, 31, 255)));
+        assertEquals(0xe0000001, bytesToBEInt(bytes(224, 0, 0, 1)));
+
+        assertEquals(0x00000000, bytesToLEInt(bytes(0, 0, 0, 0)));
+        assertEquals(0x01020304, bytesToLEInt(bytes(4, 3, 2, 1)));
+        assertEquals(0xffff0000, bytesToLEInt(bytes(0, 0, 255, 255)));
+    }
+
+    @Test
+    public void testUnsignedGetters() {
+      ByteBuffer b = ByteBuffer.allocate(4);
+      b.putInt(0xffff);
+
+      assertEquals(0x0, getUint8(b, 0));
+      assertEquals(0x0, getUint8(b, 1));
+      assertEquals(0xff, getUint8(b, 2));
+      assertEquals(0xff, getUint8(b, 3));
+
+      assertEquals(0x0, getUint16(b, 0));
+      assertEquals(0xffff, getUint16(b, 2));
+
+      b.rewind();
+      b.putInt(0xffffffff);
+      assertEquals(0xffffffffL, getUint32(b, 0));
+    }
+
+    static byte[] bytes(int b1, int b2, int b3, int b4) {
+        return new byte[] {b(b1), b(b2), b(b3), b(b4)};
+    }
+
+    static byte b(int i) {
+        return (byte) i;
+    }
+}
diff --git a/tests/net/java/com/android/server/connectivity/VpnTest.java b/tests/net/java/com/android/server/connectivity/VpnTest.java
index efe6fec..506d9e5 100644
--- a/tests/net/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/net/java/com/android/server/connectivity/VpnTest.java
@@ -27,6 +27,7 @@
 import android.app.AppOpsManager;
 import android.app.NotificationManager;
 import android.content.Context;
+import android.content.Intent;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.UserInfo;
@@ -42,6 +43,8 @@
 import android.util.ArrayMap;
 import android.util.ArraySet;
 
+import com.android.internal.net.VpnConfig;
+
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Map;
@@ -101,8 +104,10 @@
     @Override
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
+
         when(mContext.getPackageManager()).thenReturn(mPackageManager);
         setMockedPackages(mPackages);
+
         when(mContext.getPackageName()).thenReturn(Vpn.class.getPackage().getName());
         when(mContext.getSystemService(eq(Context.USER_SERVICE))).thenReturn(mUserManager);
         when(mContext.getSystemService(eq(Context.APP_OPS_SERVICE))).thenReturn(mAppOps);
@@ -258,6 +263,58 @@
     }
 
     @SmallTest
+    public void testLockdownRuleRepeatability() throws Exception {
+        final Vpn vpn = createVpn(primaryUser.id);
+
+        // Given legacy lockdown is already enabled,
+        vpn.setLockdown(true);
+        verify(mNetService, times(1)).setAllowOnlyVpnForUids(
+                eq(true), aryEq(new UidRange[] {UidRange.createForUser(primaryUser.id)}));
+
+        // Enabling legacy lockdown twice should do nothing.
+        vpn.setLockdown(true);
+        verify(mNetService, times(1)).setAllowOnlyVpnForUids(anyBoolean(), any(UidRange[].class));
+
+        // And disabling should remove the rules exactly once.
+        vpn.setLockdown(false);
+        verify(mNetService, times(1)).setAllowOnlyVpnForUids(
+                eq(false), aryEq(new UidRange[] {UidRange.createForUser(primaryUser.id)}));
+
+        // Removing the lockdown again should have no effect.
+        vpn.setLockdown(false);
+        verify(mNetService, times(2)).setAllowOnlyVpnForUids(anyBoolean(), any(UidRange[].class));
+    }
+
+    @SmallTest
+    public void testLockdownRuleReversibility() throws Exception {
+        final Vpn vpn = createVpn(primaryUser.id);
+
+        final UidRange[] entireUser = {
+            UidRange.createForUser(primaryUser.id)
+        };
+        final UidRange[] exceptPkg0 = {
+            new UidRange(entireUser[0].start, entireUser[0].start + PKG_UIDS[0] - 1),
+            new UidRange(entireUser[0].start + PKG_UIDS[0] + 1, entireUser[0].stop)
+        };
+
+        final InOrder order = inOrder(mNetService);
+
+        // Given lockdown is enabled with no package (legacy VPN),
+        vpn.setLockdown(true);
+        order.verify(mNetService).setAllowOnlyVpnForUids(eq(true), aryEq(entireUser));
+
+        // When a new VPN package is set the rules should change to cover that package.
+        vpn.prepare(null, PKGS[0]);
+        order.verify(mNetService).setAllowOnlyVpnForUids(eq(false), aryEq(entireUser));
+        order.verify(mNetService).setAllowOnlyVpnForUids(eq(true), aryEq(exceptPkg0));
+
+        // When that VPN package is unset, everything should be undone again in reverse.
+        vpn.prepare(null, VpnConfig.LEGACY_VPN);
+        order.verify(mNetService).setAllowOnlyVpnForUids(eq(false), aryEq(exceptPkg0));
+        order.verify(mNetService).setAllowOnlyVpnForUids(eq(true), aryEq(entireUser));
+    }
+
+    @SmallTest
     public void testNotificationShownForAlwaysOnApp() {
         final UserHandle userHandle = UserHandle.of(primaryUser.id);
         final Vpn vpn = createVpn(primaryUser.id);