diff --git a/services/java/com/android/server/WifiService.java b/services/java/com/android/server/WifiService.java
index 3d95bf0..a63b3d8 100644
--- a/services/java/com/android/server/WifiService.java
+++ b/services/java/com/android/server/WifiService.java
@@ -2163,12 +2163,12 @@
             if (DBG) throw new IllegalArgumentException("lockMode=" + lockMode);
             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());
         }
@@ -2227,17 +2227,18 @@
                 ++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) {
@@ -2283,9 +2284,9 @@
 
         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:
@@ -2298,13 +2299,16 @@
                         ++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;
     }