Close sockets when changing network permissions.

Bug: 23113288
Change-Id: I8dcb02c79c81244e5b7288cb50770ac6a5867fcc
diff --git a/server/SockDiagTest.cpp b/server/SockDiagTest.cpp
index f9353f3..70a199d 100644
--- a/server/SockDiagTest.cpp
+++ b/server/SockDiagTest.cpp
@@ -25,6 +25,7 @@
 
 #include <gtest/gtest.h>
 
+#include "Fwmark.h"
 #include "NetdConstants.h"
 #include "SockDiag.h"
 #include "UidRanges.h"
@@ -278,6 +279,7 @@
     UID_EXCLUDE_LOOPBACK,
     UIDRANGE,
     UIDRANGE_EXCLUDE_LOOPBACK,
+    PERMISSION,
 };
 
 const char *testTypeName(MicroBenchmarkTestType mode) {
@@ -288,10 +290,30 @@
         TO_STRING_TYPE(UID_EXCLUDE_LOOPBACK);
         TO_STRING_TYPE(UIDRANGE);
         TO_STRING_TYPE(UIDRANGE_EXCLUDE_LOOPBACK);
+        TO_STRING_TYPE(PERMISSION);
     }
 #undef TO_STRING_TYPE
 }
 
+static struct {
+    unsigned netId;
+    bool explicitlySelected;
+    Permission permission;
+} permissionTestcases[] = {
+    { 42, false, PERMISSION_NONE,    },
+    { 42, false, PERMISSION_NETWORK, },
+    { 42, false, PERMISSION_SYSTEM,  },
+    { 42, true,  PERMISSION_NONE,    },
+    { 42, true,  PERMISSION_NETWORK, },
+    { 42, true,  PERMISSION_SYSTEM,  },
+    { 43, false, PERMISSION_NONE,    },
+    { 43, false, PERMISSION_NETWORK, },
+    { 43, false, PERMISSION_SYSTEM,  },
+    { 43, true,  PERMISSION_NONE,    },
+    { 43, true,  PERMISSION_NETWORK, },
+    { 43, true,  PERMISSION_SYSTEM,  },
+};
+
 class SockDiagMicroBenchmarkTest : public ::testing::TestWithParam<MicroBenchmarkTestType> {
 
 public:
@@ -305,10 +327,15 @@
     constexpr static int MAX_SOCKETS = 500;
     constexpr static int ADDRESS_SOCKETS = 500;
     constexpr static int UID_SOCKETS = 50;
+    constexpr static int PERMISSION_SOCKETS = 16;
+
     constexpr static uid_t START_UID = 8000;  // START_UID + number of sockets must be <= 9999.
     constexpr static int CLOSE_UID = START_UID + UID_SOCKETS - 42;  // Close to the end
     static_assert(START_UID + MAX_SOCKETS < 9999, "Too many sockets");
 
+    constexpr static int TEST_NETID = 42;  // One of the OEM netIds.
+
+
     int howManySockets() {
         MicroBenchmarkTestType mode = GetParam();
         switch (mode) {
@@ -319,6 +346,30 @@
         case UIDRANGE:
         case UIDRANGE_EXCLUDE_LOOPBACK:
             return UID_SOCKETS;
+        case PERMISSION:
+            return ARRAY_SIZE(permissionTestcases);
+        }
+    }
+
+    int modifySocketForTest(int s, int i) {
+        MicroBenchmarkTestType mode = GetParam();
+        switch (mode) {
+        case UID:
+        case UID_EXCLUDE_LOOPBACK:
+        case UIDRANGE:
+        case UIDRANGE_EXCLUDE_LOOPBACK: {
+            uid_t uid = START_UID + i;
+            return fchown(s, uid, -1);
+        }
+        case PERMISSION: {
+            Fwmark fwmark;
+            fwmark.netId = permissionTestcases[i].netId;
+            fwmark.explicitlySelected = permissionTestcases[i].explicitlySelected;
+            fwmark.permission = permissionTestcases[i].permission;
+            return setsockopt(s, SOL_SOCKET, SO_MARK, &fwmark.intValue, sizeof(fwmark.intValue));
+        }
+        default:
+            return 0;
         }
     }
 
@@ -346,6 +397,11 @@
                 UidRanges uidRanges;
                 uidRanges.parseFrom(ARRAY_SIZE(uidRangeStrings), (char **) uidRangeStrings);
                 ret = mSd.destroySockets(uidRanges, skipUids, excludeLoopback);
+                break;
+            }
+            case PERMISSION: {
+                ret = mSd.destroySocketsLackingPermission(TEST_NETID, PERMISSION_NETWORK, false);
+                break;
             }
         }
         return ret;
@@ -373,6 +429,11 @@
             case UID_EXCLUDE_LOOPBACK:
             case UIDRANGE_EXCLUDE_LOOPBACK:
                 return false;
+            case PERMISSION:
+                if (permissionTestcases[i].netId != 42) return false;
+                if (permissionTestcases[i].explicitlySelected != 1) return true;
+                Permission permission = permissionTestcases[i].permission;
+                return permission != PERMISSION_NETWORK && permission != PERMISSION_SYSTEM;
         }
     }
 
@@ -424,11 +485,10 @@
     auto start = std::chrono::steady_clock::now();
     for (int i = 0; i < numSockets; i++) {
         int s = socket(AF_INET6, SOCK_STREAM, 0);
-        uid_t uid = START_UID + i;
-        ASSERT_EQ(0, fchown(s, uid, -1));
         clientlen = sizeof(client);
         ASSERT_EQ(0, connect(s, (sockaddr *) &server, sizeof(server)))
             << "Connecting socket " << i << " failed " << strerror(errno);
+        ASSERT_EQ(0, modifySocketForTest(s, i));
         serversockets[i] = accept(listensocket, (sockaddr *) &client, &clientlen);
         ASSERT_NE(-1, serversockets[i])
             << "Accepting socket " << i << " failed " << strerror(errno);
@@ -472,4 +532,5 @@
 
 INSTANTIATE_TEST_CASE_P(Address, SockDiagMicroBenchmarkTest,
                         testing::Values(ADDRESS, UID, UIDRANGE,
-                                        UID_EXCLUDE_LOOPBACK, UIDRANGE_EXCLUDE_LOOPBACK));
+                                        UID_EXCLUDE_LOOPBACK, UIDRANGE_EXCLUDE_LOOPBACK,
+                                        PERMISSION));