Add WorkSource capability to WifiManager for starting a scan.

Change-Id: I011bf3b91fbddaba7ab0128bb03d27d90b8a886f

3 files changed, 31 insertions, 13 deletions

diff --git a/wifi/java/android/net/wifi/IWifiManager.aidl b/wifi/java/android/net/wifi/IWifiManager.aidl
index 6a4bd45..8103e84 100644
--- a/wifi/java/android/net/wifi/IWifiManager.aidl
+++ b/wifi/java/android/net/wifi/IWifiManager.aidl
@@ -43,7 +43,7 @@ interface IWifiManager
 
     boolean pingSupplicant();
 
-    void startScan();
+    void startScan(in WorkSource ws);
 
     List<ScanResult> getScanResults(String callingPackage);
 
diff --git a/wifi/java/android/net/wifi/WifiManager.java b/wifi/java/android/net/wifi/WifiManager.java
index 86ca4a7..6793710 100644
--- a/wifi/java/android/net/wifi/WifiManager.java
+++ b/wifi/java/android/net/wifi/WifiManager.java
@@ -759,7 +759,18 @@ public class WifiManager {
      */
     public boolean startScan() {
         try {
-            mService.startScan();
+            final WorkSource workSource = null;
+            mService.startScan(workSource);
+            return true;
+        } catch (RemoteException e) {
+            return false;
+        }
+    }
+
+    /** @hide */
+    public boolean startScan(WorkSource workSource) {
+        try {
+            mService.startScan(workSource);
             return true;
         } catch (RemoteException e) {
             return false;
diff --git a/wifi/java/android/net/wifi/WifiStateMachine.java b/wifi/java/android/net/wifi/WifiStateMachine.java
index 858fbcc..72ebb5b 100644
--- a/wifi/java/android/net/wifi/WifiStateMachine.java
+++ b/wifi/java/android/net/wifi/WifiStateMachine.java
@@ -607,7 +607,8 @@ public class WifiStateMachine extends StateMachine {
                 new BroadcastReceiver() {
                     @Override
                     public void onReceive(Context context, Intent intent) {
-                        startScan(UNKNOWN_SCAN_SOURCE);
+                        final WorkSource workSource = null;
+                        startScan(UNKNOWN_SCAN_SOURCE, workSource);
                     }
                 },
                 new IntentFilter(ACTION_START_SCAN));
@@ -725,15 +726,21 @@ public class WifiStateMachine extends StateMachine {
     }
 
     /**
-     * TODO: doc
+     * Initiate a wifi scan.  If workSource is not null, blame is given to it,
+     * otherwise blame is given to callingUid.
+     *
+     * @param callingUid The uid initiating the wifi scan.  Blame will be given
+     *                   here unless workSource is specified.
+     * @param workSource If not null, blame is given to workSource.
      */
-    public void startScan(int callingUid) {
-        sendMessage(CMD_START_SCAN, callingUid);
+    public void startScan(int callingUid, WorkSource workSource) {
+        sendMessage(CMD_START_SCAN, callingUid, 0, workSource);
     }
 
-    private void noteScanStart(int callingUid) {
-        if (mScanWorkSource == null && callingUid != UNKNOWN_SCAN_SOURCE) {
-            mScanWorkSource = new WorkSource(callingUid);
+    // If workSource is not null, blame is given to it, otherwise blame is given to callingUid.
+    private void noteScanStart(int callingUid, WorkSource workSource) {
+        if (mScanWorkSource == null && (callingUid != UNKNOWN_SCAN_SOURCE || workSource != null)) {
+            mScanWorkSource = workSource != null ? workSource : new WorkSource(callingUid);
             try {
                 mBatteryStats.noteWifiScanStartedFromSource(mScanWorkSource);
             } catch (RemoteException e) {
@@ -2502,7 +2509,7 @@ public class WifiStateMachine extends StateMachine {
         public boolean processMessage(Message message) {
             switch(message.what) {
                 case CMD_START_SCAN:
-                    noteScanStart(message.arg1);
+                    noteScanStart(message.arg1, (WorkSource) message.obj);
                     startScanNative(WifiNative.SCAN_WITH_CONNECTION_SETUP);
                     break;
                 case CMD_SET_COUNTRY_CODE:
@@ -2788,7 +2795,7 @@ public class WifiStateMachine extends StateMachine {
                 // Handle scan. All the connection related commands are
                 // handled only in ConnectModeState
                 case CMD_START_SCAN:
-                    noteScanStart(message.arg1);
+                    noteScanStart(message.arg1, (WorkSource) message.obj);
                     startScanNative(WifiNative.SCAN_WITHOUT_CONNECTION_SETUP);
                     break;
                 default:
@@ -3054,7 +3061,7 @@ public class WifiStateMachine extends StateMachine {
                     break;
                 case CMD_START_SCAN:
                     /* Do not attempt to connect when we are already connected */
-                    noteScanStart(message.arg1);
+                    noteScanStart(message.arg1, (WorkSource) message.obj);
                     startScanNative(WifiNative.SCAN_WITHOUT_CONNECTION_SETUP);
                     break;
                     /* Ignore connection to same network */
@@ -3372,7 +3379,7 @@ public class WifiStateMachine extends StateMachine {
                     if (mP2pConnected.get()) break;
                     if (message.arg1 == mPeriodicScanToken &&
                             mWifiConfigStore.getConfiguredNetworks().size() == 0) {
-                        sendMessage(CMD_START_SCAN, UNKNOWN_SCAN_SOURCE);
+                        sendMessage(CMD_START_SCAN, UNKNOWN_SCAN_SOURCE, 0, (WorkSource) null);
                         sendMessageDelayed(obtainMessage(CMD_NO_NETWORKS_PERIODIC_SCAN,
                                     ++mPeriodicScanToken, 0), mSupplicantScanIntervalMs);
                     }