Merge "Add reference counted resources to IpSecService"
diff --git a/core/java/android/app/usage/NetworkStats.java b/core/java/android/app/usage/NetworkStats.java
index 3670b91..222e9a0 100644
--- a/core/java/android/app/usage/NetworkStats.java
+++ b/core/java/android/app/usage/NetworkStats.java
@@ -97,12 +97,12 @@
     private NetworkStatsHistory.Entry mRecycledHistoryEntry = null;
 
     /** @hide */
-    NetworkStats(Context context, NetworkTemplate template, long startTimestamp,
+    NetworkStats(Context context, NetworkTemplate template, int flags, long startTimestamp,
             long endTimestamp) throws RemoteException, SecurityException {
         final INetworkStatsService statsService = INetworkStatsService.Stub.asInterface(
                 ServiceManager.getService(Context.NETWORK_STATS_SERVICE));
         // Open network stats session
-        mSession = statsService.openSessionForUsageStats(context.getOpPackageName());
+        mSession = statsService.openSessionForUsageStats(flags, context.getOpPackageName());
         mCloseGuard.open("close");
         mTemplate = template;
         mStartTimeStamp = startTimestamp;
diff --git a/core/java/android/app/usage/NetworkStatsManager.java b/core/java/android/app/usage/NetworkStatsManager.java
index ef262e0..853b003 100644
--- a/core/java/android/app/usage/NetworkStatsManager.java
+++ b/core/java/android/app/usage/NetworkStatsManager.java
@@ -24,15 +24,14 @@
 import android.content.Context;
 import android.net.ConnectivityManager;
 import android.net.DataUsageRequest;
+import android.net.INetworkStatsService;
 import android.net.NetworkIdentity;
 import android.net.NetworkTemplate;
-import android.net.INetworkStatsService;
 import android.os.Binder;
-import android.os.Build;
-import android.os.Message;
-import android.os.Messenger;
 import android.os.Handler;
 import android.os.Looper;
+import android.os.Message;
+import android.os.Messenger;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.os.ServiceManager.ServiceNotFoundException;
@@ -79,7 +78,7 @@
  * In addition to tethering usage, usage by removed users and apps, and usage by the system
  * is also included in the results for callers with one of these higher levels of access.
  * <p />
- * <b>NOTE:</b> Prior to API level {@value Build.VERSION_CODES#N}, all calls to these APIs required
+ * <b>NOTE:</b> Prior to API level {@value android.os.Build.VERSION_CODES#N}, all calls to these APIs required
  * the above permission, even to access an app's own data usage, and carrier-privileged apps were
  * not included.
  */
@@ -96,6 +95,13 @@
     private final Context mContext;
     private final INetworkStatsService mService;
 
+    /** @hide */
+    public static final int FLAG_POLL_ON_OPEN = 1 << 0;
+    /** @hide */
+    public static final int FLAG_AUGMENT_WITH_SUBSCRIPTION_PLAN = 1 << 1;
+
+    private int mFlags;
+
     /**
      * {@hide}
      */
@@ -103,6 +109,25 @@
         mContext = context;
         mService = INetworkStatsService.Stub.asInterface(
                 ServiceManager.getServiceOrThrow(Context.NETWORK_STATS_SERVICE));
+        setPollOnOpen(true);
+    }
+
+    /** @hide */
+    public void setPollOnOpen(boolean pollOnOpen) {
+        if (pollOnOpen) {
+            mFlags |= FLAG_POLL_ON_OPEN;
+        } else {
+            mFlags &= ~FLAG_POLL_ON_OPEN;
+        }
+    }
+
+    /** @hide */
+    public void setAugmentWithSubscriptionPlan(boolean augmentWithSubscriptionPlan) {
+        if (augmentWithSubscriptionPlan) {
+            mFlags |= FLAG_AUGMENT_WITH_SUBSCRIPTION_PLAN;
+        } else {
+            mFlags &= ~FLAG_AUGMENT_WITH_SUBSCRIPTION_PLAN;
+        }
     }
 
     /**
@@ -136,7 +161,7 @@
         }
 
         Bucket bucket = null;
-        NetworkStats stats = new NetworkStats(mContext, template, startTime, endTime);
+        NetworkStats stats = new NetworkStats(mContext, template, mFlags, startTime, endTime);
         bucket = stats.getDeviceSummaryForNetwork();
 
         stats.close();
@@ -174,7 +199,7 @@
         }
 
         NetworkStats stats;
-        stats = new NetworkStats(mContext, template, startTime, endTime);
+        stats = new NetworkStats(mContext, template, mFlags, startTime, endTime);
         stats.startSummaryEnumeration();
 
         stats.close();
@@ -211,7 +236,7 @@
         }
 
         NetworkStats result;
-        result = new NetworkStats(mContext, template, startTime, endTime);
+        result = new NetworkStats(mContext, template, mFlags, startTime, endTime);
         result.startSummaryEnumeration();
 
         return result;
@@ -260,7 +285,7 @@
 
         NetworkStats result;
         try {
-            result = new NetworkStats(mContext, template, startTime, endTime);
+            result = new NetworkStats(mContext, template, mFlags, startTime, endTime);
             result.startHistoryEnumeration(uid, tag);
         } catch (RemoteException e) {
             Log.e(TAG, "Error while querying stats for uid=" + uid + " tag=" + tag, e);
@@ -305,7 +330,7 @@
         }
 
         NetworkStats result;
-        result = new NetworkStats(mContext, template, startTime, endTime);
+        result = new NetworkStats(mContext, template, mFlags, startTime, endTime);
         result.startUserUidEnumeration();
         return result;
     }
diff --git a/core/java/android/net/INetworkStatsService.aidl b/core/java/android/net/INetworkStatsService.aidl
index e693009..9180112 100644
--- a/core/java/android/net/INetworkStatsService.aidl
+++ b/core/java/android/net/INetworkStatsService.aidl
@@ -36,7 +36,7 @@
      *  PACKAGE_USAGE_STATS permission is always checked. If PACKAGE_USAGE_STATS is not granted
      *  READ_NETWORK_USAGE_STATS is checked for.
      */
-    INetworkStatsSession openSessionForUsageStats(String callingPackage);
+    INetworkStatsSession openSessionForUsageStats(int flags, String callingPackage);
 
     /** Return network layer usage total for traffic that matches template. */
     long getNetworkTotalBytes(in NetworkTemplate template, long start, long end);
diff --git a/core/java/android/net/IpSecAlgorithm.java b/core/java/android/net/IpSecAlgorithm.java
index 64f8f39..f82627b 100644
--- a/core/java/android/net/IpSecAlgorithm.java
+++ b/core/java/android/net/IpSecAlgorithm.java
@@ -15,11 +15,13 @@
  */
 package android.net;
 
+import android.annotation.NonNull;
 import android.annotation.StringDef;
 import android.os.Build;
 import android.os.Parcel;
 import android.os.Parcelable;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.HexDump;
 
 import java.lang.annotation.Retention;
@@ -27,10 +29,14 @@
 import java.util.Arrays;
 
 /**
- * IpSecAlgorithm specifies a single algorithm that can be applied to an IpSec Transform. Refer to
- * RFC 4301.
+ * This class represents a single algorithm that can be used by an {@link IpSecTransform}.
+ *
+ * @see <a href="https://tools.ietf.org/html/rfc4301">RFC 4301, Security Architecture for the
+ * Internet Protocol</a>
  */
 public final class IpSecAlgorithm implements Parcelable {
+    private static final String TAG = "IpSecAlgorithm";
+
     /**
      * AES-CBC Encryption/Ciphering Algorithm.
      *
@@ -39,17 +45,19 @@
     public static final String CRYPT_AES_CBC = "cbc(aes)";
 
     /**
-     * MD5 HMAC Authentication/Integrity Algorithm. This algorithm is not recommended for use in new
-     * applications and is provided for legacy compatibility with 3gpp infrastructure.
+     * MD5 HMAC Authentication/Integrity Algorithm. <b>This algorithm is not recommended for use in
+     * new applications and is provided for legacy compatibility with 3gpp infrastructure.</b>
      *
+     * <p>Keys for this algorithm must be 128 bits in length.
      * <p>Valid truncation lengths are multiples of 8 bits from 96 to (default) 128.
      */
     public static final String AUTH_HMAC_MD5 = "hmac(md5)";
 
     /**
-     * SHA1 HMAC Authentication/Integrity Algorithm. This algorithm is not recommended for use in
-     * new applications and is provided for legacy compatibility with 3gpp infrastructure.
+     * SHA1 HMAC Authentication/Integrity Algorithm. <b>This algorithm is not recommended for use in
+     * new applications and is provided for legacy compatibility with 3gpp infrastructure.</b>
      *
+     * <p>Keys for this algorithm must be 160 bits in length.
      * <p>Valid truncation lengths are multiples of 8 bits from 96 to (default) 160.
      */
     public static final String AUTH_HMAC_SHA1 = "hmac(sha1)";
@@ -57,6 +65,7 @@
     /**
      * SHA256 HMAC Authentication/Integrity Algorithm.
      *
+     * <p>Keys for this algorithm must be 256 bits in length.
      * <p>Valid truncation lengths are multiples of 8 bits from 96 to (default) 256.
      */
     public static final String AUTH_HMAC_SHA256 = "hmac(sha256)";
@@ -64,13 +73,15 @@
     /**
      * SHA384 HMAC Authentication/Integrity Algorithm.
      *
+     * <p>Keys for this algorithm must be 384 bits in length.
      * <p>Valid truncation lengths are multiples of 8 bits from 192 to (default) 384.
      */
     public static final String AUTH_HMAC_SHA384 = "hmac(sha384)";
 
     /**
-     * SHA512 HMAC Authentication/Integrity Algorithm
+     * SHA512 HMAC Authentication/Integrity Algorithm.
      *
+     * <p>Keys for this algorithm must be 512 bits in length.
      * <p>Valid truncation lengths are multiples of 8 bits from 256 to (default) 512.
      */
     public static final String AUTH_HMAC_SHA512 = "hmac(sha512)";
@@ -80,9 +91,9 @@
      *
      * <p>Valid lengths for keying material are {160, 224, 288}.
      *
-     * <p>As per RFC4106 (Section 8.1), keying material consists of a 128, 192, or 256 bit AES key
-     * followed by a 32-bit salt. RFC compliance requires that the salt must be unique per
-     * invocation with the same key.
+     * <p>As per <a href="https://tools.ietf.org/html/rfc4106#section-8.1">RFC4106 (Section
+     * 8.1)</a>, keying material consists of a 128, 192, or 256 bit AES key followed by a 32-bit
+     * salt. RFC compliance requires that the salt must be unique per invocation with the same key.
      *
      * <p>Valid ICV (truncation) lengths are {64, 96, 128}.
      */
@@ -105,48 +116,45 @@
     private final int mTruncLenBits;
 
     /**
-     * Specify a IpSecAlgorithm of one of the supported types including the truncation length of the
-     * algorithm
+     * Creates an IpSecAlgorithm of one of the supported types. Supported algorithm names are
+     * defined as constants in this class.
      *
-     * @param algorithm type for IpSec.
-     * @param key non-null Key padded to a multiple of 8 bits.
+     * @param algorithm name of the algorithm.
+     * @param key key padded to a multiple of 8 bits.
      */
-    public IpSecAlgorithm(String algorithm, byte[] key) {
+    public IpSecAlgorithm(@AlgorithmName String algorithm, @NonNull byte[] key) {
         this(algorithm, key, key.length * 8);
     }
 
     /**
-     * Specify a IpSecAlgorithm of one of the supported types including the truncation length of the
-     * algorithm
+     * Creates an IpSecAlgorithm of one of the supported types. Supported algorithm names are
+     * defined as constants in this class.
      *
-     * @param algoName precise name of the algorithm to be used.
-     * @param key non-null Key padded to a multiple of 8 bits.
-     * @param truncLenBits the number of bits of output hash to use; only meaningful for
-     *     Authentication or Authenticated Encryption (equivalent to ICV length).
+     * <p>This constructor only supports algorithms that use a truncation length. i.e.
+     * Authentication and Authenticated Encryption algorithms.
+     *
+     * @param algorithm name of the algorithm.
+     * @param key key padded to a multiple of 8 bits.
+     * @param truncLenBits number of bits of output hash to use.
      */
-    public IpSecAlgorithm(@AlgorithmName String algoName, byte[] key, int truncLenBits) {
-        if (!isTruncationLengthValid(algoName, truncLenBits)) {
-            throw new IllegalArgumentException("Unknown algorithm or invalid length");
-        }
-        mName = algoName;
+    public IpSecAlgorithm(@AlgorithmName String algorithm, @NonNull byte[] key, int truncLenBits) {
+        mName = algorithm;
         mKey = key.clone();
-        mTruncLenBits = Math.min(truncLenBits, key.length * 8);
+        mTruncLenBits = truncLenBits;
+        checkValidOrThrow(mName, mKey.length * 8, mTruncLenBits);
     }
 
-    /** Retrieve the algorithm name */
+    /** Get the algorithm name */
     public String getName() {
         return mName;
     }
 
-    /** Retrieve the key for this algorithm */
+    /** Get the key for this algorithm */
     public byte[] getKey() {
         return mKey.clone();
     }
 
-    /**
-     * Retrieve the truncation length, in bits, for the key in this algo. By default this will be
-     * the length in bits of the key.
-     */
+    /** Get the truncation length of this algorithm, in bits */
     public int getTruncationLengthBits() {
         return mTruncLenBits;
     }
@@ -167,7 +175,11 @@
     public static final Parcelable.Creator<IpSecAlgorithm> CREATOR =
             new Parcelable.Creator<IpSecAlgorithm>() {
                 public IpSecAlgorithm createFromParcel(Parcel in) {
-                    return new IpSecAlgorithm(in);
+                    final String name = in.readString();
+                    final byte[] key = in.createByteArray();
+                    final int truncLenBits = in.readInt();
+
+                    return new IpSecAlgorithm(name, key, truncLenBits);
                 }
 
                 public IpSecAlgorithm[] newArray(int size) {
@@ -175,30 +187,47 @@
                 }
             };
 
-    private IpSecAlgorithm(Parcel in) {
-        mName = in.readString();
-        mKey = in.createByteArray();
-        mTruncLenBits = in.readInt();
-    }
+    private static void checkValidOrThrow(String name, int keyLen, int truncLen) {
+        boolean isValidLen = true;
+        boolean isValidTruncLen = true;
 
-    private static boolean isTruncationLengthValid(String algo, int truncLenBits) {
-        switch (algo) {
+        switch(name) {
             case CRYPT_AES_CBC:
-                return (truncLenBits == 128 || truncLenBits == 192 || truncLenBits == 256);
+                isValidLen = keyLen == 128 || keyLen == 192 || keyLen == 256;
+                break;
             case AUTH_HMAC_MD5:
-                return (truncLenBits >= 96 && truncLenBits <= 128);
+                isValidLen = keyLen == 128;
+                isValidTruncLen = truncLen >= 96 && truncLen <= 128;
+                break;
             case AUTH_HMAC_SHA1:
-                return (truncLenBits >= 96 && truncLenBits <= 160);
+                isValidLen = keyLen == 160;
+                isValidTruncLen = truncLen >= 96 && truncLen <= 160;
+                break;
             case AUTH_HMAC_SHA256:
-                return (truncLenBits >= 96 && truncLenBits <= 256);
+                isValidLen = keyLen == 256;
+                isValidTruncLen = truncLen >= 96 && truncLen <= 256;
+                break;
             case AUTH_HMAC_SHA384:
-                return (truncLenBits >= 192 && truncLenBits <= 384);
+                isValidLen = keyLen == 384;
+                isValidTruncLen = truncLen >= 192 && truncLen <= 384;
+                break;
             case AUTH_HMAC_SHA512:
-                return (truncLenBits >= 256 && truncLenBits <= 512);
+                isValidLen = keyLen == 512;
+                isValidTruncLen = truncLen >= 256 && truncLen <= 512;
+                break;
             case AUTH_CRYPT_AES_GCM:
-                return (truncLenBits == 64 || truncLenBits == 96 || truncLenBits == 128);
+                // The keying material for GCM is a key plus a 32-bit salt
+                isValidLen = keyLen == 128 + 32 || keyLen == 192 + 32 || keyLen == 256 + 32;
+                break;
             default:
-                return false;
+                throw new IllegalArgumentException("Couldn't find an algorithm: " + name);
+        }
+
+        if (!isValidLen) {
+            throw new IllegalArgumentException("Invalid key material keyLength: " + keyLen);
+        }
+        if (!isValidTruncLen) {
+            throw new IllegalArgumentException("Invalid truncation keyLength: " + truncLen);
         }
     }
 
@@ -215,8 +244,9 @@
                 .toString();
     }
 
-    /** package */
-    static boolean equals(IpSecAlgorithm lhs, IpSecAlgorithm rhs) {
+    /** @hide */
+    @VisibleForTesting
+    public static boolean equals(IpSecAlgorithm lhs, IpSecAlgorithm rhs) {
         if (lhs == null || rhs == null) return (lhs == rhs);
         return (lhs.mName.equals(rhs.mName)
                 && Arrays.equals(lhs.mKey, rhs.mKey)
diff --git a/core/java/android/net/IpSecConfig.java b/core/java/android/net/IpSecConfig.java
index 61b13a9..e6cd3fc 100644
--- a/core/java/android/net/IpSecConfig.java
+++ b/core/java/android/net/IpSecConfig.java
@@ -20,7 +20,12 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 
-/** @hide */
+/**
+ * This class encapsulates all the configuration parameters needed to create IPsec transforms and
+ * policies.
+ *
+ * @hide
+ */
 public final class IpSecConfig implements Parcelable {
     private static final String TAG = "IpSecConfig";
 
@@ -38,6 +43,9 @@
     // for outbound packets. It may also be used to select packets.
     private Network mNetwork;
 
+    /**
+     * This class captures the parameters that specifically apply to inbound or outbound traffic.
+     */
     public static class Flow {
         // Minimum requirements for identifying a transform
         // SPI identifying the IPsec flow in packet processing
diff --git a/core/java/android/net/IpSecManager.java b/core/java/android/net/IpSecManager.java
index eccd5f4..a9e60ec 100644
--- a/core/java/android/net/IpSecManager.java
+++ b/core/java/android/net/IpSecManager.java
@@ -19,6 +19,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.SystemService;
+import android.annotation.TestApi;
 import android.content.Context;
 import android.os.Binder;
 import android.os.ParcelFileDescriptor;
@@ -37,22 +38,28 @@
 import java.net.Socket;
 
 /**
- * This class contains methods for managing IPsec sessions, which will perform kernel-space
- * encryption and decryption of socket or Network traffic.
+ * This class contains methods for managing IPsec sessions. Once configured, the kernel will apply
+ * confidentiality (encryption) and integrity (authentication) to IP traffic.
  *
- * <p>An IpSecManager may be obtained by calling {@link
- * android.content.Context#getSystemService(String) Context#getSystemService(String)} with {@link
- * android.content.Context#IPSEC_SERVICE Context#IPSEC_SERVICE}
+ * <p>Note that not all aspects of IPsec are permitted by this API. Applications may create
+ * transport mode security associations and apply them to individual sockets. Applications looking
+ * to create a VPN should use {@link VpnService}.
+ *
+ * @see <a href="https://tools.ietf.org/html/rfc4301">RFC 4301, Security Architecture for the
+ * Internet Protocol</a>
  */
 @SystemService(Context.IPSEC_SERVICE)
 public final class IpSecManager {
     private static final String TAG = "IpSecManager";
 
     /**
-     * The Security Parameter Index, SPI, 0 indicates an unknown or invalid index.
+     * The Security Parameter Index (SPI) 0 indicates an unknown or invalid index.
      *
      * <p>No IPsec packet may contain an SPI of 0.
+     *
+     * @hide
      */
+    @TestApi
     public static final int INVALID_SECURITY_PARAMETER_INDEX = 0;
 
     /** @hide */
@@ -66,10 +73,12 @@
     public static final int INVALID_RESOURCE_ID = 0;
 
     /**
-     * Indicates that the combination of remote InetAddress and SPI was non-unique for a given
-     * request. If encountered, selection of a new SPI is required before a transform may be
-     * created. Note, this should happen very rarely if the SPI is chosen to be sufficiently random
-     * or reserved using reserveSecurityParameterIndex.
+     * Thrown to indicate that a requested SPI is in use.
+     *
+     * <p>The combination of remote {@code InetAddress} and SPI must be unique across all apps on
+     * one device. If this error is encountered, a new SPI is required before a transform may be
+     * created. This error can be avoided by calling {@link
+     * IpSecManager#reserveSecurityParameterIndex}.
      */
     public static final class SpiUnavailableException extends AndroidException {
         private final int mSpi;
@@ -78,24 +87,26 @@
          * Construct an exception indicating that a transform with the given SPI is already in use
          * or otherwise unavailable.
          *
-         * @param msg Description indicating the colliding SPI
+         * @param msg description indicating the colliding SPI
          * @param spi the SPI that could not be used due to a collision
          */
         SpiUnavailableException(String msg, int spi) {
-            super(msg + "(spi: " + spi + ")");
+            super(msg + " (spi: " + spi + ")");
             mSpi = spi;
         }
 
-        /** Retrieve the SPI that caused a collision */
+        /** Get the SPI that caused a collision. */
         public int getSpi() {
             return mSpi;
         }
     }
 
     /**
-     * Indicates that the requested system resource for IPsec, such as a socket or other system
-     * resource is unavailable. If this exception is thrown, try releasing allocated objects of the
-     * type requested.
+     * Thrown to indicate that an IPsec resource is unavailable.
+     *
+     * <p>This could apply to resources such as sockets, {@link SecurityParameterIndex}, {@link
+     * IpSecTransform}, or other system resources. If this exception is thrown, users should release
+     * allocated objects of the type requested.
      */
     public static final class ResourceUnavailableException extends AndroidException {
 
@@ -106,6 +117,13 @@
 
     private final IIpSecService mService;
 
+    /**
+     * This class represents a reserved SPI.
+     *
+     * <p>Objects of this type are used to track reserved security parameter indices. They can be
+     * obtained by calling {@link IpSecManager#reserveSecurityParameterIndex} and must be released
+     * by calling {@link #close()} when they are no longer needed.
+     */
     public static final class SecurityParameterIndex implements AutoCloseable {
         private final IIpSecService mService;
         private final InetAddress mRemoteAddress;
@@ -113,7 +131,7 @@
         private int mSpi = INVALID_SECURITY_PARAMETER_INDEX;
         private int mResourceId;
 
-        /** Return the underlying SPI held by this object */
+        /** Get the underlying SPI held by this object. */
         public int getSpi() {
             return mSpi;
         }
@@ -135,6 +153,7 @@
             mCloseGuard.close();
         }
 
+        /** Check that the SPI was closed properly. */
         @Override
         protected void finalize() throws Throwable {
             if (mCloseGuard != null) {
@@ -197,13 +216,13 @@
     }
 
     /**
-     * Reserve an SPI for traffic bound towards the specified remote address.
+     * Reserve a random SPI for traffic bound to or from the specified remote address.
      *
      * <p>If successful, this SPI is guaranteed available until released by a call to {@link
      * SecurityParameterIndex#close()}.
      *
      * @param direction {@link IpSecTransform#DIRECTION_IN} or {@link IpSecTransform#DIRECTION_OUT}
-     * @param remoteAddress address of the remote. SPIs must be unique for each remoteAddress.
+     * @param remoteAddress address of the remote. SPIs must be unique for each remoteAddress
      * @return the reserved SecurityParameterIndex
      * @throws ResourceUnavailableException indicating that too many SPIs are currently allocated
      *     for this user
@@ -223,17 +242,18 @@
     }
 
     /**
-     * Reserve an SPI for traffic bound towards the specified remote address.
+     * Reserve the requested SPI for traffic bound to or from the specified remote address.
      *
      * <p>If successful, this SPI is guaranteed available until released by a call to {@link
      * SecurityParameterIndex#close()}.
      *
      * @param direction {@link IpSecTransform#DIRECTION_IN} or {@link IpSecTransform#DIRECTION_OUT}
-     * @param remoteAddress address of the remote. SPIs must be unique for each remoteAddress.
-     * @param requestedSpi the requested SPI, or '0' to allocate a random SPI.
+     * @param remoteAddress address of the remote. SPIs must be unique for each remoteAddress
+     * @param requestedSpi the requested SPI, or '0' to allocate a random SPI
      * @return the reserved SecurityParameterIndex
      * @throws ResourceUnavailableException indicating that too many SPIs are currently allocated
      *     for this user
+     * @throws SpiUnavailableException indicating that the requested SPI could not be reserved
      */
     public SecurityParameterIndex reserveSecurityParameterIndex(
             int direction, InetAddress remoteAddress, int requestedSpi)
@@ -245,16 +265,28 @@
     }
 
     /**
-     * Apply an active Transport Mode IPsec Transform to a stream socket to perform IPsec
-     * encapsulation of the traffic flowing between the socket and the remote InetAddress of that
-     * transform. For security reasons, attempts to send traffic to any IP address other than the
-     * address associated with that transform will throw an IOException. In addition, if the
-     * IpSecTransform is later deactivated, the socket will throw an IOException on any calls to
-     * send() or receive() until the transform is removed from the socket by calling {@link
-     * #removeTransportModeTransform(Socket, IpSecTransform)};
+     * Apply an IPsec transform to a stream socket.
+     *
+     * <p>This applies transport mode encapsulation to the given socket. Once applied, I/O on the
+     * socket will be encapsulated according to the parameters of the {@code IpSecTransform}. When
+     * the transform is removed from the socket by calling {@link #removeTransportModeTransform},
+     * unprotected traffic can resume on that socket.
+     *
+     * <p>For security reasons, the destination address of any traffic on the socket must match the
+     * remote {@code  InetAddress} of the {@code IpSecTransform}. Attempts to send traffic to any
+     * other IP address will result in an IOException. In addition, reads and writes on the socket
+     * will throw IOException if the user deactivates the transform (by calling {@link
+     * IpSecTransform#close()}) without calling {@link #removeTransportModeTransform}.
+     *
+     * <h4>Rekey Procedure</h4> <p>When applying a new tranform to a socket, the previous transform
+     * will be removed. However, inbound traffic on the old transform will continue to be decrypted
+     * until that transform is deallocated by calling {@link IpSecTransform#close()}. This overlap
+     * allows rekey procedures where both transforms are valid until both endpoints are using the
+     * new transform and all in-flight packets have been received.
      *
      * @param socket a stream socket
-     * @param transform an {@link IpSecTransform}, which must be an active Transport Mode transform.
+     * @param transform a transport mode {@code IpSecTransform}
+     * @throws IOException indicating that the transform could not be applied
      * @hide
      */
     public void applyTransportModeTransform(Socket socket, IpSecTransform transform)
@@ -265,16 +297,28 @@
     }
 
     /**
-     * Apply an active Transport Mode IPsec Transform to a datagram socket to perform IPsec
-     * encapsulation of the traffic flowing between the socket and the remote InetAddress of that
-     * transform. For security reasons, attempts to send traffic to any IP address other than the
-     * address associated with that transform will throw an IOException. In addition, if the
-     * IpSecTransform is later deactivated, the socket will throw an IOException on any calls to
-     * send() or receive() until the transform is removed from the socket by calling {@link
-     * #removeTransportModeTransform(DatagramSocket, IpSecTransform)};
+     * Apply an IPsec transform to a datagram socket.
+     *
+     * <p>This applies transport mode encapsulation to the given socket. Once applied, I/O on the
+     * socket will be encapsulated according to the parameters of the {@code IpSecTransform}. When
+     * the transform is removed from the socket by calling {@link #removeTransportModeTransform},
+     * unprotected traffic can resume on that socket.
+     *
+     * <p>For security reasons, the destination address of any traffic on the socket must match the
+     * remote {@code InetAddress} of the {@code IpSecTransform}. Attempts to send traffic to any
+     * other IP address will result in an IOException. In addition, reads and writes on the socket
+     * will throw IOException if the user deactivates the transform (by calling {@link
+     * IpSecTransform#close()}) without calling {@link #removeTransportModeTransform}.
+     *
+     * <h4>Rekey Procedure</h4> <p>When applying a new tranform to a socket, the previous transform
+     * will be removed. However, inbound traffic on the old transform will continue to be decrypted
+     * until that transform is deallocated by calling {@link IpSecTransform#close()}. This overlap
+     * allows rekey procedures where both transforms are valid until both endpoints are using the
+     * new transform and all in-flight packets have been received.
      *
      * @param socket a datagram socket
-     * @param transform an {@link IpSecTransform}, which must be an active Transport Mode transform.
+     * @param transform a transport mode {@code IpSecTransform}
+     * @throws IOException indicating that the transform could not be applied
      * @hide
      */
     public void applyTransportModeTransform(DatagramSocket socket, IpSecTransform transform)
@@ -285,16 +329,28 @@
     }
 
     /**
-     * Apply an active Transport Mode IPsec Transform to a stream socket to perform IPsec
-     * encapsulation of the traffic flowing between the socket and the remote InetAddress of that
-     * transform. For security reasons, attempts to send traffic to any IP address other than the
-     * address associated with that transform will throw an IOException. In addition, if the
-     * IpSecTransform is later deactivated, the socket will throw an IOException on any calls to
-     * send() or receive() until the transform is removed from the socket by calling {@link
-     * #removeTransportModeTransform(FileDescriptor, IpSecTransform)};
+     * Apply an IPsec transform to a socket.
+     *
+     * <p>This applies transport mode encapsulation to the given socket. Once applied, I/O on the
+     * socket will be encapsulated according to the parameters of the {@code IpSecTransform}. When
+     * the transform is removed from the socket by calling {@link #removeTransportModeTransform},
+     * unprotected traffic can resume on that socket.
+     *
+     * <p>For security reasons, the destination address of any traffic on the socket must match the
+     * remote {@code InetAddress} of the {@code IpSecTransform}. Attempts to send traffic to any
+     * other IP address will result in an IOException. In addition, reads and writes on the socket
+     * will throw IOException if the user deactivates the transform (by calling {@link
+     * IpSecTransform#close()}) without calling {@link #removeTransportModeTransform}.
+     *
+     * <h4>Rekey Procedure</h4> <p>When applying a new tranform to a socket, the previous transform
+     * will be removed. However, inbound traffic on the old transform will continue to be decrypted
+     * until that transform is deallocated by calling {@link IpSecTransform#close()}. This overlap
+     * allows rekey procedures where both transforms are valid until both endpoints are using the
+     * new transform and all in-flight packets have been received.
      *
      * @param socket a socket file descriptor
-     * @param transform an {@link IpSecTransform}, which must be an active Transport Mode transform.
+     * @param transform a transport mode {@code IpSecTransform}
+     * @throws IOException indicating that the transform could not be applied
      */
     public void applyTransportModeTransform(FileDescriptor socket, IpSecTransform transform)
             throws IOException {
@@ -323,6 +379,7 @@
      * Applications should probably not use this API directly. Instead, they should use {@link
      * VpnService} to provide VPN capability in a more generic fashion.
      *
+     * TODO: Update javadoc for tunnel mode APIs at the same time the APIs are re-worked.
      * @param net a {@link Network} that will be tunneled via IP Sec.
      * @param transform an {@link IpSecTransform}, which must be an active Tunnel Mode transform.
      * @hide
@@ -330,14 +387,19 @@
     public void applyTunnelModeTransform(Network net, IpSecTransform transform) {}
 
     /**
-     * Remove a transform from a given stream socket. Once removed, traffic on the socket will not
-     * be encypted. This allows sockets that have been used for IPsec to be reclaimed for
-     * communication in the clear in the event socket reuse is desired. This operation will succeed
-     * regardless of the underlying state of a transform. If a transform is removed, communication
-     * on all sockets to which that transform was applied will fail until this method is called.
+     * Remove an IPsec transform from a stream socket.
      *
-     * @param socket a socket that previously had a transform applied to it.
+     * <p>Once removed, traffic on the socket will not be encrypted. This operation will succeed
+     * regardless of the state of the transform. Removing a transform from a socket allows the
+     * socket to be reused for communication in the clear.
+     *
+     * <p>If an {@code IpSecTransform} object applied to this socket was deallocated by calling
+     * {@link IpSecTransform#close()}, then communication on the socket will fail until this method
+     * is called.
+     *
+     * @param socket a socket that previously had a transform applied to it
      * @param transform the IPsec Transform that was previously applied to the given socket
+     * @throws IOException indicating that the transform could not be removed from the socket
      * @hide
      */
     public void removeTransportModeTransform(Socket socket, IpSecTransform transform)
@@ -348,14 +410,19 @@
     }
 
     /**
-     * Remove a transform from a given datagram socket. Once removed, traffic on the socket will not
-     * be encypted. This allows sockets that have been used for IPsec to be reclaimed for
-     * communication in the clear in the event socket reuse is desired. This operation will succeed
-     * regardless of the underlying state of a transform. If a transform is removed, communication
-     * on all sockets to which that transform was applied will fail until this method is called.
+     * Remove an IPsec transform from a datagram socket.
      *
-     * @param socket a socket that previously had a transform applied to it.
+     * <p>Once removed, traffic on the socket will not be encrypted. This operation will succeed
+     * regardless of the state of the transform. Removing a transform from a socket allows the
+     * socket to be reused for communication in the clear.
+     *
+     * <p>If an {@code IpSecTransform} object applied to this socket was deallocated by calling
+     * {@link IpSecTransform#close()}, then communication on the socket will fail until this method
+     * is called.
+     *
+     * @param socket a socket that previously had a transform applied to it
      * @param transform the IPsec Transform that was previously applied to the given socket
+     * @throws IOException indicating that the transform could not be removed from the socket
      * @hide
      */
     public void removeTransportModeTransform(DatagramSocket socket, IpSecTransform transform)
@@ -366,14 +433,19 @@
     }
 
     /**
-     * Remove a transform from a given stream socket. Once removed, traffic on the socket will not
-     * be encypted. This allows sockets that have been used for IPsec to be reclaimed for
-     * communication in the clear in the event socket reuse is desired. This operation will succeed
-     * regardless of the underlying state of a transform. If a transform is removed, communication
-     * on all sockets to which that transform was applied will fail until this method is called.
+     * Remove an IPsec transform from a socket.
      *
-     * @param socket a socket file descriptor that previously had a transform applied to it.
+     * <p>Once removed, traffic on the socket will not be encrypted. This operation will succeed
+     * regardless of the state of the transform. Removing a transform from a socket allows the
+     * socket to be reused for communication in the clear.
+     *
+     * <p>If an {@code IpSecTransform} object applied to this socket was deallocated by calling
+     * {@link IpSecTransform#close()}, then communication on the socket will fail until this method
+     * is called.
+     *
+     * @param socket a socket that previously had a transform applied to it
      * @param transform the IPsec Transform that was previously applied to the given socket
+     * @throws IOException indicating that the transform could not be removed from the socket
      */
     public void removeTransportModeTransform(FileDescriptor socket, IpSecTransform transform)
             throws IOException {
@@ -382,7 +454,7 @@
         }
     }
 
-    /* Call down to activate a transform */
+    /* Call down to remove a transform */
     private void removeTransportModeTransform(ParcelFileDescriptor pfd, IpSecTransform transform) {
         try {
             mService.removeTransportModeTransform(pfd, transform.getResourceId());
@@ -397,6 +469,7 @@
      * all traffic that cannot be routed to the Tunnel's outbound interface. If that interface is
      * lost, all traffic will drop.
      *
+     * TODO: Update javadoc for tunnel mode APIs at the same time the APIs are re-worked.
      * @param net a network that currently has transform applied to it.
      * @param transform a Tunnel Mode IPsec Transform that has been previously applied to the given
      *     network
@@ -405,11 +478,18 @@
     public void removeTunnelModeTransform(Network net, IpSecTransform transform) {}
 
     /**
-     * Class providing access to a system-provided UDP Encapsulation Socket, which may be used for
-     * IKE signalling as well as for inbound and outbound UDP encapsulated IPsec traffic.
+     * This class provides access to a UDP encapsulation Socket.
      *
-     * <p>The socket provided by this class cannot be re-bound or closed via the inner
-     * FileDescriptor. Instead, disposing of this socket requires a call to close().
+     * <p>{@code UdpEncapsulationSocket} wraps a system-provided datagram socket intended for IKEv2
+     * signalling and UDP encapsulated IPsec traffic. Instances can be obtained by calling {@link
+     * IpSecManager#openUdpEncapsulationSocket}. The provided socket cannot be re-bound by the
+     * caller. The caller should not close the {@code FileDescriptor} returned by {@link
+     * #getSocket}, but should use {@link #close} instead.
+     *
+     * <p>Allowing the user to close or unbind a UDP encapsulation socket could impact the traffic
+     * of the next user who binds to that port. To prevent this scenario, these sockets are held
+     * open by the system so that they may only be closed by calling {@link #close} or when the user
+     * process exits.
      */
     public static final class UdpEncapsulationSocket implements AutoCloseable {
         private final ParcelFileDescriptor mPfd;
@@ -443,7 +523,7 @@
             mCloseGuard.open("constructor");
         }
 
-        /** Access the inner UDP Encapsulation Socket */
+        /** Get the wrapped socket. */
         public FileDescriptor getSocket() {
             if (mPfd == null) {
                 return null;
@@ -451,22 +531,19 @@
             return mPfd.getFileDescriptor();
         }
 
-        /** Retrieve the port number of the inner encapsulation socket */
+        /** Get the bound port of the wrapped socket. */
         public int getPort() {
             return mPort;
         }
 
-        @Override
         /**
-         * Release the resources that have been reserved for this Socket.
+         * Close this socket.
          *
-         * <p>This method closes the underlying socket, reducing a user's allocated sockets in the
-         * system. This must be done as part of cleanup following use of a socket. Failure to do so
-         * will cause the socket to count against a total allocation limit for IpSec and eventually
-         * fail due to resource limits.
-         *
-         * @param fd a file descriptor previously returned as a UDP Encapsulation socket.
+         * <p>This closes the wrapped socket. Open encapsulation sockets count against a user's
+         * resource limits, and forgetting to close them eventually will result in {@link
+         * ResourceUnavailableException} being thrown.
          */
+        @Override
         public void close() throws IOException {
             try {
                 mService.closeUdpEncapsulationSocket(mResourceId);
@@ -483,6 +560,7 @@
             mCloseGuard.close();
         }
 
+        /** Check that the socket was closed properly. */
         @Override
         protected void finalize() throws Throwable {
             if (mCloseGuard != null) {
@@ -499,21 +577,14 @@
     };
 
     /**
-     * Open a socket that is bound to a free UDP port on the system.
+     * Open a socket for UDP encapsulation and bind to the given port.
      *
-     * <p>By binding in this manner and holding the FileDescriptor, the socket cannot be un-bound by
-     * the caller. This provides safe access to a socket on a port that can later be used as a UDP
-     * Encapsulation port.
+     * <p>See {@link UdpEncapsulationSocket} for the proper way to close the returned socket.
      *
-     * <p>This socket reservation works in conjunction with IpSecTransforms, which may re-use the
-     * socket port. Explicitly opening this port is only necessary if communication is desired on
-     * that port.
-     *
-     * @param port a local UDP port to be reserved for UDP Encapsulation. is provided, then this
-     *     method will bind to the specified port or fail. To retrieve the port number, call {@link
-     *     android.system.Os#getsockname(FileDescriptor)}.
-     * @return a {@link UdpEncapsulationSocket} that is bound to the requested port for the lifetime
-     *     of the object.
+     * @param port a local UDP port
+     * @return a socket that is bound to the given port
+     * @throws IOException indicating that the socket could not be opened or bound
+     * @throws ResourceUnavailableException indicating that too many encapsulation sockets are open
      */
     // Returning a socket in this fashion that has been created and bound by the system
     // is the only safe way to ensure that a socket is both accessible to the user and
@@ -533,17 +604,16 @@
     }
 
     /**
-     * Open a socket that is bound to a port selected by the system.
+     * Open a socket for UDP encapsulation.
      *
-     * <p>By binding in this manner and holding the FileDescriptor, the socket cannot be un-bound by
-     * the caller. This provides safe access to a socket on a port that can later be used as a UDP
-     * Encapsulation port.
+     * <p>See {@link UdpEncapsulationSocket} for the proper way to close the returned socket.
      *
-     * <p>This socket reservation works in conjunction with IpSecTransforms, which may re-use the
-     * socket port. Explicitly opening this port is only necessary if communication is desired on
-     * that port.
+     * <p>The local port of the returned socket can be obtained by calling {@link
+     * UdpEncapsulationSocket#getPort()}.
      *
-     * @return a {@link UdpEncapsulationSocket} that is bound to an arbitrarily selected port
+     * @return a socket that is bound to a local port
+     * @throws IOException indicating that the socket could not be opened or bound
+     * @throws ResourceUnavailableException indicating that too many encapsulation sockets are open
      */
     // Returning a socket in this fashion that has been created and bound by the system
     // is the only safe way to ensure that a socket is both accessible to the user and
@@ -556,7 +626,7 @@
     }
 
     /**
-     * Retrieve an instance of an IpSecManager within you application context
+     * Construct an instance of IpSecManager within an application context.
      *
      * @param context the application context for this manager
      * @hide
diff --git a/core/java/android/net/IpSecTransform.java b/core/java/android/net/IpSecTransform.java
index 48b5bd5..cda4ec7 100644
--- a/core/java/android/net/IpSecTransform.java
+++ b/core/java/android/net/IpSecTransform.java
@@ -38,27 +38,29 @@
 import java.net.InetAddress;
 
 /**
- * This class represents an IpSecTransform, which encapsulates both properties and state of IPsec.
+ * This class represents an IPsec transform, which comprises security associations in one or both
+ * directions.
  *
- * <p>IpSecTransforms must be built from an IpSecTransform.Builder, and they must persist throughout
- * the lifetime of the underlying transform. If a transform object leaves scope, the underlying
- * transform may be disabled automatically, with likely undesirable results.
+ * <p>Transforms are created using {@link IpSecTransform.Builder}. Each {@code IpSecTransform}
+ * object encapsulates the properties and state of an inbound and outbound IPsec security
+ * association. That includes, but is not limited to, algorithm choice, key material, and allocated
+ * system resources.
  *
- * <p>An IpSecTransform may either represent a tunnel mode transform that operates on a wide array
- * of traffic or may represent a transport mode transform operating on a Socket or Sockets.
+ * @see <a href="https://tools.ietf.org/html/rfc4301">RFC 4301, Security Architecture for the
+ * Internet Protocol</a>
  */
 public final class IpSecTransform implements AutoCloseable {
     private static final String TAG = "IpSecTransform";
 
     /**
-     * For direction-specific attributes of an IpSecTransform, indicates that an attribute applies
-     * to traffic towards the host.
+     * For direction-specific attributes of an {@link IpSecTransform}, indicates that an attribute
+     * applies to traffic towards the host.
      */
     public static final int DIRECTION_IN = 0;
 
     /**
-     * For direction-specific attributes of an IpSecTransform, indicates that an attribute applies
-     * to traffic from the host.
+     * For direction-specific attributes of an {@link IpSecTransform}, indicates that an attribute
+     * applies to traffic from the host.
      */
     public static final int DIRECTION_OUT = 1;
 
@@ -77,16 +79,16 @@
     public static final int ENCAP_NONE = 0;
 
     /**
-     * IpSec traffic will be encapsulated within a UDP header with an additional 8-byte header pad
-     * (of '0'-value bytes) that prevents traffic from being interpreted as IKE or as ESP over UDP.
+     * IPsec traffic will be encapsulated within UDP, but with 8 zero-value bytes between the UDP
+     * header and payload. This prevents traffic from being interpreted as ESP or IKEv2.
      *
      * @hide
      */
     public static final int ENCAP_ESPINUDP_NON_IKE = 1;
 
     /**
-     * IpSec traffic will be encapsulated within UDP as per <a
-     * href="https://tools.ietf.org/html/rfc3948">RFC3498</a>.
+     * IPsec traffic will be encapsulated within UDP as per
+     * <a href="https://tools.ietf.org/html/rfc3948">RFC 3498</a>.
      *
      * @hide
      */
@@ -165,13 +167,14 @@
     }
 
     /**
-     * Deactivate an IpSecTransform and free all resources for that transform that are managed by
-     * the system for this Transform.
+     * Deactivate this {@code IpSecTransform} and free allocated resources.
      *
-     * <p>Deactivating a transform while it is still applied to any Socket will result in sockets
-     * refusing to send or receive data. This method will silently succeed if the specified
-     * transform has already been removed; thus, it is always safe to attempt cleanup when a
-     * transform is no longer needed.
+     * <p>Deactivating a transform while it is still applied to a socket will result in errors on
+     * that socket. Make sure to remove transforms by calling {@link
+     * IpSecManager#removeTransportModeTransform}. Note, removing an {@code IpSecTransform} from a
+     * socket will not deactivate it (because one transform may be applied to multiple sockets).
+     *
+     * <p>It is safe to call this method on a transform that has already been deactivated.
      */
     public void close() {
         Log.d(TAG, "Removing Transform with Id " + mResourceId);
@@ -197,6 +200,7 @@
         }
     }
 
+    /** Check that the transform was closed properly. */
     @Override
     protected void finalize() throws Throwable {
         if (mCloseGuard != null) {
@@ -264,65 +268,63 @@
     }
 
     /**
-     * Builder object to facilitate the creation of IpSecTransform objects.
-     *
-     * <p>Apply additional properties to the transform and then call a build() method to return an
-     * IpSecTransform object.
-     *
-     * @see Builder#buildTransportModeTransform(InetAddress)
+     * This class is used to build {@link IpSecTransform} objects.
      */
     public static class Builder {
         private Context mContext;
         private IpSecConfig mConfig;
 
         /**
-         * Add an encryption algorithm to the transform for the given direction.
+         * Set the encryption algorithm for the given direction.
          *
-         * <p>If encryption is set for a given direction without also providing an SPI for that
-         * direction, creation of an IpSecTransform will fail upon calling a build() method.
+         * <p>If encryption is set for a direction without also providing an SPI for that direction,
+         * creation of an {@code IpSecTransform} will fail when attempting to build the transform.
          *
-         * <p>Authenticated encryption is mutually exclusive with encryption and authentication.
+         * <p>Encryption is mutually exclusive with authenticated encryption.
          *
-         * @param direction either {@link #DIRECTION_IN or #DIRECTION_OUT}
+         * @param direction either {@link #DIRECTION_IN} or {@link #DIRECTION_OUT}
          * @param algo {@link IpSecAlgorithm} specifying the encryption to be applied.
          */
         public IpSecTransform.Builder setEncryption(
                 @TransformDirection int direction, IpSecAlgorithm algo) {
+            // TODO: throw IllegalArgumentException if algo is not an encryption algorithm.
             mConfig.setEncryption(direction, algo);
             return this;
         }
 
         /**
-         * Add an authentication/integrity algorithm to the transform.
+         * Set the authentication (integrity) algorithm for the given direction.
          *
-         * <p>If authentication is set for a given direction without also providing an SPI for that
-         * direction, creation of an IpSecTransform will fail upon calling a build() method.
+         * <p>If authentication is set for a direction without also providing an SPI for that
+         * direction, creation of an {@code IpSecTransform} will fail when attempting to build the
+         * transform.
          *
-         * <p>Authenticated encryption is mutually exclusive with encryption and authentication.
+         * <p>Authentication is mutually exclusive with authenticated encryption.
          *
-         * @param direction either {@link #DIRECTION_IN or #DIRECTION_OUT}
+         * @param direction either {@link #DIRECTION_IN} or {@link #DIRECTION_OUT}
          * @param algo {@link IpSecAlgorithm} specifying the authentication to be applied.
          */
         public IpSecTransform.Builder setAuthentication(
                 @TransformDirection int direction, IpSecAlgorithm algo) {
+            // TODO: throw IllegalArgumentException if algo is not an authentication algorithm.
             mConfig.setAuthentication(direction, algo);
             return this;
         }
 
         /**
-         * Add an authenticated encryption algorithm to the transform for the given direction.
+         * Set the authenticated encryption algorithm for the given direction.
          *
          * <p>If an authenticated encryption algorithm is set for a given direction without also
-         * providing an SPI for that direction, creation of an IpSecTransform will fail upon calling
-         * a build() method.
+         * providing an SPI for that direction, creation of an {@code IpSecTransform} will fail when
+         * attempting to build the transform.
          *
          * <p>The Authenticated Encryption (AE) class of algorithms are also known as Authenticated
          * Encryption with Associated Data (AEAD) algorithms, or Combined mode algorithms (as
-         * referred to in RFC 4301)
+         * referred to in <a href="https://tools.ietf.org/html/rfc4301">RFC 4301</a>).
          *
          * <p>Authenticated encryption is mutually exclusive with encryption and authentication.
          *
-         * @param direction either {@link #DIRECTION_IN or #DIRECTION_OUT}
+         * @param direction either {@link #DIRECTION_IN} or {@link #DIRECTION_OUT}
          * @param algo {@link IpSecAlgorithm} specifying the authenticated encryption algorithm to
          *     be applied.
          */
@@ -333,19 +335,16 @@
         }
 
         /**
-         * Set the SPI, which uniquely identifies a particular IPsec session from others. Because
-         * IPsec operates at the IP layer, this 32-bit identifier uniquely identifies packets to a
-         * given destination address.
+         * Set the SPI for the given direction.
          *
-         * <p>Care should be chosen when selecting an SPI to ensure that is is as unique as
-         * possible. To reserve a value call {@link IpSecManager#reserveSecurityParameterIndex(int,
-         * InetAddress, int)}. Otherwise, SPI collisions would prevent a transform from being
-         * activated. IpSecManager#reserveSecurityParameterIndex(int, InetAddres$s, int)}.
+         * <p>Because IPsec operates at the IP layer, this 32-bit identifier uniquely identifies
+         * packets to a given destination address. To prevent SPI collisions, values should be
+         * reserved by calling {@link IpSecManager#reserveSecurityParameterIndex}.
          *
-         * <p>Unless an SPI is set for a given direction, traffic in that direction will be
-         * sent/received without any IPsec applied.
+         * <p>If the SPI and algorithms are omitted for one direction, traffic in that direction
+         * will not be encrypted or authenticated.
          *
-         * @param direction either {@link #DIRECTION_IN or #DIRECTION_OUT}
+         * @param direction either {@link #DIRECTION_IN} or {@link #DIRECTION_OUT}
          * @param spi a unique {@link IpSecManager.SecurityParameterIndex} to identify transformed
          *     traffic
          */
@@ -356,11 +355,10 @@
         }
 
         /**
-         * Specify the network on which this transform will emit its traffic; (otherwise it will
-         * emit on the default network).
+         * Set the {@link Network} which will carry tunneled traffic.
          *
-         * <p>Restricts the transformed traffic to a particular {@link Network}. This is required in
-         * tunnel mode.
+         * <p>Restricts the transformed traffic to a particular {@link Network}. This is required
+         * for tunnel mode, otherwise tunneled traffic would be sent on the default network.
          *
          * @hide
          */
@@ -371,15 +369,18 @@
         }
 
         /**
-         * Add UDP encapsulation to an IPv4 transform
+         * Add UDP encapsulation to an IPv4 transform.
          *
-         * <p>This option allows IPsec traffic to pass through NAT. Refer to RFC 3947 and 3948 for
-         * details on how UDP should be applied to IPsec.
+         * <p>This allows IPsec traffic to pass through a NAT.
          *
-         * @param localSocket a {@link IpSecManager.UdpEncapsulationSocket} for sending and
-         *     receiving encapsulating traffic.
-         * @param remotePort the UDP port number of the remote that will send and receive
-         *     encapsulated traffic. In the case of IKE, this is likely port 4500.
+         * @see <a href="https://tools.ietf.org/html/rfc3948">RFC 3948, UDP Encapsulation of IPsec
+         * ESP Packets</a>
+         * @see <a href="https://tools.ietf.org/html/rfc7296#section-2.23">RFC 7296 section 2.23,
+         * NAT Traversal of IKEv2</a>
+         *
+         * @param localSocket a socket for sending and receiving encapsulated traffic
+         * @param remotePort the UDP port number of the remote host that will send and receive
+         *     encapsulated traffic. In the case of IKEv2, this should be port 4500.
          */
         public IpSecTransform.Builder setIpv4Encapsulation(
                 IpSecManager.UdpEncapsulationSocket localSocket, int remotePort) {
@@ -393,12 +394,15 @@
         // TODO: Probably a better exception to throw for NATTKeepalive failure
         // TODO: Specify the needed NATT keepalive permission.
         /**
-         * Send a NATT Keepalive packet with a given maximum interval. This will create an offloaded
-         * request to do power-efficient NATT Keepalive. If NATT keepalive is requested but cannot
-         * be activated, then the transform will fail to activate and throw an IOException.
+         * Set NAT-T keepalives to be sent with a given interval.
          *
-         * @param intervalSeconds the maximum number of seconds between keepalive packets, no less
-         *     than 20s and no more than 3600s.
+         * <p>This will set power-efficient keepalive packets to be sent by the system. If NAT-T
+         * keepalive is requested but cannot be activated, then creation of an {@link
+         * IpSecTransform} will fail when calling the build method.
+         *
+         * @param intervalSeconds the maximum number of seconds between keepalive packets. Must be
+         *     between 20s and 3600s.
+         *
          * @hide
          */
         @SystemApi
@@ -408,36 +412,29 @@
         }
 
         /**
-         * Build and return an active {@link IpSecTransform} object as a Transport Mode Transform.
-         * Some parameters have interdependencies that are checked at build time. If a well-formed
-         * transform cannot be created from the supplied parameters, this method will throw an
-         * Exception.
+         * Build a transport mode {@link IpSecTransform}.
          *
-         * <p>Upon a successful return from this call, the provided IpSecTransform will be active
-         * and may be applied to sockets. If too many IpSecTransform objects are active for a given
-         * user this operation will fail and throw ResourceUnavailableException. To avoid these
-         * exceptions, unused Transform objects must be cleaned up by calling {@link
-         * IpSecTransform#close()} when they are no longer needed.
+         * <p>This builds and activates a transport mode transform. Note that an active transform
+         * will not affect any network traffic until it has been applied to one or more sockets.
          *
-         * @param remoteAddress the {@link InetAddress} that, when matched on traffic to/from this
-         *     socket will cause the transform to be applied.
-         *     <p>Note that an active transform will not impact any network traffic until it has
-         *     been applied to one or more Sockets. Calling this method is a necessary precondition
-         *     for applying it to a socket, but is not sufficient to actually apply IPsec.
+         * @see IpSecManager#applyTransportModeTransform
+         *
+         * @param remoteAddress the remote {@code InetAddress} of traffic on sockets that will use
+         *     this transform
          * @throws IllegalArgumentException indicating that a particular combination of transform
-         *     properties is invalid.
-         * @throws IpSecManager.ResourceUnavailableException in the event that no more Transforms
-         *     may be allocated
-         * @throws SpiUnavailableException if the SPI collides with an existing transform
-         *     (unlikely).
-         * @throws ResourceUnavailableException if the current user currently has exceeded the
-         *     number of allowed active transforms.
+         *     properties is invalid
+         * @throws IpSecManager.ResourceUnavailableException indicating that too many transforms are
+         *     active
+         * @throws IpSecManager.SpiUnavailableException indicating the rare case where an SPI
+         *     collides with an existing transform
+         * @throws IOException indicating other errors
          */
         public IpSecTransform buildTransportModeTransform(InetAddress remoteAddress)
                 throws IpSecManager.ResourceUnavailableException,
                         IpSecManager.SpiUnavailableException, IOException {
             mConfig.setMode(MODE_TRANSPORT);
             mConfig.setRemoteAddress(remoteAddress.getHostAddress());
+            // FIXME: modifying a builder after calling build can change the built transform.
             return new IpSecTransform(mContext, mConfig).activate();
         }
 
@@ -465,9 +462,9 @@
         }
 
         /**
-         * Create a new IpSecTransform.Builder to construct an IpSecTransform
+         * Create a new IpSecTransform.Builder.
          *
-         * @param context current Context
+         * @param context current context
          */
         public Builder(@NonNull Context context) {
             Preconditions.checkNotNull(context);
diff --git a/core/java/android/net/NetworkIdentity.java b/core/java/android/net/NetworkIdentity.java
index df404b7..d3b3599 100644
--- a/core/java/android/net/NetworkIdentity.java
+++ b/core/java/android/net/NetworkIdentity.java
@@ -157,7 +157,7 @@
      * Scrub given IMSI on production builds.
      */
     public static String scrubSubscriberId(String subscriberId) {
-        if ("eng".equals(Build.TYPE)) {
+        if (Build.IS_ENG) {
             return subscriberId;
         } else if (subscriberId != null) {
             // TODO: parse this as MCC+MNC instead of hard-coding
diff --git a/core/java/android/net/NetworkStatsHistory.java b/core/java/android/net/NetworkStatsHistory.java
index 5f521de..433f941 100644
--- a/core/java/android/net/NetworkStatsHistory.java
+++ b/core/java/android/net/NetworkStatsHistory.java
@@ -27,6 +27,7 @@
 import static android.net.NetworkStatsHistory.ParcelUtils.readLongArray;
 import static android.net.NetworkStatsHistory.ParcelUtils.writeLongArray;
 import static android.text.format.DateUtils.SECOND_IN_MILLIS;
+
 import static com.android.internal.util.ArrayUtils.total;
 
 import android.os.Parcel;
@@ -282,6 +283,24 @@
         return entry;
     }
 
+    public void setValues(int i, Entry entry) {
+        // Unwind old values
+        if (rxBytes != null) totalBytes -= rxBytes[i];
+        if (txBytes != null) totalBytes -= txBytes[i];
+
+        bucketStart[i] = entry.bucketStart;
+        setLong(activeTime, i, entry.activeTime);
+        setLong(rxBytes, i, entry.rxBytes);
+        setLong(rxPackets, i, entry.rxPackets);
+        setLong(txBytes, i, entry.txBytes);
+        setLong(txPackets, i, entry.txPackets);
+        setLong(operations, i, entry.operations);
+
+        // Apply new values
+        if (rxBytes != null) totalBytes += rxBytes[i];
+        if (txBytes != null) totalBytes += txBytes[i];
+    }
+
     /**
      * Record that data traffic occurred in the given time range. Will
      * distribute across internal buckets, creating new buckets as needed.
diff --git a/core/java/android/net/NetworkTemplate.java b/core/java/android/net/NetworkTemplate.java
index 0d2fcd0..b307c5d 100644
--- a/core/java/android/net/NetworkTemplate.java
+++ b/core/java/android/net/NetworkTemplate.java
@@ -326,6 +326,10 @@
         }
     }
 
+    public boolean matchesSubscriberId(String subscriberId) {
+        return ArrayUtils.contains(mMatchSubscriberIds, subscriberId);
+    }
+
     /**
      * Check if mobile network with matching IMSI.
      */
diff --git a/core/java/android/net/TrafficStats.java b/core/java/android/net/TrafficStats.java
index f934616..c339856 100644
--- a/core/java/android/net/TrafficStats.java
+++ b/core/java/android/net/TrafficStats.java
@@ -109,25 +109,26 @@
      */
     public static final int TAG_SYSTEM_RESTORE = 0xFFFFFF04;
 
-    /** @hide */
-    public static final int TAG_SYSTEM_DHCP = 0xFFFFFF05;
-    /** @hide */
-    public static final int TAG_SYSTEM_NTP = 0xFFFFFF06;
-    /** @hide */
-    public static final int TAG_SYSTEM_PROBE = 0xFFFFFF07;
-    /** @hide */
-    public static final int TAG_SYSTEM_NEIGHBOR = 0xFFFFFF08;
-    /** @hide */
-    public static final int TAG_SYSTEM_GPS = 0xFFFFFF09;
-    /** @hide */
-    public static final int TAG_SYSTEM_PAC = 0xFFFFFF0A;
-
     /**
-     * Sockets that are strictly local on device; never hits network.
+     * Default tag value for code (typically APKs) downloaded by an app store on
+     * behalf of the app, such as updates.
      *
      * @hide
      */
-    public static final int TAG_SYSTEM_LOCAL = 0xFFFFFFAA;
+    public static final int TAG_SYSTEM_APP = 0xFFFFFF05;
+
+    /** @hide */
+    public static final int TAG_SYSTEM_DHCP = 0xFFFFFF40;
+    /** @hide */
+    public static final int TAG_SYSTEM_NTP = 0xFFFFFF41;
+    /** @hide */
+    public static final int TAG_SYSTEM_PROBE = 0xFFFFFF42;
+    /** @hide */
+    public static final int TAG_SYSTEM_NEIGHBOR = 0xFFFFFF43;
+    /** @hide */
+    public static final int TAG_SYSTEM_GPS = 0xFFFFFF44;
+    /** @hide */
+    public static final int TAG_SYSTEM_PAC = 0xFFFFFF45;
 
     private static INetworkStatsService sStatsService;
 
@@ -210,6 +211,19 @@
     }
 
     /**
+     * Set active tag to use when accounting {@link Socket} traffic originating
+     * from the current thread. The tag used internally is well-defined to
+     * distinguish all code (typically APKs) downloaded by an app store on
+     * behalf of the app, such as updates.
+     *
+     * @hide
+     */
+    @SystemApi
+    public static void setThreadStatsTagApp() {
+        setThreadStatsTag(TAG_SYSTEM_APP);
+    }
+
+    /**
      * Get the active tag used when accounting {@link Socket} traffic originating
      * from the current thread. Only one active tag per thread is supported.
      * {@link #tagSocket(Socket)}.
diff --git a/services/core/java/com/android/server/net/NetworkStatsCollection.java b/services/core/java/com/android/server/net/NetworkStatsCollection.java
index 0354300..4ceb592 100644
--- a/services/core/java/com/android/server/net/NetworkStatsCollection.java
+++ b/services/core/java/com/android/server/net/NetworkStatsCollection.java
@@ -28,6 +28,8 @@
 import static android.net.TrafficStats.UID_REMOVED;
 import static android.text.format.DateUtils.WEEK_IN_MILLIS;
 
+import static com.android.server.net.NetworkStatsService.TAG;
+
 import android.net.NetworkIdentity;
 import android.net.NetworkStats;
 import android.net.NetworkStatsHistory;
@@ -37,20 +39,24 @@
 import android.service.NetworkStatsCollectionKeyProto;
 import android.service.NetworkStatsCollectionProto;
 import android.service.NetworkStatsCollectionStatsProto;
+import android.telephony.SubscriptionPlan;
 import android.util.ArrayMap;
 import android.util.AtomicFile;
 import android.util.IntArray;
+import android.util.Pair;
+import android.util.Slog;
 import android.util.proto.ProtoOutputStream;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.FileRotator;
 import com.android.internal.util.IndentingPrintWriter;
 
+import libcore.io.IoUtils;
+
 import com.google.android.collect.Lists;
 import com.google.android.collect.Maps;
 
-import libcore.io.IoUtils;
-
 import java.io.BufferedInputStream;
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
@@ -60,9 +66,11 @@
 import java.io.InputStream;
 import java.io.PrintWriter;
 import java.net.ProtocolException;
+import java.time.ZonedDateTime;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.Objects;
 
 /**
@@ -140,6 +148,63 @@
         return mStartMillis == Long.MAX_VALUE && mEndMillis == Long.MIN_VALUE;
     }
 
+    @VisibleForTesting
+    public long roundUp(long time) {
+        if (time == Long.MIN_VALUE || time == Long.MAX_VALUE
+                || time == SubscriptionPlan.TIME_UNKNOWN) {
+            return time;
+        } else {
+            final long mod = time % mBucketDuration;
+            if (mod > 0) {
+                time -= mod;
+                time += mBucketDuration;
+            }
+            return time;
+        }
+    }
+
+    @VisibleForTesting
+    public long roundDown(long time) {
+        if (time == Long.MIN_VALUE || time == Long.MAX_VALUE
+                || time == SubscriptionPlan.TIME_UNKNOWN) {
+            return time;
+        } else {
+            final long mod = time % mBucketDuration;
+            if (mod > 0) {
+                time -= mod;
+            }
+            return time;
+        }
+    }
+
+    /**
+     * Safely multiple a value by a rational.
+     * <p>
+     * Internally it uses integer-based math whenever possible, but switches
+     * over to double-based math if values would overflow.
+     */
+    @VisibleForTesting
+    public static long multiplySafe(long value, long num, long den) {
+        long x = value;
+        long y = num;
+
+        // Logic shamelessly borrowed from Math.multiplyExact()
+        long r = x * y;
+        long ax = Math.abs(x);
+        long ay = Math.abs(y);
+        if (((ax | ay) >>> 31 != 0)) {
+            // Some bits greater than 2^31 that might cause overflow
+            // Check the result using the divide operator
+            // and check for the special case of Long.MIN_VALUE * -1
+            if (((y != 0) && (r / y != x)) ||
+                    (x == Long.MIN_VALUE && y == -1)) {
+                // Use double math to avoid overflowing
+                return (long) (((double) num / den) * value);
+            }
+        }
+        return r / den;
+    }
+
     public int[] getRelevantUids(@NetworkStatsAccess.Level int accessLevel) {
         return getRelevantUids(accessLevel, Binder.getCallingUid());
     }
@@ -165,60 +230,110 @@
      * Combine all {@link NetworkStatsHistory} in this collection which match
      * the requested parameters.
      */
-    public NetworkStatsHistory getHistory(
-            NetworkTemplate template, int uid, int set, int tag, int fields,
-            @NetworkStatsAccess.Level int accessLevel) {
-        return getHistory(template, uid, set, tag, fields, Long.MIN_VALUE, Long.MAX_VALUE,
-                accessLevel);
-    }
-
-    /**
-     * Combine all {@link NetworkStatsHistory} in this collection which match
-     * the requested parameters.
-     */
-    public NetworkStatsHistory getHistory(
-            NetworkTemplate template, int uid, int set, int tag, int fields, long start, long end,
-            @NetworkStatsAccess.Level int accessLevel) {
-        return getHistory(template, uid, set, tag, fields, start, end, accessLevel,
-                Binder.getCallingUid());
-    }
-
-    /**
-     * Combine all {@link NetworkStatsHistory} in this collection which match
-     * the requested parameters.
-     */
-    public NetworkStatsHistory getHistory(
-            NetworkTemplate template, int uid, int set, int tag, int fields, long start, long end,
+    public NetworkStatsHistory getHistory(NetworkTemplate template, SubscriptionPlan augmentPlan,
+            int uid, int set, int tag, int fields, long start, long end,
             @NetworkStatsAccess.Level int accessLevel, int callerUid) {
         if (!NetworkStatsAccess.isAccessibleToUser(uid, callerUid, accessLevel)) {
             throw new SecurityException("Network stats history of uid " + uid
                     + " is forbidden for caller " + callerUid);
         }
 
+        final int bucketEstimate = (int) ((end - start) / mBucketDuration);
         final NetworkStatsHistory combined = new NetworkStatsHistory(
-                mBucketDuration, start == end ? 1 : estimateBuckets(), fields);
+                mBucketDuration, bucketEstimate, fields);
 
         // shortcut when we know stats will be empty
         if (start == end) return combined;
 
+        // Figure out the window of time that we should be augmenting (if any)
+        long augmentStart = SubscriptionPlan.TIME_UNKNOWN;
+        long augmentEnd = (augmentPlan != null) ? augmentPlan.getDataUsageTime()
+                : SubscriptionPlan.TIME_UNKNOWN;
+        // And if augmenting, we might need to collect more data to adjust with
+        long collectStart = start;
+        long collectEnd = end;
+
+        if (augmentEnd != SubscriptionPlan.TIME_UNKNOWN) {
+            final Iterator<Pair<ZonedDateTime, ZonedDateTime>> it = augmentPlan.cycleIterator();
+            while (it.hasNext()) {
+                final Pair<ZonedDateTime, ZonedDateTime> cycle = it.next();
+                final long cycleStart = cycle.first.toInstant().toEpochMilli();
+                final long cycleEnd = cycle.second.toInstant().toEpochMilli();
+                if (cycleStart <= augmentEnd && augmentEnd < cycleEnd) {
+                    augmentStart = cycleStart;
+                    collectStart = Long.min(collectStart, augmentStart);
+                    collectEnd = Long.max(collectEnd, augmentEnd);
+                    break;
+                }
+            }
+        }
+
+        if (augmentStart != SubscriptionPlan.TIME_UNKNOWN) {
+            // Shrink augmentation window so we don't risk undercounting.
+            augmentStart = roundUp(augmentStart);
+            augmentEnd = roundDown(augmentEnd);
+            // Grow collection window so we get all the stats needed.
+            collectStart = roundDown(collectStart);
+            collectEnd = roundUp(collectEnd);
+        }
+
         for (int i = 0; i < mStats.size(); i++) {
             final Key key = mStats.keyAt(i);
             if (key.uid == uid && NetworkStats.setMatches(set, key.set) && key.tag == tag
                     && templateMatches(template, key.ident)) {
                 final NetworkStatsHistory value = mStats.valueAt(i);
-                combined.recordHistory(value, start, end);
+                combined.recordHistory(value, collectStart, collectEnd);
             }
         }
-        return combined;
-    }
 
-    /**
-     * Summarize all {@link NetworkStatsHistory} in this collection which match
-     * the requested parameters.
-     */
-    public NetworkStats getSummary(NetworkTemplate template, long start, long end,
-            @NetworkStatsAccess.Level int accessLevel) {
-        return getSummary(template, start, end, accessLevel, Binder.getCallingUid());
+        if (augmentStart != SubscriptionPlan.TIME_UNKNOWN) {
+            final NetworkStatsHistory.Entry entry = combined.getValues(
+                    augmentStart, augmentEnd, null);
+
+            // If we don't have any recorded data for this time period, give
+            // ourselves something to scale with.
+            if (entry.rxBytes == 0 || entry.txBytes == 0) {
+                combined.recordData(augmentStart, augmentEnd,
+                        new NetworkStats.Entry(1, 0, 1, 0, 0));
+                combined.getValues(augmentStart, augmentEnd, entry);
+            }
+
+            final long rawBytes = entry.rxBytes + entry.txBytes;
+            final long rawRxBytes = entry.rxBytes;
+            final long rawTxBytes = entry.txBytes;
+            final long targetBytes = augmentPlan.getDataUsageBytes();
+            final long targetRxBytes = multiplySafe(targetBytes, rawRxBytes, rawBytes);
+            final long targetTxBytes = multiplySafe(targetBytes, rawTxBytes, rawBytes);
+
+            // Scale all matching buckets to reach anchor target
+            final long beforeTotal = combined.getTotalBytes();
+            for (int i = 0; i < combined.size(); i++) {
+                combined.getValues(i, entry);
+                if (entry.bucketStart >= augmentStart
+                        && entry.bucketStart + entry.bucketDuration <= augmentEnd) {
+                    entry.rxBytes = multiplySafe(targetRxBytes, entry.rxBytes, rawRxBytes);
+                    entry.txBytes = multiplySafe(targetTxBytes, entry.txBytes, rawTxBytes);
+                    // We purposefully clear out packet counters to indicate
+                    // that this data has been augmented.
+                    entry.rxPackets = 0;
+                    entry.txPackets = 0;
+                    combined.setValues(i, entry);
+                }
+            }
+
+            final long deltaTotal = combined.getTotalBytes() - beforeTotal;
+            if (deltaTotal != 0) {
+                Slog.d(TAG, "Augmented network usage by " + deltaTotal + " bytes");
+            }
+
+            // Finally we can slice data as originally requested
+            final NetworkStatsHistory sliced = new NetworkStatsHistory(
+                    mBucketDuration, bucketEstimate, fields);
+            sliced.recordHistory(combined, start, end);
+            return sliced;
+        } else {
+            return combined;
+        }
     }
 
     /**
@@ -230,6 +345,7 @@
         final long now = System.currentTimeMillis();
 
         final NetworkStats stats = new NetworkStats(end - start, 24);
+
         // shortcut when we know stats will be empty
         if (start == end) return stats;
 
diff --git a/services/core/java/com/android/server/net/NetworkStatsObservers.java b/services/core/java/com/android/server/net/NetworkStatsObservers.java
index a256cbc..741c206 100644
--- a/services/core/java/com/android/server/net/NetworkStatsObservers.java
+++ b/services/core/java/com/android/server/net/NetworkStatsObservers.java
@@ -17,28 +17,26 @@
 package com.android.server.net;
 
 import static android.net.TrafficStats.MB_IN_BYTES;
+
 import static com.android.internal.util.Preconditions.checkArgument;
 
 import android.app.usage.NetworkStatsManager;
 import android.net.DataUsageRequest;
 import android.net.NetworkStats;
-import android.net.NetworkStats.NonMonotonicObserver;
 import android.net.NetworkStatsHistory;
 import android.net.NetworkTemplate;
-import android.os.Binder;
 import android.os.Bundle;
-import android.os.Looper;
-import android.os.Message;
-import android.os.Messenger;
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.IBinder;
+import android.os.Looper;
+import android.os.Message;
+import android.os.Messenger;
 import android.os.Process;
 import android.os.RemoteException;
 import android.util.ArrayMap;
-import android.util.IntArray;
-import android.util.SparseArray;
 import android.util.Slog;
+import android.util.SparseArray;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.net.VpnInfo;
@@ -410,7 +408,7 @@
          */
         private long getTotalBytesForNetworkUid(NetworkTemplate template, int uid) {
             try {
-                NetworkStatsHistory history = mCollection.getHistory(template, uid,
+                NetworkStatsHistory history = mCollection.getHistory(template, null, uid,
                         NetworkStats.SET_ALL, NetworkStats.TAG_NONE,
                         NetworkStatsHistory.FIELD_ALL,
                         Long.MIN_VALUE /* start */, Long.MAX_VALUE /* end */,
diff --git a/services/core/java/com/android/server/net/NetworkStatsRecorder.java b/services/core/java/com/android/server/net/NetworkStatsRecorder.java
index 80309e1..4bee55e 100644
--- a/services/core/java/com/android/server/net/NetworkStatsRecorder.java
+++ b/services/core/java/com/android/server/net/NetworkStatsRecorder.java
@@ -20,6 +20,7 @@
 import static android.net.TrafficStats.KB_IN_BYTES;
 import static android.net.TrafficStats.MB_IN_BYTES;
 import static android.text.format.DateUtils.YEAR_IN_MILLIS;
+
 import static com.android.internal.util.Preconditions.checkNotNull;
 
 import android.annotation.Nullable;
@@ -28,6 +29,7 @@
 import android.net.NetworkStatsHistory;
 import android.net.NetworkTemplate;
 import android.net.TrafficStats;
+import android.os.Binder;
 import android.os.DropBoxManager;
 import android.service.NetworkStatsRecorderProto;
 import android.util.Log;
@@ -38,6 +40,9 @@
 import com.android.internal.net.VpnInfo;
 import com.android.internal.util.FileRotator;
 import com.android.internal.util.IndentingPrintWriter;
+
+import libcore.io.IoUtils;
+
 import com.google.android.collect.Sets;
 
 import java.io.ByteArrayOutputStream;
@@ -52,8 +57,6 @@
 import java.util.HashSet;
 import java.util.Map;
 
-import libcore.io.IoUtils;
-
 /**
  * Logic to record deltas between periodic {@link NetworkStats} snapshots into
  * {@link NetworkStatsHistory} that belong to {@link NetworkStatsCollection}.
@@ -150,7 +153,7 @@
 
     public NetworkStats.Entry getTotalSinceBootLocked(NetworkTemplate template) {
         return mSinceBoot.getSummary(template, Long.MIN_VALUE, Long.MAX_VALUE,
-                NetworkStatsAccess.Level.DEVICE).getTotal(null);
+                NetworkStatsAccess.Level.DEVICE, Binder.getCallingUid()).getTotal(null);
     }
 
     public NetworkStatsCollection getSinceBoot() {
diff --git a/services/core/java/com/android/server/net/NetworkStatsService.java b/services/core/java/com/android/server/net/NetworkStatsService.java
index 421db40..3af5265 100644
--- a/services/core/java/com/android/server/net/NetworkStatsService.java
+++ b/services/core/java/com/android/server/net/NetworkStatsService.java
@@ -18,7 +18,6 @@
 
 import static android.Manifest.permission.ACCESS_NETWORK_STATE;
 import static android.Manifest.permission.CONNECTIVITY_INTERNAL;
-import static android.Manifest.permission.MODIFY_NETWORK_ACCOUNTING;
 import static android.Manifest.permission.READ_NETWORK_USAGE_HISTORY;
 import static android.content.Intent.ACTION_SHUTDOWN;
 import static android.content.Intent.ACTION_UID_REMOVED;
@@ -27,6 +26,8 @@
 import static android.net.ConnectivityManager.ACTION_TETHER_STATE_CHANGED;
 import static android.net.ConnectivityManager.isNetworkTypeMobile;
 import static android.net.NetworkStats.IFACE_ALL;
+import static android.net.NetworkStats.METERED_ALL;
+import static android.net.NetworkStats.ROAMING_ALL;
 import static android.net.NetworkStats.SET_ALL;
 import static android.net.NetworkStats.SET_DEFAULT;
 import static android.net.NetworkStats.SET_FOREGROUND;
@@ -34,10 +35,12 @@
 import static android.net.NetworkStats.STATS_PER_UID;
 import static android.net.NetworkStats.TAG_NONE;
 import static android.net.NetworkStats.UID_ALL;
+import static android.net.NetworkStatsHistory.FIELD_ALL;
 import static android.net.NetworkTemplate.buildTemplateMobileWildcard;
 import static android.net.NetworkTemplate.buildTemplateWifiWildcard;
 import static android.net.TrafficStats.KB_IN_BYTES;
 import static android.net.TrafficStats.MB_IN_BYTES;
+import static android.provider.Settings.Global.NETSTATS_AUGMENT_ENABLED;
 import static android.provider.Settings.Global.NETSTATS_DEV_BUCKET_DURATION;
 import static android.provider.Settings.Global.NETSTATS_DEV_DELETE_AGE;
 import static android.provider.Settings.Global.NETSTATS_DEV_PERSIST_BYTES;
@@ -66,6 +69,7 @@
 
 import android.app.AlarmManager;
 import android.app.PendingIntent;
+import android.app.usage.NetworkStatsManager;
 import android.content.BroadcastReceiver;
 import android.content.ContentResolver;
 import android.content.Context;
@@ -105,6 +109,8 @@
 import android.provider.Settings.Global;
 import android.service.NetworkInterfaceProto;
 import android.service.NetworkStatsServiceDumpProto;
+import android.telephony.SubscriptionManager;
+import android.telephony.SubscriptionPlan;
 import android.telephony.TelephonyManager;
 import android.text.format.DateUtils;
 import android.util.ArrayMap;
@@ -118,6 +124,7 @@
 import android.util.TrustedTime;
 import android.util.proto.ProtoOutputStream;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.net.VpnInfo;
 import com.android.internal.util.ArrayUtils;
@@ -140,8 +147,8 @@
  * other system services.
  */
 public class NetworkStatsService extends INetworkStatsService.Stub {
-    private static final String TAG = "NetworkStats";
-    private static final boolean LOGV = false;
+    static final String TAG = "NetworkStats";
+    static final boolean LOGV = false;
 
     private static final int MSG_PERFORM_POLL = 1;
     private static final int MSG_UPDATE_IFACES = 2;
@@ -195,6 +202,7 @@
         public long getPollInterval();
         public long getTimeCacheMaxAge();
         public boolean getSampleEnabled();
+        public boolean getAugmentEnabled();
 
         public static class Config {
             public final long bucketDuration;
@@ -234,12 +242,17 @@
     private final DropBoxNonMonotonicObserver mNonMonotonicObserver =
             new DropBoxNonMonotonicObserver();
 
+    @GuardedBy("mStatsLock")
     private NetworkStatsRecorder mDevRecorder;
+    @GuardedBy("mStatsLock")
     private NetworkStatsRecorder mXtRecorder;
+    @GuardedBy("mStatsLock")
     private NetworkStatsRecorder mUidRecorder;
+    @GuardedBy("mStatsLock")
     private NetworkStatsRecorder mUidTagRecorder;
 
     /** Cached {@link #mXtRecorder} stats. */
+    @GuardedBy("mStatsLock")
     private NetworkStatsCollection mXtStatsCached;
 
     /** Current counter sets for each UID. */
@@ -321,15 +334,15 @@
             return;
         }
 
-        // create data recorders along with historical rotators
-        mDevRecorder = buildRecorder(PREFIX_DEV, mSettings.getDevConfig(), false);
-        mXtRecorder = buildRecorder(PREFIX_XT, mSettings.getXtConfig(), false);
-        mUidRecorder = buildRecorder(PREFIX_UID, mSettings.getUidConfig(), false);
-        mUidTagRecorder = buildRecorder(PREFIX_UID_TAG, mSettings.getUidTagConfig(), true);
-
-        updatePersistThresholds();
-
         synchronized (mStatsLock) {
+            // create data recorders along with historical rotators
+            mDevRecorder = buildRecorder(PREFIX_DEV, mSettings.getDevConfig(), false);
+            mXtRecorder = buildRecorder(PREFIX_XT, mSettings.getXtConfig(), false);
+            mUidRecorder = buildRecorder(PREFIX_UID, mSettings.getUidConfig(), false);
+            mUidTagRecorder = buildRecorder(PREFIX_UID_TAG, mSettings.getUidTagConfig(), true);
+
+            updatePersistThresholdsLocked();
+
             // upgrade any legacy stats, migrating them to rotated files
             maybeUpgradeLegacyStatsLocked();
 
@@ -467,18 +480,20 @@
 
     @Override
     public INetworkStatsSession openSession() {
-        return createSession(null, /* poll on create */ false);
+        // NOTE: if callers want to get non-augmented data, they should go
+        // through the public API
+        return openSessionInternal(NetworkStatsManager.FLAG_AUGMENT_WITH_SUBSCRIPTION_PLAN, null);
     }
 
     @Override
-    public INetworkStatsSession openSessionForUsageStats(final String callingPackage) {
-        return createSession(callingPackage, /* poll on create */ true);
+    public INetworkStatsSession openSessionForUsageStats(int flags, String callingPackage) {
+        return openSessionInternal(flags, callingPackage);
     }
 
-    private INetworkStatsSession createSession(final String callingPackage, boolean pollOnCreate) {
+    private INetworkStatsSession openSessionInternal(final int flags, final String callingPackage) {
         assertBandwidthControlEnabled();
 
-        if (pollOnCreate) {
+        if ((flags & NetworkStatsManager.FLAG_POLL_ON_OPEN) != 0) {
             final long ident = Binder.clearCallingIdentity();
             try {
                 performPoll(FLAG_PERSIST_ALL);
@@ -491,9 +506,13 @@
         // for its lifetime; when caller closes only weak references remain.
 
         return new INetworkStatsSession.Stub() {
+            private final int mCallingUid = Binder.getCallingUid();
+            private final String mCallingPackage = callingPackage;
+            private final @NetworkStatsAccess.Level int mAccessLevel = checkAccessLevel(
+                    callingPackage);
+
             private NetworkStatsCollection mUidComplete;
             private NetworkStatsCollection mUidTagComplete;
-            private String mCallingPackage = callingPackage;
 
             private NetworkStatsCollection getUidComplete() {
                 synchronized (mStatsLock) {
@@ -515,55 +534,38 @@
 
             @Override
             public int[] getRelevantUids() {
-                return getUidComplete().getRelevantUids(checkAccessLevel(mCallingPackage));
+                return getUidComplete().getRelevantUids(mAccessLevel);
             }
 
             @Override
-            public NetworkStats getDeviceSummaryForNetwork(NetworkTemplate template, long start,
-                    long end) {
-                @NetworkStatsAccess.Level int accessLevel = checkAccessLevel(mCallingPackage);
-                if (accessLevel < NetworkStatsAccess.Level.DEVICESUMMARY) {
-                    throw new SecurityException("Calling package " + mCallingPackage
-                            + " cannot access device summary network stats");
-                }
-                NetworkStats result = new NetworkStats(end - start, 1);
-                final long ident = Binder.clearCallingIdentity();
-                try {
-                    // Using access level higher than the one we checked for above.
-                    // Reason is that we are combining usage data in a way that is not PII
-                    // anymore.
-                    result.combineAllValues(
-                            internalGetSummaryForNetwork(template, start, end,
-                                    NetworkStatsAccess.Level.DEVICE));
-                } finally {
-                    Binder.restoreCallingIdentity(ident);
-                }
-                return result;
+            public NetworkStats getDeviceSummaryForNetwork(
+                    NetworkTemplate template, long start, long end) {
+                return internalGetSummaryForNetwork(template, flags, start, end, mAccessLevel,
+                        mCallingUid);
             }
 
             @Override
             public NetworkStats getSummaryForNetwork(
                     NetworkTemplate template, long start, long end) {
-                @NetworkStatsAccess.Level int accessLevel = checkAccessLevel(mCallingPackage);
-                return internalGetSummaryForNetwork(template, start, end, accessLevel);
+                return internalGetSummaryForNetwork(template, flags, start, end, mAccessLevel,
+                        mCallingUid);
             }
 
             @Override
             public NetworkStatsHistory getHistoryForNetwork(NetworkTemplate template, int fields) {
-                @NetworkStatsAccess.Level int accessLevel = checkAccessLevel(mCallingPackage);
-                return internalGetHistoryForNetwork(template, fields, accessLevel);
+                return internalGetHistoryForNetwork(template, flags, fields, mAccessLevel,
+                        mCallingUid);
             }
 
             @Override
             public NetworkStats getSummaryForAllUid(
                     NetworkTemplate template, long start, long end, boolean includeTags) {
                 try {
-                    @NetworkStatsAccess.Level int accessLevel = checkAccessLevel(mCallingPackage);
-                    final NetworkStats stats =
-                            getUidComplete().getSummary(template, start, end, accessLevel);
+                    final NetworkStats stats = getUidComplete()
+                            .getSummary(template, start, end, mAccessLevel, mCallingUid);
                     if (includeTags) {
                         final NetworkStats tagStats = getUidTagComplete()
-                                .getSummary(template, start, end, accessLevel);
+                                .getSummary(template, start, end, mAccessLevel, mCallingUid);
                         stats.combineAllValues(tagStats);
                     }
                     return stats;
@@ -577,13 +579,13 @@
             @Override
             public NetworkStatsHistory getHistoryForUid(
                     NetworkTemplate template, int uid, int set, int tag, int fields) {
-                @NetworkStatsAccess.Level int accessLevel = checkAccessLevel(mCallingPackage);
+                // NOTE: We don't augment UID-level statistics
                 if (tag == TAG_NONE) {
-                    return getUidComplete().getHistory(template, uid, set, tag, fields,
-                            accessLevel);
+                    return getUidComplete().getHistory(template, null, uid, set, tag, fields,
+                            Long.MIN_VALUE, Long.MAX_VALUE, mAccessLevel, mCallingUid);
                 } else {
-                    return getUidTagComplete().getHistory(template, uid, set, tag, fields,
-                            accessLevel);
+                    return getUidTagComplete().getHistory(template, null, uid, set, tag, fields,
+                            Long.MIN_VALUE, Long.MAX_VALUE, mAccessLevel, mCallingUid);
                 }
             }
 
@@ -591,13 +593,13 @@
             public NetworkStatsHistory getHistoryIntervalForUid(
                     NetworkTemplate template, int uid, int set, int tag, int fields,
                     long start, long end) {
-                @NetworkStatsAccess.Level int accessLevel = checkAccessLevel(mCallingPackage);
+                // NOTE: We don't augment UID-level statistics
                 if (tag == TAG_NONE) {
-                    return getUidComplete().getHistory(template, uid, set, tag, fields, start, end,
-                            accessLevel);
+                    return getUidComplete().getHistory(template, null, uid, set, tag, fields,
+                            start, end, mAccessLevel, mCallingUid);
                 } else if (uid == Binder.getCallingUid()) {
-                    return getUidTagComplete().getHistory(template, uid, set, tag, fields,
-                            start, end, accessLevel);
+                    return getUidTagComplete().getHistory(template, null, uid, set, tag, fields,
+                            start, end, mAccessLevel, mCallingUid);
                 } else {
                     throw new SecurityException("Calling package " + mCallingPackage
                             + " cannot access tag information from a different uid");
@@ -618,36 +620,87 @@
     }
 
     /**
+     * Find the most relevant {@link SubscriptionPlan} for the given
+     * {@link NetworkTemplate} and flags. This is typically used to augment
+     * local measurement results to match a known anchor from the carrier.
+     */
+    private SubscriptionPlan resolveSubscriptionPlan(NetworkTemplate template, int flags) {
+        SubscriptionPlan plan = null;
+        if ((flags & NetworkStatsManager.FLAG_AUGMENT_WITH_SUBSCRIPTION_PLAN) != 0
+                && (template.getMatchRule() == NetworkTemplate.MATCH_MOBILE_ALL)
+                && mSettings.getAugmentEnabled()) {
+            Slog.d(TAG, "Resolving plan for " + template);
+            final long token = Binder.clearCallingIdentity();
+            try {
+                final SubscriptionManager sm = mContext.getSystemService(SubscriptionManager.class);
+                final TelephonyManager tm = mContext.getSystemService(TelephonyManager.class);
+                for (int subId : sm.getActiveSubscriptionIdList()) {
+                    if (template.matchesSubscriberId(tm.getSubscriberId(subId))) {
+                        Slog.d(TAG, "Found active matching subId " + subId);
+                        final List<SubscriptionPlan> plans = sm.getSubscriptionPlans(subId);
+                        if (!plans.isEmpty()) {
+                            plan = plans.get(0);
+                        }
+                    }
+                }
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
+            Slog.d(TAG, "Resolved to plan " + plan);
+        }
+        return plan;
+    }
+
+    /**
      * Return network summary, splicing between DEV and XT stats when
      * appropriate.
      */
-    private NetworkStats internalGetSummaryForNetwork(
-            NetworkTemplate template, long start, long end,
-            @NetworkStatsAccess.Level int accessLevel) {
+    private NetworkStats internalGetSummaryForNetwork(NetworkTemplate template, int flags,
+            long start, long end, @NetworkStatsAccess.Level int accessLevel, int callingUid) {
         // We've been using pure XT stats long enough that we no longer need to
         // splice DEV and XT together.
-        return mXtStatsCached.getSummary(template, start, end, accessLevel);
+        final NetworkStatsHistory history = internalGetHistoryForNetwork(template, flags, FIELD_ALL,
+                accessLevel, callingUid);
+
+        final long now = System.currentTimeMillis();
+        final NetworkStatsHistory.Entry entry = history.getValues(start, end, now, null);
+
+        final NetworkStats stats = new NetworkStats(end - start, 1);
+        stats.addValues(new NetworkStats.Entry(IFACE_ALL, UID_ALL, SET_ALL, TAG_NONE, METERED_ALL,
+                ROAMING_ALL, entry.rxBytes, entry.rxPackets, entry.txBytes, entry.txPackets,
+                entry.operations));
+        return stats;
     }
 
     /**
      * Return network history, splicing between DEV and XT stats when
      * appropriate.
      */
-    private NetworkStatsHistory internalGetHistoryForNetwork(NetworkTemplate template, int fields,
-            @NetworkStatsAccess.Level int accessLevel) {
+    private NetworkStatsHistory internalGetHistoryForNetwork(NetworkTemplate template,
+            int flags, int fields, @NetworkStatsAccess.Level int accessLevel, int callingUid) {
         // We've been using pure XT stats long enough that we no longer need to
         // splice DEV and XT together.
-        return mXtStatsCached.getHistory(template, UID_ALL, SET_ALL, TAG_NONE, fields, accessLevel);
+        final SubscriptionPlan augmentPlan = resolveSubscriptionPlan(template, flags);
+        synchronized (mStatsLock) {
+            return mXtStatsCached.getHistory(template, augmentPlan,
+                    UID_ALL, SET_ALL, TAG_NONE, fields, Long.MIN_VALUE, Long.MAX_VALUE,
+                    accessLevel, callingUid);
+        }
     }
 
     @Override
     public long getNetworkTotalBytes(NetworkTemplate template, long start, long end) {
-        // Special case - since this is for internal use only, don't worry about a full access level
-        // check and just require the signature/privileged permission.
+        // Special case - since this is for internal use only, don't worry about
+        // a full access level check and just require the signature/privileged
+        // permission.
         mContext.enforceCallingOrSelfPermission(READ_NETWORK_USAGE_HISTORY, TAG);
         assertBandwidthControlEnabled();
-        return internalGetSummaryForNetwork(template, start, end, NetworkStatsAccess.Level.DEVICE)
-                .getTotalBytes();
+
+        // NOTE: if callers want to get non-augmented data, they should go
+        // through the public API
+        return internalGetSummaryForNetwork(template,
+                NetworkStatsManager.FLAG_AUGMENT_WITH_SUBSCRIPTION_PLAN, start, end,
+                NetworkStatsAccess.Level.DEVICE, Binder.getCallingUid()).getTotalBytes();
     }
 
     @Override
@@ -691,7 +744,8 @@
     @Override
     public void incrementOperationCount(int uid, int tag, int operationCount) {
         if (Binder.getCallingUid() != uid) {
-            mContext.enforceCallingOrSelfPermission(MODIFY_NETWORK_ACCOUNTING, TAG);
+            mContext.enforceCallingOrSelfPermission(
+                    android.Manifest.permission.UPDATE_DEVICE_STATS, TAG);
         }
 
         if (operationCount < 0) {
@@ -712,7 +766,7 @@
 
     @Override
     public void setUidForeground(int uid, boolean uidForeground) {
-        mContext.enforceCallingOrSelfPermission(MODIFY_NETWORK_ACCOUNTING, TAG);
+        mContext.enforceCallingOrSelfPermission(CONNECTIVITY_INTERNAL, TAG);
 
         synchronized (mStatsLock) {
             final int set = uidForeground ? SET_FOREGROUND : SET_DEFAULT;
@@ -752,7 +806,7 @@
 
     @Override
     public void advisePersistThreshold(long thresholdBytes) {
-        mContext.enforceCallingOrSelfPermission(MODIFY_NETWORK_ACCOUNTING, TAG);
+        mContext.enforceCallingOrSelfPermission(CONNECTIVITY_INTERNAL, TAG);
         assertBandwidthControlEnabled();
 
         // clamp threshold into safe range
@@ -768,7 +822,7 @@
         synchronized (mStatsLock) {
             if (!mSystemReady) return;
 
-            updatePersistThresholds();
+            updatePersistThresholdsLocked();
 
             mDevRecorder.maybePersistLocked(currentTime);
             mXtRecorder.maybePersistLocked(currentTime);
@@ -824,7 +878,7 @@
      * reflect current {@link #mPersistThreshold} value. Always defers to
      * {@link Global} values when defined.
      */
-    private void updatePersistThresholds() {
+    private void updatePersistThresholdsLocked() {
         mDevRecorder.setPersistThreshold(mSettings.getDevPersistBytes(mPersistThreshold));
         mXtRecorder.setPersistThreshold(mSettings.getXtPersistBytes(mPersistThreshold));
         mUidRecorder.setPersistThreshold(mSettings.getUidPersistBytes(mPersistThreshold));
@@ -1266,7 +1320,7 @@
         synchronized (mStatsLock) {
             if (args.length > 0 && "--proto".equals(args[0])) {
                 // In this case ignore all other arguments.
-                dumpProto(fd);
+                dumpProtoLocked(fd);
                 return;
             }
 
@@ -1342,7 +1396,7 @@
         }
     }
 
-    private void dumpProto(FileDescriptor fd) {
+    private void dumpProtoLocked(FileDescriptor fd) {
         final ProtoOutputStream proto = new ProtoOutputStream(fd);
 
         // TODO Right now it writes all history.  Should it limit to the "since-boot" log?
@@ -1530,6 +1584,10 @@
             return getGlobalBoolean(NETSTATS_SAMPLE_ENABLED, true);
         }
         @Override
+        public boolean getAugmentEnabled() {
+            return getGlobalBoolean(NETSTATS_AUGMENT_ENABLED, true);
+        }
+        @Override
         public Config getDevConfig() {
             return new Config(getGlobalLong(NETSTATS_DEV_BUCKET_DURATION, HOUR_IN_MILLIS),
                     getGlobalLong(NETSTATS_DEV_ROTATE_AGE, 15 * DAY_IN_MILLIS),