diff --git a/services/java/com/android/server/WifiService.java b/services/java/com/android/server/WifiService.java
index 540c7fe406e3b..55b7c9e5824f5 100644
--- a/services/java/com/android/server/WifiService.java
+++ b/services/java/com/android/server/WifiService.java
@@ -1148,12 +1148,12 @@ public class WifiService extends IWifiManager.Stub {
         if (lockMode != WifiManager.WIFI_MODE_FULL && lockMode != WifiManager.WIFI_MODE_SCAN_ONLY) {
             return false;
         }
-        if (ws != null) {
-            enforceWakeSourcePermission(Binder.getCallingUid(), Binder.getCallingPid());
-        }
         if (ws != null && ws.size() == 0) {
             ws = null;
         }
+        if (ws != null) {
+            enforceWakeSourcePermission(Binder.getCallingUid(), Binder.getCallingPid());
+        }
         if (ws == null) {
             ws = new WorkSource(Binder.getCallingUid());
         }
@@ -1201,17 +1201,18 @@ public class WifiService extends IWifiManager.Stub {
                 ++mScanLocksAcquired;
                 break;
             }
+
+            // Be aggressive about adding new locks into the accounted state...
+            // we want to over-report rather than under-report.
+            reportStartWorkSource();
+
+            updateWifiState();
+            return true;
         } catch (RemoteException e) {
+            return false;
         } finally {
             Binder.restoreCallingIdentity(ident);
         }
-
-        // Be aggressive about adding new locks into the accounted state...
-        // we want to over-report rather than under-report.
-        reportStartWorkSource();
-
-        updateWifiState();
-        return true;
     }
 
     public void updateWifiLockWorkSource(IBinder lock, WorkSource ws) {
@@ -1257,9 +1258,9 @@ public class WifiService extends IWifiManager.Stub {
 
         hadLock = (wifiLock != null);
 
-        if (hadLock) {
-            long ident = Binder.clearCallingIdentity();
-            try {
+        long ident = Binder.clearCallingIdentity();
+        try {
+            if (hadLock) {
                 noteAcquireWifiLock(wifiLock);
                 switch(wifiLock.mMode) {
                     case WifiManager.WIFI_MODE_FULL:
@@ -1269,13 +1270,16 @@ public class WifiService extends IWifiManager.Stub {
                         ++mScanLocksReleased;
                         break;
                 }
-            } catch (RemoteException e) {
-            } finally {
-                Binder.restoreCallingIdentity(ident);
             }
+
+            // TODO - should this only happen if you hadLock?
+            updateWifiState();
+
+        } catch (RemoteException e) {
+        } finally {
+            Binder.restoreCallingIdentity(ident);
         }
-        // TODO - should this only happen if you hadLock?
-        updateWifiState();
+
         return hadLock;
     }