Add optimistic locking to account updates

service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java

@@ -439,12 +439,12 @@ public class AccountController {
       return;
     }
 
-    device.setApnId(null);
-    device.setVoipApnId(null);
-    device.setGcmId(registrationId.getGcmRegistrationId());
-    device.setFetchesMessages(false);
-
-    accounts.update(account);
+    account = accounts.updateDevice(account, device.getId(), d -> {
+      d.setApnId(null);
+      d.setVoipApnId(null);
+      d.setGcmId(registrationId.getGcmRegistrationId());
+      d.setFetchesMessages(false);
+    });
 
     if (!wasAccountEnabled && account.isEnabled()) {
       directoryQueue.refreshRegisteredUser(account);
@@ -457,11 +457,12 @@ public class AccountController {
   public void deleteGcmRegistrationId(@Auth DisabledPermittedAccount disabledPermittedAccount) {
     Account account = disabledPermittedAccount.getAccount();
     Device  device  = account.getAuthenticatedDevice().get();
-    device.setGcmId(null);
-    device.setFetchesMessages(false);
-    device.setUserAgent("OWA");
 
-    accounts.update(account);
+    account = accounts.updateDevice(account, device.getId(), d -> {
+      d.setGcmId(null);
+      d.setFetchesMessages(false);
+      d.setUserAgent("OWA");
+    });
     directoryQueue.refreshRegisteredUser(account);
   }
 
@@ -474,11 +475,12 @@ public class AccountController {
     Device  device            = account.getAuthenticatedDevice().get();
     boolean wasAccountEnabled = account.isEnabled();
 
-    device.setApnId(registrationId.getApnRegistrationId());
-    device.setVoipApnId(registrationId.getVoipRegistrationId());
-    device.setGcmId(null);
-    device.setFetchesMessages(false);
-    accounts.update(account);
+    account = accounts.updateDevice(account, device.getId(), d -> {
+      d.setApnId(registrationId.getApnRegistrationId());
+      d.setVoipApnId(registrationId.getVoipRegistrationId());
+      d.setGcmId(null);
+      d.setFetchesMessages(false);
+    });
 
     if (!wasAccountEnabled && account.isEnabled()) {
       directoryQueue.refreshRegisteredUser(account);
@@ -491,15 +493,16 @@ public class AccountController {
   public void deleteApnRegistrationId(@Auth DisabledPermittedAccount disabledPermittedAccount) {
     Account account = disabledPermittedAccount.getAccount();
     Device  device  = account.getAuthenticatedDevice().get();
-    device.setApnId(null);
-    device.setFetchesMessages(false);
-    if (device.getId() == 1) {
-      device.setUserAgent("OWI");
-    } else {
-      device.setUserAgent("OWP");
-    }
 
-    accounts.update(account);
+    accounts.updateDevice(account, device.getId(), d -> {
+      d.setApnId(null);
+      d.setFetchesMessages(false);
+      if (d.getId() == 1) {
+        d.setUserAgent("OWI");
+      } else {
+        d.setUserAgent("OWP");
+      }
+    });
     directoryQueue.refreshRegisteredUser(account);
   }
 
@@ -509,18 +512,18 @@ public class AccountController {
   @Path("/registration_lock")
   public void setRegistrationLock(@Auth Account account, @Valid RegistrationLock accountLock) {
     AuthenticationCredentials credentials = new AuthenticationCredentials(accountLock.getRegistrationLock());
-    account.setRegistrationLock(credentials.getHashedAuthenticationToken(), credentials.getSalt());
-    account.setPin(null);
 
-    accounts.update(account);
+    accounts.update(account, a -> {
+      a.setRegistrationLock(credentials.getHashedAuthenticationToken(), credentials.getSalt());
+      a.setPin(null);
+    });
   }
 
   @Timed
   @DELETE
   @Path("/registration_lock")
   public void removeRegistrationLock(@Auth Account account) {
-    account.setRegistrationLock(null, null);
-    accounts.update(account);
+    accounts.update(account, a -> a.setRegistrationLock(null, null));
   }
 
   @Timed
@@ -531,21 +534,21 @@ public class AccountController {
     // TODO Remove once PIN-based reglocks have been deprecated
     logger.info("PIN set by User-Agent: {}", userAgent);
 
-    account.setPin(accountLock.getPin());
-    account.setRegistrationLock(null, null);
-
-    accounts.update(account);
+    accounts.update(account, a -> {
+      a.setPin(accountLock.getPin());
+      a.setRegistrationLock(null, null);
+    });
   }
 
   @Timed
   @DELETE
   @Path("/pin/")
+
   public void removePin(@Auth Account account, @HeaderParam("User-Agent") String userAgent) {
     // TODO Remove once PIN-based reglocks have been deprecated
     logger.info("PIN removed by User-Agent: {}", userAgent);
 
-    account.setPin(null);
-    accounts.update(account);
+    accounts.update(account, a -> a.setPin(null));
   }
 
   @Timed
@@ -553,8 +556,8 @@ public class AccountController {
   @Path("/name/")
   public void setName(@Auth DisabledPermittedAccount disabledPermittedAccount, @Valid DeviceName deviceName) {
     Account account = disabledPermittedAccount.getAccount();
-    account.getAuthenticatedDevice().get().setName(deviceName.getDeviceName());
-    accounts.update(account);
+    Device device = account.getAuthenticatedDevice().get();
+    accounts.updateDevice(account, device.getId(), d -> d.setName(deviceName.getDeviceName()));
   }
 
   @Timed
@@ -572,25 +575,29 @@ public class AccountController {
                                    @Valid AccountAttributes attributes)
   {
     Account account = disabledPermittedAccount.getAccount();
-    Device  device  = account.getAuthenticatedDevice().get();
+    long deviceId = account.getAuthenticatedDevice().get().getId();
 
-    device.setFetchesMessages(attributes.getFetchesMessages());
-    device.setName(attributes.getName());
-    device.setLastSeen(Util.todayInMillis());
-    device.setCapabilities(attributes.getCapabilities());
-    device.setRegistrationId(attributes.getRegistrationId());
-    device.setUserAgent(userAgent);
+    account = accounts.update(account, a-> {
 
-    setAccountRegistrationLockFromAttributes(account, attributes);
+      a.getDevice(deviceId).ifPresent(d -> {
+            d.setFetchesMessages(attributes.getFetchesMessages());
+            d.setName(attributes.getName());
+            d.setLastSeen(Util.todayInMillis());
+            d.setCapabilities(attributes.getCapabilities());
+            d.setRegistrationId(attributes.getRegistrationId());
+            d.setUserAgent(userAgent);
+          });
+
+      setAccountRegistrationLockFromAttributes(a, attributes);
+
+      a.setUnidentifiedAccessKey(attributes.getUnidentifiedAccessKey());
+      a.setUnrestrictedUnidentifiedAccess(attributes.isUnrestrictedUnidentifiedAccess());
+      a.setDiscoverableByPhoneNumber(attributes.isDiscoverableByPhoneNumber());
+
+    });
 
     final boolean hasDiscoverabilityChange = (account.isDiscoverableByPhoneNumber() != attributes.isDiscoverableByPhoneNumber());
 
-    account.setUnidentifiedAccessKey(attributes.getUnidentifiedAccessKey());
-    account.setUnrestrictedUnidentifiedAccess(attributes.isUnrestrictedUnidentifiedAccess());
-    account.setDiscoverableByPhoneNumber(attributes.isDiscoverableByPhoneNumber());
-
-    accounts.update(account);
-
     if (hasDiscoverabilityChange) {
       directoryQueue.refreshRegisteredUser(account);
     }

service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java

@@ -100,8 +100,7 @@ public class DeviceController {
     }
 
     messages.clear(account.getUuid(), deviceId);
-    account.removeDevice(deviceId);
-    accounts.update(account);
+    account = accounts.update(account, a -> a.removeDevice(deviceId));
     directoryQueue.refreshRegisteredUser(account);
     // ensure any messages that came in after the first clear() are also removed
     messages.clear(account.getUuid(), deviceId);
@@ -192,15 +191,16 @@ public class DeviceController {
       device.setName(accountAttributes.getName());
       device.setAuthenticationCredentials(new AuthenticationCredentials(password));
       device.setFetchesMessages(accountAttributes.getFetchesMessages());
-      device.setId(account.get().getNextDeviceId());
       device.setRegistrationId(accountAttributes.getRegistrationId());
       device.setLastSeen(Util.todayInMillis());
       device.setCreated(System.currentTimeMillis());
       device.setCapabilities(accountAttributes.getCapabilities());
 
-      account.get().addDevice(device);
-      messages.clear(account.get().getUuid(), device.getId());
-      accounts.update(account.get());
+      accounts.update(account.get(), a -> {
+        device.setId(account.get().getNextDeviceId());
+        messages.clear(account.get().getUuid(), device.getId());
+        a.addDevice(device);
+      });;
 
       pendingDevices.remove(number);
 
@@ -224,8 +224,8 @@ public class DeviceController {
   @Path("/capabilities")
   public void setCapabiltities(@Auth Account account, @Valid DeviceCapabilities capabilities) {
     assert(account.getAuthenticatedDevice().isPresent());
-    account.getAuthenticatedDevice().get().setCapabilities(capabilities);
-    accounts.update(account);
+    final long deviceId = account.getAuthenticatedDevice().get().getId();
+    accounts.updateDevice(account, deviceId, d -> d.setCapabilities(capabilities));
   }
 
   @VisibleForTesting protected VerificationCode generateVerificationCode() {

service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java

@@ -104,17 +104,18 @@ public class KeysController {
     boolean updateAccount     = false;
 
     if (!preKeys.getSignedPreKey().equals(device.getSignedPreKey())) {
-      device.setSignedPreKey(preKeys.getSignedPreKey());
       updateAccount = true;
     }
 
     if (!preKeys.getIdentityKey().equals(account.getIdentityKey())) {
-      account.setIdentityKey(preKeys.getIdentityKey());
       updateAccount = true;
     }
 
     if (updateAccount) {
-      accounts.update(account);
+      account = accounts.update(account, a -> {
+        a.getDevice(device.getId()).ifPresent(d -> d.setSignedPreKey(preKeys.getSignedPreKey()));
+        a.setIdentityKey(preKeys.getIdentityKey());
+      });
 
       if (!wasAccountEnabled && account.isEnabled()) {
         directoryQueue.refreshRegisteredUser(account);
@@ -200,8 +201,7 @@ public class KeysController {
     Device  device            = account.getAuthenticatedDevice().get();
     boolean wasAccountEnabled = account.isEnabled();
 
-    device.setSignedPreKey(signedPreKey);
-    accounts.update(account);
+    account = accounts.updateDevice(account, device.getId(), d -> d.setSignedPreKey(signedPreKey));
 
     if (!wasAccountEnabled && account.isEnabled()) {
       directoryQueue.refreshRegisteredUser(account);

service/src/main/java/org/whispersystems/textsecuregcm/controllers/ProfileController.java

@@ -156,10 +156,11 @@ public class ProfileController {
       response = Optional.of(generateAvatarUploadForm(avatar));
     }
 
-    account.setProfileName(request.getName());
-    account.setAvatar(avatar);
-    account.setCurrentProfileVersion(request.getVersion());
-    accountsManager.update(account);
+    accountsManager.update(account, a -> {
+      a.setProfileName(request.getName());
+      a.setAvatar(avatar);
+      a.setCurrentProfileVersion(request.getVersion());
+    });
 
     if (response.isPresent()) return Response.ok(response).build();
     else                      return Response.ok().build();
@@ -317,8 +318,7 @@ public class ProfileController {
   @Produces(MediaType.APPLICATION_JSON)
   @Path("/name/{name}")
   public void setProfile(@Auth Account account, @PathParam("name") @ExactlySize(value = {72, 108}, payload = {Unwrapping.Unwrap.class}) Optional<String> name) {
-    account.setProfileName(name.orElse(null));
-    accountsManager.update(account);
+    accountsManager.update(account, a -> a.setProfileName(name.orElse(null)));
   }
 
   @Deprecated
@@ -382,8 +382,7 @@ public class ProfileController {
           .build());
     }
 
-    account.setAvatar(objectName);
-    accountsManager.update(account);
+    accountsManager.update(account, a -> a.setAvatar(objectName));
 
     return profileAvatarUploadAttributes;
   }