Refresh accounts from storage when checking for device state changes after requests

service/src/main/java/org/whispersystems/textsecuregcm/auth/AuthEnablementRefreshRequirementProvider.java

@@ -14,9 +14,11 @@ import java.util.Set;
 import java.util.UUID;
 import java.util.stream.Collectors;
 import org.glassfish.jersey.server.ContainerRequest;
+import org.glassfish.jersey.server.monitoring.RequestEvent;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.util.Pair;
 
@@ -33,34 +35,48 @@ import org.whispersystems.textsecuregcm.util.Pair;
  */
 public class AuthEnablementRefreshRequirementProvider implements WebsocketRefreshRequirementProvider {
 
+  private final AccountsManager accountsManager;
+
   private static final Logger logger = LoggerFactory.getLogger(AuthEnablementRefreshRequirementProvider.class);
 
+  private static final String ACCOUNT_UUID = AuthEnablementRefreshRequirementProvider.class.getName() + ".accountUuid";
   private static final String DEVICES_ENABLED = AuthEnablementRefreshRequirementProvider.class.getName() + ".devicesEnabled";
 
+  public AuthEnablementRefreshRequirementProvider(final AccountsManager accountsManager) {
+    this.accountsManager = accountsManager;
+  }
+
   @VisibleForTesting
-  Map<Long, Boolean> buildDevicesEnabledMap(final Account account) {
+  static Map<Long, Boolean> buildDevicesEnabledMap(final Account account) {
     return account.getDevices().stream().collect(Collectors.toMap(Device::getId, Device::isEnabled));
   }
 
   @Override
-  public void handleRequestFiltered(final ContainerRequest request) {
-    // The authenticated principal, if any, will be available after filters have run.
-    // Now that the account is known, capture a snapshot of `isEnabled` for the account's devices before carrying out
-    // the request’s business logic.
-    ContainerRequestUtil.getAuthenticatedAccount(request)
-        .ifPresent(account -> request.setProperty(DEVICES_ENABLED, buildDevicesEnabledMap(account)));
+  public void handleRequestFiltered(final RequestEvent requestEvent) {
+    if (requestEvent.getUriInfo().getMatchedResourceMethod().getInvocable().getHandlingMethod().getAnnotation(ChangesDeviceEnabledState.class) != null) {
+      // The authenticated principal, if any, will be available after filters have run.
+      // Now that the account is known, capture a snapshot of `isEnabled` for the account's devices before carrying out
+      // the request’s business logic.
+      ContainerRequestUtil.getAuthenticatedAccount(requestEvent.getContainerRequest()).ifPresent(account ->
+          setAccount(requestEvent.getContainerRequest(), account));
+    }
+  }
+
+  public static void setAccount(final ContainerRequest containerRequest, final Account account) {
+    containerRequest.setProperty(ACCOUNT_UUID, account.getUuid());
+    containerRequest.setProperty(DEVICES_ENABLED, buildDevicesEnabledMap(account));
   }
 
   @Override
-  public List<Pair<UUID, Long>> handleRequestFinished(final ContainerRequest request) {
+  public List<Pair<UUID, Long>> handleRequestFinished(final RequestEvent requestEvent) {
     // Now that the request is finished, check whether `isEnabled` changed for any of the devices. If the value did
     // change or if a devices was added or removed, all devices must disconnect and reauthenticate.
-    if (request.getProperty(DEVICES_ENABLED) != null) {
+    if (requestEvent.getContainerRequest().getProperty(DEVICES_ENABLED) != null) {
 
       @SuppressWarnings("unchecked") final Map<Long, Boolean> initialDevicesEnabled =
-          (Map<Long, Boolean>) request.getProperty(DEVICES_ENABLED);
+          (Map<Long, Boolean>) requestEvent.getContainerRequest().getProperty(DEVICES_ENABLED);
 
-      return ContainerRequestUtil.getAuthenticatedAccount(request).map(account -> {
+      return accountsManager.get((UUID) requestEvent.getContainerRequest().getProperty(ACCOUNT_UUID)).map(account -> {
         final Set<Long> deviceIdsToDisplace;
         final Map<Long, Boolean> currentDevicesEnabled = buildDevicesEnabledMap(account);
 

service/src/main/java/org/whispersystems/textsecuregcm/auth/ChangesDeviceEnabledState.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2013-2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.auth;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Indicates that an endpoint may change the "enabled" state of one or more devices associated with an account, and that
+ * any websockets associated with the account may need to be refreshed after a call to that endpoint.
+ */
+@Target(ElementType.METHOD)
+@Retention(RetentionPolicy.RUNTIME)
+public @interface ChangesDeviceEnabledState {
+}

service/src/main/java/org/whispersystems/textsecuregcm/auth/PhoneNumberChangeRefreshRequirementProvider.java

@@ -11,6 +11,7 @@ import java.util.Optional;
 import java.util.UUID;
 import java.util.stream.Collectors;
 import org.glassfish.jersey.server.ContainerRequest;
+import org.glassfish.jersey.server.monitoring.RequestEvent;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.util.Pair;
 
@@ -20,17 +21,18 @@ public class PhoneNumberChangeRefreshRequirementProvider implements WebsocketRef
       PhoneNumberChangeRefreshRequirementProvider.class.getName() + ".initialNumber";
 
   @Override
-  public void handleRequestFiltered(final ContainerRequest request) {
-    ContainerRequestUtil.getAuthenticatedAccount(request)
-        .ifPresent(account -> request.setProperty(INITIAL_NUMBER_KEY, account.getNumber()));
+  public void handleRequestFiltered(final RequestEvent requestEvent) {
+    ContainerRequestUtil.getAuthenticatedAccount(requestEvent.getContainerRequest())
+        .ifPresent(account -> requestEvent.getContainerRequest().setProperty(INITIAL_NUMBER_KEY, account.getNumber()));
   }
 
   @Override
-  public List<Pair<UUID, Long>> handleRequestFinished(final ContainerRequest request) {
-    final String initialNumber = (String) request.getProperty(INITIAL_NUMBER_KEY);
+  public List<Pair<UUID, Long>> handleRequestFinished(final RequestEvent requestEvent) {
+    final String initialNumber = (String) requestEvent.getContainerRequest().getProperty(INITIAL_NUMBER_KEY);
 
     if (initialNumber != null) {
-      final Optional<Account> maybeAuthenticatedAccount = ContainerRequestUtil.getAuthenticatedAccount(request);
+      final Optional<Account> maybeAuthenticatedAccount =
+          ContainerRequestUtil.getAuthenticatedAccount(requestEvent.getContainerRequest());
 
       return maybeAuthenticatedAccount
           .filter(account -> !initialNumber.equals(account.getNumber()))

service/src/main/java/org/whispersystems/textsecuregcm/auth/WebsocketRefreshApplicationEventListener.java

@@ -10,6 +10,7 @@ import org.glassfish.jersey.server.monitoring.ApplicationEventListener;
 import org.glassfish.jersey.server.monitoring.RequestEvent;
 import org.glassfish.jersey.server.monitoring.RequestEventListener;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
+import org.whispersystems.textsecuregcm.storage.AccountsManager;
 
 /**
  * Delegates request events to a listener that watches for intra-request changes that require websocket refreshes
@@ -18,9 +19,11 @@ public class WebsocketRefreshApplicationEventListener implements ApplicationEven
 
   private final WebsocketRefreshRequestEventListener websocketRefreshRequestEventListener;
 
-  public WebsocketRefreshApplicationEventListener(final ClientPresenceManager clientPresenceManager) {
+  public WebsocketRefreshApplicationEventListener(final AccountsManager accountsManager,
+      final ClientPresenceManager clientPresenceManager) {
+
     this.websocketRefreshRequestEventListener = new WebsocketRefreshRequestEventListener(clientPresenceManager,
-        new AuthEnablementRefreshRequirementProvider(),
+        new AuthEnablementRefreshRequirementProvider(accountsManager),
         new PhoneNumberChangeRefreshRequirementProvider());
   }
 

service/src/main/java/org/whispersystems/textsecuregcm/auth/WebsocketRefreshRequestEventListener.java

@@ -16,6 +16,9 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 
+import javax.ws.rs.container.ResourceInfo;
+import javax.ws.rs.core.Context;
+
 import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
 
 public class WebsocketRefreshRequestEventListener implements RequestEventListener {
@@ -39,17 +42,20 @@ public class WebsocketRefreshRequestEventListener implements RequestEventListene
     this.providers = providers;
   }
 
+  @Context
+  private ResourceInfo resourceInfo;
+
   @Override
   public void onEvent(final RequestEvent event) {
     if (event.getType() == Type.REQUEST_FILTERED) {
       for (final WebsocketRefreshRequirementProvider provider : providers) {
-        provider.handleRequestFiltered(event.getContainerRequest());
+        provider.handleRequestFiltered(event);
       }
     } else if (event.getType() == Type.FINISHED) {
       final AtomicInteger displacedDevices = new AtomicInteger(0);
 
       Arrays.stream(providers)
-          .flatMap(provider -> provider.handleRequestFinished(event.getContainerRequest()).stream())
+          .flatMap(provider -> provider.handleRequestFinished(event).stream())
           .distinct()
           .forEach(pair -> {
             try {

service/src/main/java/org/whispersystems/textsecuregcm/auth/WebsocketRefreshRequirementProvider.java

@@ -8,6 +8,7 @@ package org.whispersystems.textsecuregcm.auth;
 import java.util.List;
 import java.util.UUID;
 import org.glassfish.jersey.server.ContainerRequest;
+import org.glassfish.jersey.server.monitoring.RequestEvent;
 import org.whispersystems.textsecuregcm.util.Pair;
 
 /**
@@ -19,16 +20,16 @@ public interface WebsocketRefreshRequirementProvider {
   /**
    * Processes a request after filters have run and the request has been mapped to a destination controller.
    *
-   * @param request the request to observe
+   * @param requestEvent the request event to observe
    */
-  void handleRequestFiltered(ContainerRequest request);
+  void handleRequestFiltered(RequestEvent requestEvent);
 
   /**
    * Processes a request after all normal request handling has been completed.
    *
-   * @param request the request to observe
+   * @param requestEvent the request event to observe
    * @return a list of pairs of account UUID/device ID pairs identifying websockets that need to be refreshed as a
    * result of the observed request
    */
-  List<Pair<UUID, Long>> handleRequestFinished(ContainerRequest request);
+  List<Pair<UUID, Long>> handleRequestFinished(RequestEvent requestEvent);
 }