2 files changed, 34 insertions, 2 deletions
diff --git a/wifi/java/android/net/wifi/IWifiManager.aidl b/wifi/java/android/net/wifi/IWifiManager.aidl
index 6e0bc9d..0ee559e 100644
--- a/wifi/java/android/net/wifi/IWifiManager.aidl
+++ b/wifi/java/android/net/wifi/IWifiManager.aidl
@@ -21,6 +21,8 @@ import android.net.wifi.WifiConfiguration;
 import android.net.wifi.ScanResult;
 import android.net.DhcpInfo;
 
+import android.os.WorkSource;
+
 /**
  * Interface that allows controlling and querying Wi-Fi connectivity.
  *
@@ -66,7 +68,9 @@ interface IWifiManager
 
     DhcpInfo getDhcpInfo();
 
-    boolean acquireWifiLock(IBinder lock, int lockType, String tag);
+    boolean acquireWifiLock(IBinder lock, int lockType, String tag, in WorkSource ws);
+
+    void updateWifiLockWorkSource(IBinder lock, in WorkSource ws);
 
     boolean releaseWifiLock(IBinder lock);
 
diff --git a/wifi/java/android/net/wifi/WifiManager.java b/wifi/java/android/net/wifi/WifiManager.java
index 9d21521..dd162f2 100644
--- a/wifi/java/android/net/wifi/WifiManager.java
+++ b/wifi/java/android/net/wifi/WifiManager.java
@@ -23,6 +23,7 @@ import android.os.Binder;
 import android.os.IBinder;
 import android.os.Handler;
 import android.os.RemoteException;
+import android.os.WorkSource;
 
 import java.util.List;
 
@@ -872,6 +873,7 @@ public class WifiManager {
         int mLockType;
         private boolean mRefCounted;
         private boolean mHeld;
+        private WorkSource mWorkSource;
 
         private WifiLock(int lockType, String tag) {
             mTag = tag;
@@ -897,7 +899,7 @@ public class WifiManager {
             synchronized (mBinder) {
                 if (mRefCounted ? (++mRefCount > 0) : (!mHeld)) {
                     try {
-                        mService.acquireWifiLock(mBinder, mLockType, mTag);
+                        mService.acquireWifiLock(mBinder, mLockType, mTag, mWorkSource);
                         synchronized (WifiManager.this) {
                             if (mActiveLockCount >= MAX_ACTIVE_LOCKS) {
                                 mService.releaseWifiLock(mBinder);
@@ -969,6 +971,32 @@ public class WifiManager {
             }
         }
 
+        public void setWorkSource(WorkSource ws) {
+            synchronized (mBinder) {
+                if (ws != null && ws.size() == 0) {
+                    ws = null;
+                }
+                boolean changed = true;
+                if (ws == null) {
+                    mWorkSource = null;
+                } else if (mWorkSource == null) {
+                    changed = mWorkSource != null;
+                    mWorkSource = new WorkSource(ws);
+                } else {
+                    changed = mWorkSource.diff(ws);
+                    if (changed) {
+                        mWorkSource.set(ws);
+                    }
+                }
+                if (changed && mHeld) {
+                    try {
+                        mService.updateWifiLockWorkSource(mBinder, mWorkSource);
+                    } catch (RemoteException e) {
+                    }
+                }
+            }
+        }
+
         public String toString() {
             String s1, s2, s3;
             synchronized (mBinder) {