Add WorkSource capability to LocationManager

Change-Id: I0fbbad0879b87ecc75a503bf7963356595bf4b96
diff --git a/location/java/android/location/LocationRequest.java b/location/java/android/location/LocationRequest.java
index 68f540b..5ebba93 100644
--- a/location/java/android/location/LocationRequest.java
+++ b/location/java/android/location/LocationRequest.java
@@ -19,6 +19,7 @@
 import android.os.Parcel;
 import android.os.Parcelable;
 import android.os.SystemClock;
+import android.os.WorkSource;
 import android.util.TimeUtils;
 
 
@@ -145,6 +146,7 @@
     private long mExpireAt = Long.MAX_VALUE;  // no expiry
     private int mNumUpdates = Integer.MAX_VALUE;  // no expiry
     private float mSmallestDisplacement = 0.0f;    // meters
+    private WorkSource mWorkSource = new WorkSource();
 
     private String mProvider = LocationManager.FUSED_PROVIDER;  // for deprecated APIs that explicitly request a provider
 
@@ -233,6 +235,7 @@
         mNumUpdates = src.mNumUpdates;
         mSmallestDisplacement = src.mSmallestDisplacement;
         mProvider = src.mProvider;
+        mWorkSource = src.mWorkSource;
     }
 
     /**
@@ -493,6 +496,16 @@
         return mSmallestDisplacement;
     }
 
+    /** @hide */
+    public void setWorkSource(WorkSource workSource) {
+        mWorkSource = workSource;
+    }
+
+    /** @hide */
+    public WorkSource getWorkSource() {
+        return mWorkSource;
+    }
+
     private static void checkInterval(long millis) {
         if (millis < 0) {
             throw new IllegalArgumentException("invalid interval: " + millis);
@@ -538,6 +551,8 @@
             request.setSmallestDisplacement(in.readFloat());
             String provider = in.readString();
             if (provider != null) request.setProvider(provider);
+            WorkSource workSource = in.readParcelable(WorkSource.class.getClassLoader());
+            if (workSource != null) request.setWorkSource(workSource);
             return request;
         }
         @Override
@@ -560,6 +575,7 @@
         parcel.writeInt(mNumUpdates);
         parcel.writeFloat(mSmallestDisplacement);
         parcel.writeString(mProvider);
+        parcel.writeParcelable(mWorkSource, 0);
     }
 
     /** @hide */
diff --git a/services/java/com/android/server/LocationManagerService.java b/services/java/com/android/server/LocationManagerService.java
index bde9e1c..d3dcd69 100644
--- a/services/java/com/android/server/LocationManagerService.java
+++ b/services/java/com/android/server/LocationManagerService.java
@@ -458,6 +458,7 @@
 
         final ILocationListener mListener;
         final PendingIntent mPendingIntent;
+        final WorkSource mWorkSource; // WorkSource for battery blame, or null to assign to caller.
         final Object mKey;
 
         final HashMap<String,UpdateRecord> mUpdateRecords = new HashMap<String,UpdateRecord>();
@@ -467,7 +468,7 @@
         PowerManager.WakeLock mWakeLock;
 
         Receiver(ILocationListener listener, PendingIntent intent, int pid, int uid,
-                String packageName) {
+                String packageName, WorkSource workSource) {
             mListener = listener;
             mPendingIntent = intent;
             if (listener != null) {
@@ -479,12 +480,19 @@
             mUid = uid;
             mPid = pid;
             mPackageName = packageName;
+            if (workSource != null && workSource.size() <= 0) {
+                workSource = null;
+            }
+            mWorkSource = workSource;
 
             updateMonitoring(true);
 
             // construct/configure wakelock
             mWakeLock = mPowerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, WAKELOCK_KEY);
-            mWakeLock.setWorkSource(new WorkSource(mUid, mPackageName));
+            if (workSource == null) {
+                workSource = new WorkSource(mUid, mPackageName);
+            }
+            mWakeLock.setWorkSource(workSource);
         }
 
         @Override
@@ -883,6 +891,15 @@
         }
     }
 
+    /**
+     * Throw SecurityException if WorkSource use is not allowed (i.e. can't blame other packages
+     * for battery).
+     */
+    private void checkWorkSourceAllowed() {
+        mContext.enforceCallingOrSelfPermission(
+                android.Manifest.permission.UPDATE_DEVICE_STATS, null);
+    }
+
     public static int resolutionLevelToOp(int allowedResolutionLevel) {
         if (allowedResolutionLevel != RESOLUTION_LEVEL_NONE) {
             if (allowedResolutionLevel == RESOLUTION_LEVEL_COARSE) {
@@ -1124,7 +1141,15 @@
                     if (UserHandle.getUserId(record.mReceiver.mUid) == mCurrentUserId) {
                         LocationRequest locationRequest = record.mRequest;
                         if (locationRequest.getInterval() <= thresholdInterval) {
-                            worksource.add(record.mReceiver.mUid, record.mReceiver.mPackageName);
+                            if (record.mReceiver.mWorkSource != null) {
+                                // Assign blame to another work source.
+                                worksource.add(record.mReceiver.mWorkSource);
+                            } else {
+                                // Assign blame to caller.
+                                worksource.add(
+                                        record.mReceiver.mUid,
+                                        record.mReceiver.mPackageName);
+                            }
                         }
                     }
                 }
@@ -1199,11 +1224,11 @@
     }
 
     private Receiver getReceiverLocked(ILocationListener listener, int pid, int uid,
-            String packageName) {
+            String packageName, WorkSource workSource) {
         IBinder binder = listener.asBinder();
         Receiver receiver = mReceivers.get(binder);
         if (receiver == null) {
-            receiver = new Receiver(listener, null, pid, uid, packageName);
+            receiver = new Receiver(listener, null, pid, uid, packageName, workSource);
             mReceivers.put(binder, receiver);
 
             try {
@@ -1216,10 +1241,11 @@
         return receiver;
     }
 
-    private Receiver getReceiverLocked(PendingIntent intent, int pid, int uid, String packageName) {
+    private Receiver getReceiverLocked(PendingIntent intent, int pid, int uid, String packageName,
+            WorkSource workSource) {
         Receiver receiver = mReceivers.get(intent);
         if (receiver == null) {
-            receiver = new Receiver(null, intent, pid, uid, packageName);
+            receiver = new Receiver(null, intent, pid, uid, packageName, workSource);
             mReceivers.put(intent, receiver);
         }
         return receiver;
@@ -1281,16 +1307,16 @@
     }
 
     private Receiver checkListenerOrIntentLocked(ILocationListener listener, PendingIntent intent,
-            int pid, int uid, String packageName) {
+            int pid, int uid, String packageName, WorkSource workSource) {
         if (intent == null && listener == null) {
             throw new IllegalArgumentException("need either listener or intent");
         } else if (intent != null && listener != null) {
             throw new IllegalArgumentException("cannot register both listener and intent");
         } else if (intent != null) {
             checkPendingIntent(intent);
-            return getReceiverLocked(intent, pid, uid, packageName);
+            return getReceiverLocked(intent, pid, uid, packageName, workSource);
         } else {
-            return getReceiverLocked(listener, pid, uid, packageName);
+            return getReceiverLocked(listener, pid, uid, packageName, workSource);
         }
     }
 
@@ -1302,6 +1328,10 @@
         int allowedResolutionLevel = getCallerAllowedResolutionLevel();
         checkResolutionLevelIsSufficientForProviderUse(allowedResolutionLevel,
                 request.getProvider());
+        WorkSource workSource = request.getWorkSource();
+        if (workSource != null && workSource.size() > 0) {
+            checkWorkSourceAllowed();
+        }
         LocationRequest sanitizedRequest = createSanitizedRequest(request, allowedResolutionLevel);
 
         final int pid = Binder.getCallingPid();
@@ -1315,7 +1345,7 @@
 
             synchronized (mLock) {
                 Receiver recevier = checkListenerOrIntentLocked(listener, intent, pid, uid,
-                        packageName);
+                        packageName, workSource);
                 requestLocationUpdatesLocked(sanitizedRequest, recevier, pid, uid, packageName);
             }
         } finally {
@@ -1364,7 +1394,9 @@
         final int uid = Binder.getCallingUid();
 
         synchronized (mLock) {
-            Receiver receiver = checkListenerOrIntentLocked(listener, intent, pid, uid, packageName);
+            WorkSource workSource = null;
+            Receiver receiver = checkListenerOrIntentLocked(listener, intent, pid, uid, packageName,
+                    workSource);
 
             // providers may use public location API's, need to clear identity
             long identity = Binder.clearCallingIdentity();