From cc7b030a41da4b5c0fa55c7d5496f47bc2366b76 Mon Sep 17 00:00:00 2001 From: Jonathan Klabunde Tomer <125505367+jkt-signal@users.noreply.github.com> Date: Tue, 6 May 2025 13:36:41 -0700 Subject: [PATCH] Send disconnection requests after non-API device unlinks --- .../textsecuregcm/storage/AccountsManager.java | 9 +++++++++ .../textsecuregcm/storage/MessagePersister.java | 11 ++++++++--- .../textsecuregcm/workers/CommandDependencies.java | 2 ++ .../workers/MessagePersisterServiceCommand.java | 1 + .../workers/RemoveExpiredLinkedDevicesCommand.java | 5 ++++- .../storage/MessagePersisterIntegrationTest.java | 4 +++- .../textsecuregcm/storage/MessagePersisterTest.java | 10 ++++++++-- .../FinishPushNotificationExperimentCommandTest.java | 1 + ...LockAccountsWithoutPniIdentityKeysCommandTest.java | 1 + .../workers/LockAccountsWithoutPqKeysCommandTest.java | 1 + .../workers/NotifyIdleDevicesCommandTest.java | 1 + ...moveAccountsWithoutPniIdentityKeysCommandTest.java | 1 + .../RemoveAccountsWithoutPqKeysCommandTest.java | 1 + .../RemoveLinkedDevicesWithoutPniKeysCommandTest.java | 1 + .../RemoveLinkedDevicesWithoutPqKeysCommandTest.java | 1 + .../StartPushNotificationExperimentCommandTest.java | 1 + 16 files changed, 44 insertions(+), 7 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java index 5c2207076..520bbf8b1 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java @@ -580,6 +580,15 @@ public class AccountsManager extends RedisPubSubAdapter implemen return Optional.of(aci); } + /** + * Unlink a device from the given account. The device will be immediately disconnected if it is + * connected to any chat frontend, but it is the caller's responsibility to make sure that the + * account's *other* devices are disconnected, either by use of + * {@link org.whispersystems.textsecuregcm.auth.LinkedDeviceRefreshRequirementProvider} or + * directly by calling {@link DeviceDisconnectionManager#requestDisconnection}. + * + * @returns the updated Account + */ public CompletableFuture removeDevice(final Account account, final byte deviceId) { if (deviceId == Device.PRIMARY_ID) { throw new IllegalArgumentException("Cannot remove primary device"); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagePersister.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagePersister.java index e815ca9f9..0ac62e9f3 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagePersister.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagePersister.java @@ -25,6 +25,7 @@ import java.util.concurrent.atomic.AtomicLong; import org.apache.commons.lang3.tuple.Pair; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.whispersystems.textsecuregcm.auth.DisconnectionRequestManager; import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration; import org.whispersystems.textsecuregcm.entities.MessageProtos; import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager; @@ -44,6 +45,7 @@ public class MessagePersister implements Managed { private final AccountsManager accountsManager; private final DynamicConfigurationManager dynamicConfigurationManager; private final ExperimentEnrollmentManager experimentEnrollmentManager; + private final DisconnectionRequestManager disconnectionRequestManager; private final Duration persistDelay; @@ -82,6 +84,7 @@ public class MessagePersister implements Managed { final AccountsManager accountsManager, final DynamicConfigurationManager dynamicConfigurationManager, final ExperimentEnrollmentManager experimentEnrollmentManager, + final DisconnectionRequestManager disconnectionRequestManager, final Duration persistDelay, final int dedicatedProcessWorkerThreadCount) { @@ -90,6 +93,7 @@ public class MessagePersister implements Managed { this.accountsManager = accountsManager; this.dynamicConfigurationManager = dynamicConfigurationManager; this.experimentEnrollmentManager = experimentEnrollmentManager; + this.disconnectionRequestManager = disconnectionRequestManager; this.persistDelay = persistDelay; this.workerThreads = new Thread[dedicatedProcessWorkerThreadCount]; @@ -257,9 +261,10 @@ public class MessagePersister implements Managed { trimQueue(account, deviceId); throw new MessagePersistenceException("Could not persist due to an overfull queue. Trimmed primary queue, a subsequent retry may succeed"); } else { - logger.warn("Failed to persist queue {}::{} due to overfull queue; will unlink device", account.getUuid(), - deviceId); - accountsManager.removeDevice(account, deviceId).join(); + logger.warn("Failed to persist queue {}::{} due to overfull queue; will unlink device", accountUuid, deviceId); + accountsManager.removeDevice(account, deviceId) + .thenRun(() -> disconnectionRequestManager.requestDisconnection(accountUuid)) + .join(); } } finally { messagesCache.unlockQueueForPersistence(accountUuid, deviceId); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java index e1b24b5f7..43cdb492b 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java @@ -77,6 +77,7 @@ record CommandDependencies( AccountsManager accountsManager, ProfilesManager profilesManager, ReportMessageManager reportMessageManager, + DisconnectionRequestManager disconnectionRequestManager, MessagesCache messagesCache, MessagesManager messagesManager, KeysManager keysManager, @@ -289,6 +290,7 @@ record CommandDependencies( accountsManager, profilesManager, reportMessageManager, + disconnectionRequestManager, messagesCache, messagesManager, keys, diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/MessagePersisterServiceCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/MessagePersisterServiceCommand.java index 7a4663567..cf9221dab 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/MessagePersisterServiceCommand.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/MessagePersisterServiceCommand.java @@ -66,6 +66,7 @@ public class MessagePersisterServiceCommand extends ServerCommand accountUpdate = dryRun ? Mono.just((long) expiredDevices.size()) - : deleteDevices(account, expiredDevices, maxRetries); + : deleteDevices(account, expiredDevices, maxRetries) + .flatMap(count -> + Mono.fromCompletionStage(getCommandDependencies().disconnectionRequestManager().requestDisconnection(account.getUuid())) + .then(Mono.just(count))); return accountUpdate .doOnNext(successCounter::increment) diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterIntegrationTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterIntegrationTest.java index 2300f90fd..459a82459 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterIntegrationTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterIntegrationTest.java @@ -30,6 +30,7 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; +import org.whispersystems.textsecuregcm.auth.DisconnectionRequestManager; import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration; import org.whispersystems.textsecuregcm.entities.MessageProtos; import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager; @@ -98,7 +99,8 @@ class MessagePersisterIntegrationTest { webSocketConnectionEventManager.start(); messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager, - dynamicConfigurationManager, mock(ExperimentEnrollmentManager.class), PERSIST_DELAY, 1); + dynamicConfigurationManager, mock(ExperimentEnrollmentManager.class), mock(DisconnectionRequestManager.class), + PERSIST_DELAY, 1); account = mock(Account.class); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterTest.java index 87bf184a2..a62f32250 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/MessagePersisterTest.java @@ -10,7 +10,6 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyByte; -import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.ArgumentMatchers.eq; @@ -21,6 +20,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; import static org.whispersystems.textsecuregcm.util.MockUtils.exactly; @@ -51,6 +51,7 @@ import org.junit.jupiter.api.Timeout; import org.junit.jupiter.api.extension.RegisterExtension; import org.mockito.ArgumentCaptor; import org.mockito.stubbing.Answer; +import org.whispersystems.textsecuregcm.auth.DisconnectionRequestManager; import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration; import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicMessagePersisterConfiguration; import org.whispersystems.textsecuregcm.entities.MessageProtos; @@ -77,6 +78,7 @@ class MessagePersisterTest { private MessagePersister messagePersister; private AccountsManager accountsManager; private MessagesManager messagesManager; + private DisconnectionRequestManager disconnectionRequestManager; private Account destinationAccount; private static final UUID DESTINATION_ACCOUNT_UUID = UUID.randomUUID(); @@ -97,6 +99,7 @@ class MessagePersisterTest { messagesDynamoDb = mock(MessagesDynamoDb.class); accountsManager = mock(AccountsManager.class); + disconnectionRequestManager = mock(DisconnectionRequestManager.class); destinationAccount = mock(Account.class); when(accountsManager.getByAccountIdentifier(DESTINATION_ACCOUNT_UUID)).thenReturn(Optional.of(destinationAccount)); @@ -119,7 +122,8 @@ class MessagePersisterTest { messagesCache = new MessagesCache(REDIS_CLUSTER_EXTENSION.getRedisCluster(), messageDeliveryScheduler, sharedExecutorService, Clock.systemUTC()); messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager, - dynamicConfigurationManager, mock(ExperimentEnrollmentManager.class), PERSIST_DELAY, 1); + dynamicConfigurationManager, mock(ExperimentEnrollmentManager.class), disconnectionRequestManager, + PERSIST_DELAY, 1); when(messagesManager.clear(any(UUID.class), anyByte())).thenReturn(CompletableFuture.completedFuture(null)); @@ -301,6 +305,7 @@ class MessagePersisterTest { assertTimeoutPreemptively(Duration.ofSeconds(1), () -> messagePersister.persistQueue(destinationAccount, DESTINATION_DEVICE, "test")); verify(accountsManager, exactly()).removeDevice(destinationAccount, DESTINATION_DEVICE_ID); + verify(disconnectionRequestManager, exactly()).requestDisconnection(DESTINATION_ACCOUNT_UUID); } @Test @@ -402,6 +407,7 @@ class MessagePersisterTest { when(accountsManager.removeDevice(destinationAccount, DESTINATION_DEVICE_ID)).thenReturn(CompletableFuture.failedFuture(new TimeoutException())); assertThrows(CompletionException.class, () -> messagePersister.persistQueue(destinationAccount, DESTINATION_DEVICE, "test")); + verifyNoMoreInteractions(disconnectionRequestManager); } @SuppressWarnings("SameParameterValue") diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/FinishPushNotificationExperimentCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/FinishPushNotificationExperimentCommandTest.java index 3daf76f31..7605b2b76 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/workers/FinishPushNotificationExperimentCommandTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/FinishPushNotificationExperimentCommandTest.java @@ -74,6 +74,7 @@ class FinishPushNotificationExperimentCommandTest { null, null, null, + null, pushNotificationExperimentSamples, null, null, diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPniIdentityKeysCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPniIdentityKeysCommandTest.java index 0880d9b27..5c84f33a0 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPniIdentityKeysCommandTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPniIdentityKeysCommandTest.java @@ -59,6 +59,7 @@ class LockAccountsWithoutPniIdentityKeysCommandTest { null, null, null, + null, null); namespace = new Namespace(Map.of( diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPqKeysCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPqKeysCommandTest.java index ddd645ab9..99cbf700f 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPqKeysCommandTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPqKeysCommandTest.java @@ -50,6 +50,7 @@ class LockAccountsWithoutPqKeysCommandTest { null, null, null, + null, keysManager, null, null, diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/NotifyIdleDevicesCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/NotifyIdleDevicesCommandTest.java index dd3efbf03..70044e5bc 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/workers/NotifyIdleDevicesCommandTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/NotifyIdleDevicesCommandTest.java @@ -50,6 +50,7 @@ class NotifyIdleDevicesCommandTest { null, null, null, + null, messagesManager, null, null, diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPniIdentityKeysCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPniIdentityKeysCommandTest.java index 76d45ffaa..8505e9839 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPniIdentityKeysCommandTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPniIdentityKeysCommandTest.java @@ -54,6 +54,7 @@ class RemoveAccountsWithoutPniIdentityKeysCommandTest { null, null, null, + null, null); namespace = new Namespace(Map.of( diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPqKeysCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPqKeysCommandTest.java index c89a5bc49..62d21e058 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPqKeysCommandTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPqKeysCommandTest.java @@ -50,6 +50,7 @@ class RemoveAccountsWithoutPqKeysCommandTest { null, null, null, + null, keysManager, null, null, diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPniKeysCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPniKeysCommandTest.java index 7e3e8bc6d..48ab797f2 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPniKeysCommandTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPniKeysCommandTest.java @@ -48,6 +48,7 @@ class RemoveLinkedDevicesWithoutPniKeysCommandTest { null, null, null, + null, keysManager, null, null, diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPqKeysCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPqKeysCommandTest.java index 03e617764..913dfa80d 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPqKeysCommandTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPqKeysCommandTest.java @@ -47,6 +47,7 @@ class RemoveLinkedDevicesWithoutPqKeysCommandTest { null, null, null, + null, keysManager, null, null, diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/StartPushNotificationExperimentCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/StartPushNotificationExperimentCommandTest.java index 8b71ec2db..0a87de49b 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/workers/StartPushNotificationExperimentCommandTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/StartPushNotificationExperimentCommandTest.java @@ -63,6 +63,7 @@ class StartPushNotificationExperimentCommandTest { null, null, null, + null, pushNotificationExperimentSamples, null, null,