Allow primary to set and provide new signed prekeys for linked devices (#950)

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -169,6 +169,7 @@ import org.whispersystems.textsecuregcm.storage.AccountDatabaseCrawlerCache;
 import org.whispersystems.textsecuregcm.storage.AccountDatabaseCrawlerListener;
 import org.whispersystems.textsecuregcm.storage.Accounts;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.ChangeNumberManager;
 import org.whispersystems.textsecuregcm.storage.ContactDiscoveryWriter;
 import org.whispersystems.textsecuregcm.storage.DeletedAccounts;
 import org.whispersystems.textsecuregcm.storage.DeletedAccountsDirectoryReconciler;
@@ -496,6 +497,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         new RateLimitChallengeOptionManager(dynamicRateLimiters, dynamicConfigurationManager);
 
     MessagePersister messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager, dynamicConfigurationManager, Duration.ofMinutes(config.getMessageCacheConfiguration().getPersistDelayMinutes()));
+    ChangeNumberManager changeNumberManager = new ChangeNumberManager(messageSender, accountsManager);
 
     final List<AccountDatabaseCrawlerListener> directoryReconciliationAccountDatabaseCrawlerListeners = new ArrayList<>();
     final List<DeletedAccountsDirectoryReconciler> deletedAccountsDirectoryReconcilers = new ArrayList<>();
@@ -622,8 +624,8 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     environment.jersey().register(
         new AccountController(pendingAccountsManager, accountsManager, abusiveHostRules, rateLimiters,
             smsSender, dynamicConfigurationManager, turnTokenGenerator, config.getTestDevices(),
-            recaptchaClient, gcmSender, apnSender, backupCredentialsGenerator,
-            verifyExperimentEnrollmentManager));
+            recaptchaClient, gcmSender, apnSender, verifyExperimentEnrollmentManager,
+            changeNumberManager, backupCredentialsGenerator));
     environment.jersey().register(new KeysController(rateLimiters, keys, accountsManager));
 
     final List<Object> commonControllers = Lists.newArrayList(

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

@@ -33,6 +33,7 @@ import javax.validation.constraints.NotNull;
 import javax.ws.rs.BadRequestException;
 import javax.ws.rs.Consumes;
 import javax.ws.rs.DELETE;
+import javax.ws.rs.ForbiddenException;
 import javax.ws.rs.GET;
 import javax.ws.rs.HEAD;
 import javax.ws.rs.HeaderParam;
@@ -69,8 +70,11 @@ import org.whispersystems.textsecuregcm.entities.ApnRegistrationId;
 import org.whispersystems.textsecuregcm.entities.ChangePhoneNumberRequest;
 import org.whispersystems.textsecuregcm.entities.DeviceName;
 import org.whispersystems.textsecuregcm.entities.GcmRegistrationId;
+import org.whispersystems.textsecuregcm.entities.IncomingMessage;
+import org.whispersystems.textsecuregcm.entities.MismatchedDevices;
 import org.whispersystems.textsecuregcm.entities.RegistrationLock;
 import org.whispersystems.textsecuregcm.entities.RegistrationLockFailure;
+import org.whispersystems.textsecuregcm.entities.StaleDevices;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
 import org.whispersystems.textsecuregcm.push.APNSender;
@@ -84,6 +88,7 @@ import org.whispersystems.textsecuregcm.storage.AbusiveHostRule;
 import org.whispersystems.textsecuregcm.storage.AbusiveHostRules;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.ChangeNumberManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 import org.whispersystems.textsecuregcm.storage.StoredVerificationCodeManager;
@@ -92,6 +97,7 @@ import org.whispersystems.textsecuregcm.util.Constants;
 import org.whispersystems.textsecuregcm.util.ForwardedIpUtil;
 import org.whispersystems.textsecuregcm.util.Hex;
 import org.whispersystems.textsecuregcm.util.ImpossiblePhoneNumberException;
+import org.whispersystems.textsecuregcm.util.MessageValidation;
 import org.whispersystems.textsecuregcm.util.NonNormalizedPhoneNumberException;
 import org.whispersystems.textsecuregcm.util.Username;
 import org.whispersystems.textsecuregcm.util.Util;
@@ -138,6 +144,7 @@ public class AccountController {
   private final ExternalServiceCredentialGenerator backupServiceCredentialGenerator;
 
   private final TwilioVerifyExperimentEnrollmentManager verifyExperimentEnrollmentManager;
+  private final ChangeNumberManager changeNumberManager;
 
   public AccountController(StoredVerificationCodeManager pendingAccounts,
                            AccountsManager accounts,
@@ -150,8 +157,9 @@ public class AccountController {
                            RecaptchaClient recaptchaClient,
                            GCMSender gcmSender,
                            APNSender apnSender,
-                           ExternalServiceCredentialGenerator backupServiceCredentialGenerator,
-                           TwilioVerifyExperimentEnrollmentManager verifyExperimentEnrollmentManager)
+                           TwilioVerifyExperimentEnrollmentManager verifyExperimentEnrollmentManager,
+                           ChangeNumberManager changeNumberManager,
+                           ExternalServiceCredentialGenerator backupServiceCredentialGenerator)
   {
     this.pendingAccounts                   = pendingAccounts;
     this.accounts                          = accounts;
@@ -164,8 +172,9 @@ public class AccountController {
     this.recaptchaClient = recaptchaClient;
     this.gcmSender                         = gcmSender;
     this.apnSender                         = apnSender;
-    this.backupServiceCredentialGenerator  = backupServiceCredentialGenerator;
     this.verifyExperimentEnrollmentManager = verifyExperimentEnrollmentManager;
+    this.backupServiceCredentialGenerator = backupServiceCredentialGenerator;
+    this.changeNumberManager = changeNumberManager;
   }
 
   @Timed
@@ -403,38 +412,75 @@ public class AccountController {
   public AccountIdentityResponse changeNumber(@Auth final AuthenticatedAccount authenticatedAccount, @NotNull @Valid final ChangePhoneNumberRequest request)
       throws RateLimitExceededException, InterruptedException, ImpossiblePhoneNumberException, NonNormalizedPhoneNumberException {
 
-    final Account updatedAccount;
+    if (!authenticatedAccount.getAuthenticatedDevice().isMaster()) {
+      throw new ForbiddenException();
+    }
 
-    if (request.getNumber().equals(authenticatedAccount.getAccount().getNumber())) {
-      // This may be a request that got repeated due to poor network conditions or other client error; take no action,
-      // but report success since the account is in the desired state
-      updatedAccount = authenticatedAccount.getAccount();
-    } else {
-      Util.requireNormalizedNumber(request.getNumber());
+    if (request.getDeviceSignedPrekeys() != null && !request.getDeviceSignedPrekeys().isEmpty()) {
+      if (request.getDeviceMessages() == null || request.getDeviceMessages().size() != request.getDeviceSignedPrekeys().size() - 1) {
+        // device_messages should exist and be one shorter than device_signed_prekeys, since it doesn't have the primary's key.
+        throw new WebApplicationException(Response.status(400).build());
+      }
+      try {
+        // Checks that all except master ID are in device messages
+        MessageValidation.validateCompleteDeviceList(
+            authenticatedAccount.getAccount(), request.getDeviceMessages(),
+            IncomingMessage::getDestinationDeviceId, true, Optional.of(Device.MASTER_ID));
+        MessageValidation.validateRegistrationIds(
+            authenticatedAccount.getAccount(), request.getDeviceMessages(),
+            IncomingMessage::getDestinationDeviceId, IncomingMessage::getDestinationRegistrationId);
+        // Checks that all including master ID are in signed prekeys
+        MessageValidation.validateCompleteDeviceList(
+            authenticatedAccount.getAccount(), request.getDeviceSignedPrekeys().entrySet(),
+            e -> e.getKey(), false, Optional.empty());
+      } catch (MismatchedDevicesException e) {
+        throw new WebApplicationException(Response.status(409)
+            .type(MediaType.APPLICATION_JSON_TYPE)
+            .entity(new MismatchedDevices(e.getMissingDevices(),
+                e.getExtraDevices()))
+            .build());
+      } catch (StaleDevicesException e) {
+        throw new WebApplicationException(Response.status(410)
+            .type(MediaType.APPLICATION_JSON)
+            .entity(new StaleDevices(e.getStaleDevices()))
+            .build());
+      }
+    } else if (request.getDeviceMessages() != null && !request.getDeviceMessages().isEmpty()) {
+      // device_messages shouldn't exist without device_signed_prekeys.
+      throw new WebApplicationException(Response.status(400).build());
+    }
 
-      rateLimiters.getVerifyLimiter().validate(request.getNumber());
+    final String number = request.getNumber();
+    if (!authenticatedAccount.getAccount().getNumber().equals(number)) {
+      Util.requireNormalizedNumber(number);
+
+      rateLimiters.getVerifyLimiter().validate(number);
 
       final Optional<StoredVerificationCode> storedVerificationCode =
-          pendingAccounts.getCodeForNumber(request.getNumber());
+          pendingAccounts.getCodeForNumber(number);
 
       if (storedVerificationCode.isEmpty() || !storedVerificationCode.get().isValid(request.getCode())) {
-        throw new WebApplicationException(Response.status(403).build());
+        throw new ForbiddenException();
       }
 
       storedVerificationCode.flatMap(StoredVerificationCode::getTwilioVerificationSid)
           .ifPresent(smsSender::reportVerificationSucceeded);
 
-      final Optional<Account> existingAccount = accounts.getByE164(request.getNumber());
+      final Optional<Account> existingAccount = accounts.getByE164(number);
 
       if (existingAccount.isPresent()) {
         verifyRegistrationLock(existingAccount.get(), request.getRegistrationLock());
       }
 
-      rateLimiters.getVerifyLimiter().clear(request.getNumber());
-
-      updatedAccount = accounts.changeNumber(authenticatedAccount.getAccount(), request.getNumber());
+      rateLimiters.getVerifyLimiter().clear(number);
     }
 
+    final Account updatedAccount = changeNumberManager.changeNumber(
+        authenticatedAccount.getAccount(),
+        request.getNumber(),
+        Optional.ofNullable(request.getDeviceSignedPrekeys()).orElse(Collections.emptyMap()),
+        Optional.ofNullable(request.getDeviceMessages()).orElse(Collections.emptyList()));
+
     return new AccountIdentityResponse(
         updatedAccount.getUuid(),
         updatedAccount.getNumber(),

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

@@ -92,6 +92,7 @@ import org.whispersystems.textsecuregcm.storage.DeletedAccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
+import org.whispersystems.textsecuregcm.util.MessageValidation;
 import org.whispersystems.textsecuregcm.util.Pair;
 import org.whispersystems.textsecuregcm.util.Util;
 import org.whispersystems.textsecuregcm.util.ua.ClientPlatform;
@@ -225,10 +226,10 @@ public class MessageController {
         checkRateLimit(source.get(), destination.get(), userAgent);
       }
 
-      validateCompleteDeviceList(destination.get(), messages.getMessages(),
+      MessageValidation.validateCompleteDeviceList(destination.get(), messages.getMessages(),
           IncomingMessage::getDestinationDeviceId, isSyncMessage,
           source.map(AuthenticatedAccount::getAuthenticatedDevice).map(Device::getId));
-      validateRegistrationIds(destination.get(), messages.getMessages(),
+      MessageValidation.validateRegistrationIds(destination.get(), messages.getMessages(),
           IncomingMessage::getDestinationDeviceId, IncomingMessage::getDestinationRegistrationId);
 
       final List<Tag> tags = List.of(UserAgentTagUtil.getPlatformTag(userAgent),
@@ -319,10 +320,10 @@ public class MessageController {
       }
 
       final List<IncomingDeviceMessage> messagesAsList = Arrays.asList(messages);
-      validateCompleteDeviceList(destination.get(), messagesAsList,
+      MessageValidation.validateCompleteDeviceList(destination.get(), messagesAsList,
           IncomingDeviceMessage::getDeviceId, isSyncMessage,
           source.map(AuthenticatedAccount::getAuthenticatedDevice).map(Device::getId));
-      validateRegistrationIds(destination.get(), messagesAsList,
+      MessageValidation.validateRegistrationIds(destination.get(), messagesAsList,
           IncomingDeviceMessage::getDeviceId,
           IncomingDeviceMessage::getRegistrationId);
 
@@ -402,8 +403,8 @@ public class MessageController {
       final Set<Pair<Long, Integer>> deviceIdAndRegistrationIdSet = accountToDeviceIdAndRegistrationIdMap.get(account);
       final Set<Long> deviceIds = deviceIdAndRegistrationIdSet.stream().map(Pair::first).collect(Collectors.toSet());
       try {
-        validateCompleteDeviceList(account, deviceIds, false, Optional.empty());
-        validateRegistrationIds(account, deviceIdAndRegistrationIdSet.stream());
+        MessageValidation.validateCompleteDeviceList(account, deviceIds, false, Optional.empty());
+        MessageValidation.validateRegistrationIds(account, deviceIdAndRegistrationIdSet.stream());
       } catch (MismatchedDevicesException e) {
         accountMismatchedDevices.add(new AccountMismatchedDevices(account.getUuid(),
             new MismatchedDevices(e.getMissingDevices(), e.getExtraDevices())));
@@ -731,72 +732,6 @@ public class MessageController {
     }
   }
 
-  @VisibleForTesting
-  public static <T> void validateRegistrationIds(Account account, List<T> messages, Function<T, Long> getDeviceId, Function<T, Integer> getRegistrationId)
-      throws StaleDevicesException {
-    final Stream<Pair<Long, Integer>> deviceIdAndRegistrationIdStream = messages
-        .stream()
-        .map(message -> new Pair<>(getDeviceId.apply(message), getRegistrationId.apply(message)));
-    validateRegistrationIds(account, deviceIdAndRegistrationIdStream);
-  }
-
-  @VisibleForTesting
-  public static void validateRegistrationIds(Account account, Stream<Pair<Long, Integer>> deviceIdAndRegistrationIdStream)
-      throws StaleDevicesException {
-    final List<Long> staleDevices = deviceIdAndRegistrationIdStream
-        .filter(deviceIdAndRegistrationId -> deviceIdAndRegistrationId.second() > 0)
-        .filter(deviceIdAndRegistrationId -> {
-          Optional<Device> device = account.getDevice(deviceIdAndRegistrationId.first());
-          return device.isPresent() && deviceIdAndRegistrationId.second() != device.get().getRegistrationId();
-        })
-        .map(Pair::first)
-        .collect(Collectors.toList());
-
-    if (!staleDevices.isEmpty()) {
-      throw new StaleDevicesException(staleDevices);
-    }
-  }
-
-  @VisibleForTesting
-  public static <T> void validateCompleteDeviceList(Account account, List<T> messages, Function<T, Long> getDeviceId, boolean isSyncMessage,
-      Optional<Long> authenticatedDeviceId)
-      throws MismatchedDevicesException {
-    Set<Long> messageDeviceIds = messages.stream().map(getDeviceId)
-        .collect(Collectors.toSet());
-    validateCompleteDeviceList(account, messageDeviceIds, isSyncMessage, authenticatedDeviceId);
-  }
-
-  @VisibleForTesting
-  public static void validateCompleteDeviceList(Account account, Set<Long> messageDeviceIds, boolean isSyncMessage,
-      Optional<Long> authenticatedDeviceId)
-      throws MismatchedDevicesException {
-    Set<Long> accountDeviceIds = new HashSet<>();
-
-    List<Long> missingDeviceIds = new LinkedList<>();
-    List<Long> extraDeviceIds = new LinkedList<>();
-
-    for (Device device : account.getDevices()) {
-      if (device.isEnabled() &&
-          !(isSyncMessage && device.getId() == authenticatedDeviceId.get())) {
-        accountDeviceIds.add(device.getId());
-
-        if (!messageDeviceIds.contains(device.getId())) {
-          missingDeviceIds.add(device.getId());
-        }
-      }
-    }
-
-    for (Long deviceId : messageDeviceIds) {
-      if (!accountDeviceIds.contains(deviceId)) {
-        extraDeviceIds.add(deviceId);
-      }
-    }
-
-    if (!missingDeviceIds.isEmpty() || !extraDeviceIds.isEmpty()) {
-      throw new MismatchedDevicesException(missingDeviceIds, extraDeviceIds);
-    }
-  }
-
   private void validateContentLength(final int contentLength, final String userAgent) {
     Metrics.summary(CONTENT_SIZE_DISTRIBUTION_NAME, Tags.of(UserAgentTagUtil.getPlatformTag(userAgent)))
         .record(contentLength);
@@ -818,7 +753,7 @@ public class MessageController {
     }
   }
 
-  private Optional<byte[]> getMessageContent(IncomingMessage message) {
+  public static Optional<byte[]> getMessageContent(IncomingMessage message) {
     if (Util.isEmpty(message.getContent())) return Optional.empty();
 
     try {

service/src/main/java/org/whispersystems/textsecuregcm/entities/ChangePhoneNumberRequest.java

@@ -9,6 +9,8 @@ import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import javax.annotation.Nullable;
 import javax.validation.constraints.NotBlank;
+import java.util.List;
+import java.util.Map;
 
 public class ChangePhoneNumberRequest {
 
@@ -24,14 +26,26 @@ public class ChangePhoneNumberRequest {
   @Nullable
   final String registrationLock;
 
+  @JsonProperty("device_messages")
+  @Nullable
+  final List<IncomingMessage> deviceMessages;
+
+  @JsonProperty("device_signed_prekeys")
+  @Nullable
+  final Map<Long, SignedPreKey> deviceSignedPrekeys;
+
   @JsonCreator
   public ChangePhoneNumberRequest(@JsonProperty("number") final String number,
       @JsonProperty("code") final String code,
-      @JsonProperty("reglock") @Nullable final String registrationLock) {
+      @JsonProperty("reglock") @Nullable final String registrationLock,
+      @JsonProperty("device_messages") @Nullable final List<IncomingMessage> deviceMessages,
+      @JsonProperty("device_signed_prekeys") @Nullable final Map<Long, SignedPreKey> deviceSignedPrekeys) {
 
     this.number = number;
     this.code = code;
     this.registrationLock = registrationLock;
+    this.deviceMessages = deviceMessages;
+    this.deviceSignedPrekeys = deviceSignedPrekeys;
   }
 
   public String getNumber() {
@@ -46,4 +60,14 @@ public class ChangePhoneNumberRequest {
   public String getRegistrationLock() {
     return registrationLock;
   }
+
+  @Nullable
+  public List<IncomingMessage> getDeviceMessages() {
+    return deviceMessages;
+  }
+
+  @Nullable
+  public Map<Long, SignedPreKey> getDeviceSignedPrekeys() {
+    return deviceSignedPrekeys;
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/storage/ChangeNumberManager.java

@@ -0,0 +1,95 @@
+/*
+ * Copyright 2013-2022 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.textsecuregcm.storage;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.protobuf.ByteString;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.controllers.AccountController;
+import org.whispersystems.textsecuregcm.controllers.MessageController;
+import org.whispersystems.textsecuregcm.entities.IncomingMessage;
+import org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
+import org.whispersystems.textsecuregcm.entities.SignedPreKey;
+import org.whispersystems.textsecuregcm.push.MessageSender;
+import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;
+import javax.validation.constraints.NotNull;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+public class ChangeNumberManager {
+  private static final Logger logger = LoggerFactory.getLogger(AccountController.class);
+  private final MessageSender messageSender;
+  private final AccountsManager accountsManager;
+
+  public ChangeNumberManager(
+      final MessageSender messageSender,
+      final AccountsManager accountsManager) {
+    this.messageSender = messageSender;
+    this.accountsManager = accountsManager;
+  }
+
+  public Account changeNumber(
+      @NotNull Account account,
+      @NotNull final String number,
+      @NotNull final Map<Long, SignedPreKey> deviceSignedPrekeys,
+      @NotNull final List<IncomingMessage> deviceMessages) throws InterruptedException {
+
+    final Account updatedAccount;
+    if (number.equals(account.getNumber())) {
+      // This may be a request that got repeated due to poor network conditions or other client error; take no action,
+      // but report success since the account is in the desired state
+      updatedAccount = account;
+    } else {
+      updatedAccount = accountsManager.changeNumber(account, number);
+    }
+
+    // Whether the account already has this number or not, we reset signed prekeys and resend messages.
+    // This makes it so the client can resend a request they didn't get a response for (timeout, etc)
+    // to make sure their messages sent and prekeys were updated, even if the first time around the
+    // server crashed at/above this point.
+    if (deviceSignedPrekeys != null && !deviceSignedPrekeys.isEmpty()) {
+      for (Map.Entry<Long, SignedPreKey> entry : deviceSignedPrekeys.entrySet()) {
+        accountsManager.updateDevice(updatedAccount, entry.getKey(),
+            d -> d.setPhoneNumberIdentitySignedPreKey(entry.getValue()));
+      }
+
+      for (IncomingMessage message : deviceMessages) {
+        sendMessageToSelf(updatedAccount, updatedAccount.getDevice(message.getDestinationDeviceId()), message);
+      }
+    }
+    return updatedAccount;
+  }
+
+  @VisibleForTesting
+  void sendMessageToSelf(
+      Account sourceAndDestinationAccount, Optional<Device> destinationDevice, IncomingMessage message) {
+    Optional<byte[]> contents = MessageController.getMessageContent(message);
+    if (!contents.isPresent()) {
+      logger.debug("empty message contents sending to self, ignoring");
+      return;
+    } else if (!destinationDevice.isPresent()) {
+      logger.debug("destination device not present");
+      return;
+    }
+    try {
+      long serverTimestamp = System.currentTimeMillis();
+      Envelope envelope = Envelope.newBuilder()
+          .setType(Envelope.Type.forNumber(message.getType()))
+          .setTimestamp(serverTimestamp)
+          .setServerTimestamp(serverTimestamp)
+          .setDestinationUuid(sourceAndDestinationAccount.getUuid().toString())
+          .setContent(ByteString.copyFrom(contents.get()))
+          .setSource(sourceAndDestinationAccount.getNumber())
+          .setSourceUuid(sourceAndDestinationAccount.getUuid().toString())
+          .setSourceDevice((int) Device.MASTER_ID)
+          .build();
+      messageSender.sendMessage(sourceAndDestinationAccount, destinationDevice.get(), envelope, false);
+    } catch (NotPushRegisteredException e) {
+      logger.debug("Not registered", e);
+    }
+  }
+}

service/src/main/java/org/whispersystems/textsecuregcm/util/MessageValidation.java

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2013-2022 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.textsecuregcm.util;
+
+import org.whispersystems.textsecuregcm.controllers.MismatchedDevicesException;
+import org.whispersystems.textsecuregcm.controllers.StaleDevicesException;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.Device;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+public class MessageValidation {
+  public static <T> void validateRegistrationIds(Account account, List<T> messages, Function<T, Long> getDeviceId, Function<T, Integer> getRegistrationId)
+      throws StaleDevicesException {
+    final Stream<Pair<Long, Integer>> deviceIdAndRegistrationIdStream = messages
+        .stream()
+        .map(message -> new Pair<>(getDeviceId.apply(message), getRegistrationId.apply(message)));
+    validateRegistrationIds(account, deviceIdAndRegistrationIdStream);
+  }
+
+  public static void validateRegistrationIds(Account account, Stream<Pair<Long, Integer>> deviceIdAndRegistrationIdStream)
+      throws StaleDevicesException {
+    final List<Long> staleDevices = deviceIdAndRegistrationIdStream
+        .filter(deviceIdAndRegistrationId -> deviceIdAndRegistrationId.second() > 0)
+        .filter(deviceIdAndRegistrationId -> {
+          Optional<Device> device = account.getDevice(deviceIdAndRegistrationId.first());
+          return device.isPresent() && deviceIdAndRegistrationId.second() != device.get().getRegistrationId();
+        })
+        .map(Pair::first)
+        .collect(Collectors.toList());
+
+    if (!staleDevices.isEmpty()) {
+      throw new StaleDevicesException(staleDevices);
+    }
+  }
+
+  public static <T> void validateCompleteDeviceList(Account account, Collection<T> messages, Function<T, Long> getDeviceId, boolean isSyncMessage,
+      Optional<Long> authenticatedDeviceId)
+      throws MismatchedDevicesException {
+    Set<Long> messageDeviceIds = messages.stream().map(getDeviceId)
+        .collect(Collectors.toSet());
+    validateCompleteDeviceList(account, messageDeviceIds, isSyncMessage, authenticatedDeviceId);
+  }
+
+  public static void validateCompleteDeviceList(Account account, Set<Long> messageDeviceIds, boolean isSyncMessage,
+      Optional<Long> authenticatedDeviceId)
+      throws MismatchedDevicesException {
+    Set<Long> accountDeviceIds = new HashSet<>();
+
+    List<Long> missingDeviceIds = new LinkedList<>();
+    List<Long> extraDeviceIds = new LinkedList<>();
+
+    for (Device device : account.getDevices()) {
+      if (device.isEnabled() &&
+          !(isSyncMessage && device.getId() == authenticatedDeviceId.get())) {
+        accountDeviceIds.add(device.getId());
+
+        if (!messageDeviceIds.contains(device.getId())) {
+          missingDeviceIds.add(device.getId());
+        }
+      }
+    }
+
+    for (Long deviceId : messageDeviceIds) {
+      if (!accountDeviceIds.contains(deviceId)) {
+        extraDeviceIds.add(deviceId);
+      }
+    }
+
+    if (!missingDeviceIds.isEmpty() || !extraDeviceIds.isEmpty()) {
+      throw new MismatchedDevicesException(missingDeviceIds, extraDeviceIds);
+    }
+  }
+
+}