Keep track of Calling UID for WifiLock acquisition

In current implementation, when processing a request to acquire a WifiLock,
the calling UID is lost when switching context to the wifi thread. This
results in acquisition is attributed to the wifi handler thread.

This commit keeps track of the calling UID so attribution is done
correctly.

Bug: 130798029
Test: atest com.android.wifi.server
Test: cts-tradefed run cts-dev -t android.cts.statsd.atom.UidAtomTests#testWifiLockHighPerf
Test: cts-tradefed run cts-dev -t android.cts.statsd.atom.UidAtomTests#testWifiLockLowLatency
Change-Id: I37f714ee2bfcf2c4cf47b8c1821837b98a1d187b

1 files changed, 11 insertions, 2 deletions

diff --git a/service/java/com/android/server/wifi/WifiServiceImpl.java b/service/java/com/android/server/wifi/WifiServiceImpl.java
index 82b7bab97..2f2eeaf20 100644
--- a/service/java/com/android/server/wifi/WifiServiceImpl.java
+++ b/service/java/com/android/server/wifi/WifiServiceImpl.java
@@ -2816,10 +2816,15 @@ public class WifiServiceImpl extends BaseWifiService {
                 .c(Binder.getCallingUid())
                 .c(lockMode).flush();
 
+        // If no UID is provided in worksource, use the calling UID
+        WorkSource updatedWs = (ws == null || ws.isEmpty())
+                ? new WorkSource(Binder.getCallingUid()) : ws;
+
         Mutable<Boolean> lockSuccess = new Mutable<>();
         boolean runWithScissorsSuccess = mWifiInjector.getClientModeImplHandler().runWithScissors(
                 () -> {
-                    lockSuccess.value = mWifiLockManager.acquireWifiLock(lockMode, tag, binder, ws);
+                    lockSuccess.value = mWifiLockManager.acquireWifiLock(
+                            lockMode, tag, binder, updatedWs);
                 }, RUN_WITH_SCISSORS_TIMEOUT_MILLIS);
         if (!runWithScissorsSuccess) {
             Log.e(TAG, "Failed to post runnable to acquireWifiLock");
@@ -2833,9 +2838,13 @@ public class WifiServiceImpl extends BaseWifiService {
     public void updateWifiLockWorkSource(IBinder binder, WorkSource ws) {
         mLog.info("updateWifiLockWorkSource uid=%").c(Binder.getCallingUid()).flush();
 
+        // If no UID is provided in worksource, use the calling UID
+        WorkSource updatedWs = (ws == null || ws.isEmpty())
+                ? new WorkSource(Binder.getCallingUid()) : ws;
+
         boolean runWithScissorsSuccess = mWifiInjector.getClientModeImplHandler().runWithScissors(
                 () -> {
-                    mWifiLockManager.updateWifiLockWorkSource(binder, ws);
+                    mWifiLockManager.updateWifiLockWorkSource(binder, updatedWs);
                 }, RUN_WITH_SCISSORS_TIMEOUT_MILLIS);
         if (!runWithScissorsSuccess) {
             Log.e(TAG, "Failed to post runnable to updateWifiLockWorkSource");