Move socket tagging from libcore.

Change-Id: I7515896936c5bbd212c88e2801b831914219a925

5 files changed, 184 insertions, 17 deletions

diff --git a/core/java/android/net/TrafficStats.java b/core/java/android/net/TrafficStats.java
index 040489e664f4..2b59dba15c0e 100644
--- a/core/java/android/net/TrafficStats.java
+++ b/core/java/android/net/TrafficStats.java
@@ -25,8 +25,9 @@ import android.os.INetworkManagementService;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 
-import dalvik.system.BlockGuard;
+import com.android.server.NetworkManagementSocketTagger;
 
+import dalvik.system.SocketTagger;
 import java.net.Socket;
 import java.net.SocketException;
 
@@ -92,7 +93,7 @@ public class TrafficStats {
      * {@link #tagSocket(Socket)}.
      */
     public static void setThreadStatsTag(int tag) {
-        BlockGuard.setThreadSocketStatsTag(tag);
+        NetworkManagementSocketTagger.setThreadSocketStatsTag(tag);
     }
 
     /**
@@ -104,7 +105,7 @@ public class TrafficStats {
     }
 
     public static void clearThreadStatsTag() {
-        BlockGuard.setThreadSocketStatsTag(-1);
+        NetworkManagementSocketTagger.setThreadSocketStatsTag(-1);
     }
 
     /**
@@ -121,12 +122,12 @@ public class TrafficStats {
      * {@hide}
      */
     public static void setThreadStatsUid(int uid) {
-        BlockGuard.setThreadSocketStatsUid(uid);
+        NetworkManagementSocketTagger.setThreadSocketStatsUid(uid);
     }
 
     /** {@hide} */
     public static void clearThreadStatsUid() {
-        BlockGuard.setThreadSocketStatsUid(-1);
+        NetworkManagementSocketTagger.setThreadSocketStatsUid(-1);
     }
 
     /**
@@ -139,14 +140,14 @@ public class TrafficStats {
      * @see #setThreadStatsUid(int)
      */
     public static void tagSocket(Socket socket) throws SocketException {
-        BlockGuard.tagSocketFd(socket.getFileDescriptor$());
+        SocketTagger.get().tag(socket);
     }
 
     /**
      * Remove any statistics parameters from the given {@link Socket}.
      */
     public static void untagSocket(Socket socket) throws SocketException {
-        BlockGuard.untagSocketFd(socket.getFileDescriptor$());
+        SocketTagger.get().untag(socket);
     }
 
     /**
diff --git a/core/java/com/android/internal/os/RuntimeInit.java b/core/java/com/android/internal/os/RuntimeInit.java
index 5e9cd23a82cd..f13e77067d5b 100644
--- a/core/java/com/android/internal/os/RuntimeInit.java
+++ b/core/java/com/android/internal/os/RuntimeInit.java
@@ -25,16 +25,13 @@ import android.os.Process;
 import android.os.SystemProperties;
 import android.util.Log;
 import android.util.Slog;
-
 import com.android.internal.logging.AndroidConfig;
-
+import com.android.server.NetworkManagementSocketTagger;
 import dalvik.system.VMRuntime;
-
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
-import java.util.logging.LogManager;
 import java.util.TimeZone;
-
+import java.util.logging.LogManager;
 import org.apache.harmony.luni.internal.util.TimezoneGetter;
 
 /**
@@ -129,6 +126,11 @@ public class RuntimeInit {
         System.setProperty("http.agent", userAgent);
 
         /*
+         * Wire socket tagging to traffic stats.
+         */
+        NetworkManagementSocketTagger.install();
+
+        /*
          * If we're running in an emulator launched with "-trace", put the
          * VM into emulator trace profiling mode so that the user can hit
          * F9/F10 at any time to capture traces.  This has performance
diff --git a/core/java/com/android/server/NetworkManagementSocketTagger.java b/core/java/com/android/server/NetworkManagementSocketTagger.java
new file mode 100644
index 000000000000..306d2236693a
--- /dev/null
+++ b/core/java/com/android/server/NetworkManagementSocketTagger.java
@@ -0,0 +1,165 @@
+/*
+ * Copyright (C) 2011 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.server;
+
+import dalvik.system.SocketTagger;
+import java.io.FileDescriptor;
+import java.io.FileNotFoundException;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.net.SocketException;
+import java.nio.charset.Charsets;
+
+/**
+ * Assigns tags to sockets for traffic stats.
+ */
+public final class NetworkManagementSocketTagger extends SocketTagger {
+
+    private static final boolean LOGI = false;
+
+    private static ThreadLocal<SocketTags> threadSocketTags = new ThreadLocal<SocketTags>() {
+        @Override protected SocketTags initialValue() {
+            return new SocketTags();
+        }
+    };
+
+    public static void install() {
+        SocketTagger.set(new NetworkManagementSocketTagger());
+    }
+
+    public static void setThreadSocketStatsTag(int tag) {
+        threadSocketTags.get().statsTag = tag;
+    }
+
+    public static void setThreadSocketStatsUid(int uid) {
+        threadSocketTags.get().statsUid = uid;
+    }
+
+    @Override public void tag(FileDescriptor fd) throws SocketException {
+        final SocketTags options = threadSocketTags.get();
+        if (LOGI) {
+            System.logI("tagSocket(" + fd.getInt$() + ") with statsTag="
+                    + options.statsTag + ", statsUid=" + options.statsUid);
+        }
+        try {
+            // TODO: skip tagging when options would be no-op
+            tagSocketFd(fd, options.statsTag, options.statsUid);
+        } catch (IOException e) {
+            throw new SocketException("Problem tagging socket", e);
+        }
+    }
+
+    private void tagSocketFd(FileDescriptor fd, int tag, int uid) throws IOException {
+        final int fdNum = fd.getInt$();
+        if (fdNum == -1 || (tag == -1 && uid == -1)) return;
+
+        String cmd = "t " + fdNum;
+        if (tag == -1) {
+            // Case where just the uid needs adjusting. But probably the caller
+            // will want to track his own name here, just in case.
+            cmd += " 0";
+        } else {
+            cmd += " " + tagToKernel(tag);
+        }
+        if (uid != -1) {
+            cmd += " " + uid;
+        }
+        internalModuleCtrl(cmd);
+    }
+
+    @Override public void untag(FileDescriptor fd) throws SocketException {
+        if (LOGI) {
+            System.logI("untagSocket(" + fd.getInt$() + ")");
+        }
+        try {
+            unTagSocketFd(fd);
+        } catch (IOException e) {
+            throw new SocketException("Problem untagging socket", e);
+        }
+    }
+
+    private void unTagSocketFd(FileDescriptor fd) throws IOException {
+        int fdNum = fd.getInt$();
+        if (fdNum == -1) return;
+        String cmd = "u " + fdNum;
+        internalModuleCtrl(cmd);
+    }
+
+    public static class SocketTags {
+        public int statsTag = -1;
+        public int statsUid = -1;
+    }
+
+    /**
+     * Sends commands to the kernel netfilter module.
+     *
+     * @param cmd command string for the qtaguid netfilter module. May not be null.
+     *   <p>Supports:
+     *     <ul><li>tag a socket:<br>
+     *        <code>t <i>sock_fd</i> <i>acct_tag</i> [<i>uid_in_case_caller_is_acting_on_behalf_of</i>]</code><br>
+     *     <code>*_tag</code> defaults to default_policy_tag_from_uid(uid_of_caller)<br>
+     *     <code>acct_tag</code> is either 0 or greater that 2^32.<br>
+     *     <code>uid_*</code> is only settable by privileged UIDs (DownloadManager,...)
+     *     </li>
+     *     <li>untag a socket, preserving counters:<br>
+     *       <code>u <i>sock_fd</i></code>
+     *     </li></ul>
+     *   <p>Notes:<br>
+     *   <ul><li><i>sock_fd</i> is withing the callers process space.</li>
+     *   <li><i>*_tag</i> are 64bit values</li></ul>
+     *
+     */
+    private void internalModuleCtrl(String cmd) throws IOException {
+        final FileOutputStream procOut;
+        // TODO: Use something like
+        //  android.os.SystemProperties.getInt("persist.bandwidth.enable", 0)
+        // to see if tagging should happen or not.
+        try {
+            procOut = new FileOutputStream("/proc/net/xt_qtaguid/ctrl");
+        } catch (FileNotFoundException e) {
+            if (LOGI) {
+                System.logI("Can't talk to kernel module:" + e);
+            }
+            return;
+        }
+        try {
+            procOut.write(cmd.getBytes(Charsets.US_ASCII));
+        } finally {
+            procOut.close();
+        }
+    }
+
+    /**
+     * Convert {@link Integer} tag to {@code /proc/} format. Assumes unsigned
+     * base-10 format like {@code 2147483647}. Currently strips signed bit to
+     * avoid using {@link java.math.BigInteger}.
+     */
+    public static String tagToKernel(int tag) {
+        // TODO: eventually write in hex, since that's what proc exports
+        // TODO: migrate to direct integer instead of odd shifting
+        return Long.toString((((long) tag) << 32) & 0x7FFFFFFF00000000L);
+    }
+
+    /**
+     * Convert {@code /proc/} tag format to {@link Integer}. Assumes incoming
+     * format like {@code 0x7fffffff00000000}.
+     */
+    public static int kernelToTag(String string) {
+        // TODO: migrate to direct integer instead of odd shifting
+        return (int) (Long.decode(string) >> 32);
+    }
+}
diff --git a/services/java/com/android/server/NetworkManagementService.java b/services/java/com/android/server/NetworkManagementService.java
index 630aaf9edafe..b31d128037f8 100644
--- a/services/java/com/android/server/NetworkManagementService.java
+++ b/services/java/com/android/server/NetworkManagementService.java
@@ -41,8 +41,6 @@ import android.util.Slog;
 import com.google.android.collect.Lists;
 import com.google.android.collect.Maps;
 
-import dalvik.system.BlockGuard;
-
 import java.io.BufferedReader;
 import java.io.DataInputStream;
 import java.io.File;
@@ -958,7 +956,8 @@ class NetworkManagementService extends INetworkManagementService.Stub {
 
                 try {
                     final String iface = parsed.get(KEY_IFACE);
-                    final int tag = BlockGuard.kernelToTag(parsed.get(KEY_TAG_HEX));
+                    final int tag = NetworkManagementSocketTagger.kernelToTag(
+                            parsed.get(KEY_TAG_HEX));
                     final int uid = Integer.parseInt(parsed.get(KEY_UID));
                     final long rx = Long.parseLong(parsed.get(KEY_RX));
                     final long tx = Long.parseLong(parsed.get(KEY_TX));
diff --git a/services/tests/servicestests/src/com/android/server/NetworkManagementServiceTest.java b/services/tests/servicestests/src/com/android/server/NetworkManagementServiceTest.java
index 8b752ee47468..ac7cb5afd53a 100644
--- a/services/tests/servicestests/src/com/android/server/NetworkManagementServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/NetworkManagementServiceTest.java
@@ -16,8 +16,8 @@
 
 package com.android.server;
 
-import static dalvik.system.BlockGuard.kernelToTag;
-import static dalvik.system.BlockGuard.tagToKernel;
+import static com.android.server.NetworkManagementSocketTagger.kernelToTag;
+import static com.android.server.NetworkManagementSocketTagger.tagToKernel;
 
 import android.content.res.Resources;
 import android.net.NetworkStats;