From d316c72beb95e46e2700475bf9a90a13e06be632 Mon Sep 17 00:00:00 2001 From: Jon Chambers <63609320+jon-signal@users.noreply.github.com> Date: Mon, 5 May 2025 12:10:47 -0400 Subject: [PATCH] Add commands for removing accounts/devices without PNI key material --- .../textsecuregcm/WhisperServerService.java | 7 + ...AccountsWithoutPniIdentityKeysCommand.java | 104 +++++++++++++ ...AccountsWithoutPniIdentityKeysCommand.java | 112 ++++++++++++++ ...oveLinkedDevicesWithoutPniKeysCommand.java | 122 +++++++++++++++ ...untsWithoutPniIdentityKeysCommandTest.java | 119 +++++++++++++++ ...untsWithoutPniIdentityKeysCommandTest.java | 112 ++++++++++++++ ...inkedDevicesWithoutPniKeysCommandTest.java | 142 ++++++++++++++++++ 7 files changed, 718 insertions(+) create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPniIdentityKeysCommand.java create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPniIdentityKeysCommand.java create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPniKeysCommand.java create mode 100644 service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPniIdentityKeysCommandTest.java create mode 100644 service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPniIdentityKeysCommandTest.java create mode 100644 service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPniKeysCommandTest.java diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index b666c5db5..f123052cb 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java @@ -266,15 +266,18 @@ import org.whispersystems.textsecuregcm.workers.CertificateCommand; import org.whispersystems.textsecuregcm.workers.CheckDynamicConfigurationCommand; import org.whispersystems.textsecuregcm.workers.DeleteUserCommand; import org.whispersystems.textsecuregcm.workers.IdleDeviceNotificationSchedulerFactory; +import org.whispersystems.textsecuregcm.workers.LockAccountsWithoutPniIdentityKeysCommand; import org.whispersystems.textsecuregcm.workers.LockAccountsWithoutPqKeysCommand; import org.whispersystems.textsecuregcm.workers.MessagePersisterServiceCommand; import org.whispersystems.textsecuregcm.workers.NotifyIdleDevicesCommand; import org.whispersystems.textsecuregcm.workers.ProcessScheduledJobsServiceCommand; +import org.whispersystems.textsecuregcm.workers.RemoveAccountsWithoutPniIdentityKeysCommand; import org.whispersystems.textsecuregcm.workers.RemoveAccountsWithoutPqKeysCommand; import org.whispersystems.textsecuregcm.workers.RemoveExpiredAccountsCommand; import org.whispersystems.textsecuregcm.workers.RemoveExpiredBackupsCommand; import org.whispersystems.textsecuregcm.workers.RemoveExpiredLinkedDevicesCommand; import org.whispersystems.textsecuregcm.workers.RemoveExpiredUsernameHoldsCommand; +import org.whispersystems.textsecuregcm.workers.RemoveLinkedDevicesWithoutPniKeysCommand; import org.whispersystems.textsecuregcm.workers.RemoveLinkedDevicesWithoutPqKeysCommand; import org.whispersystems.textsecuregcm.workers.ScheduledApnPushNotificationSenderServiceCommand; import org.whispersystems.textsecuregcm.workers.ServerVersionCommand; @@ -342,6 +345,10 @@ public class WhisperServerService extends Application accounts) { + final boolean dryRun = getNamespace().getBoolean(DRY_RUN_ARGUMENT); + final int maxConcurrency = getNamespace().getInt(MAX_CONCURRENCY_ARGUMENT); + final int maxRetries = getNamespace().getInt(RETRIES_ARGUMENT); + + final AccountsManager accountsManager = getCommandDependencies().accountsManager(); + + accounts + .filter(account -> account.getIdentityKey(IdentityType.PNI) == null) + .flatMap(accountWithoutPniIdentityKey -> { + final String platform = DevicePlatformUtil.getDevicePlatform(accountWithoutPniIdentityKey.getPrimaryDevice()) + .map(Enum::name) + .orElse("unknown"); + + return dryRun + ? Mono.just(platform) + : Mono.fromFuture(() -> accountsManager.updateAsync(accountWithoutPniIdentityKey, Account::lockAuthTokenHash)) + .retryWhen(Retry.backoff(maxRetries, Duration.ofSeconds(1)) + .onRetryExhaustedThrow((spec, rs) -> rs.failure())) + .thenReturn(platform) + .onErrorResume(throwable -> { + log.warn("Failed to lock account without PNI identity key: {}", + accountWithoutPniIdentityKey.getIdentifier(IdentityType.ACI), throwable); + + return Mono.empty(); + }); + }, maxConcurrency) + .doOnNext(deletedAccountPlatform -> { + Metrics.counter(LOCKED_ACCOUNT_COUNTER_NAME, + "dryRun", String.valueOf(dryRun), + "platform", deletedAccountPlatform) + .increment(); + }) + .then() + .block(); + } +} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPniIdentityKeysCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPniIdentityKeysCommand.java new file mode 100644 index 000000000..5673c97d3 --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPniIdentityKeysCommand.java @@ -0,0 +1,112 @@ +/* + * Copyright 2025 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.workers; + +import com.google.common.annotations.VisibleForTesting; +import io.micrometer.core.instrument.Metrics; +import net.sourceforge.argparse4j.inf.Subparser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.whispersystems.textsecuregcm.identity.IdentityType; +import org.whispersystems.textsecuregcm.metrics.DevicePlatformUtil; +import org.whispersystems.textsecuregcm.metrics.MetricsUtil; +import org.whispersystems.textsecuregcm.storage.Account; +import org.whispersystems.textsecuregcm.storage.AccountsManager; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.util.retry.Retry; +import java.time.Duration; + +public class RemoveAccountsWithoutPniIdentityKeysCommand extends AbstractSinglePassCrawlAccountsCommand { + + @VisibleForTesting + static final String DRY_RUN_ARGUMENT = "dry-run"; + + @VisibleForTesting + static final String MAX_CONCURRENCY_ARGUMENT = "max-concurrency"; + + @VisibleForTesting + static final String RETRIES_ARGUMENT = "retries"; + + private static final String REMOVED_ACCOUNT_COUNTER_NAME = + MetricsUtil.name(LockAccountsWithoutPqKeysCommand.class, "removedAccount"); + + private static final Logger log = LoggerFactory.getLogger(RemoveAccountsWithoutPniIdentityKeysCommand.class); + + public RemoveAccountsWithoutPniIdentityKeysCommand() { + super("remove-accounts-without-pni-identity-keys", "Deletes accounts without PNI identity keys"); + } + + @Override + public void configure(final Subparser subparser) { + super.configure(subparser); + + subparser.addArgument("--dry-run") + .type(Boolean.class) + .dest(DRY_RUN_ARGUMENT) + .required(false) + .setDefault(true) + .help("If true, don’t actually lock accounts with expired linked devices"); + + subparser.addArgument("--max-concurrency") + .type(Integer.class) + .dest(MAX_CONCURRENCY_ARGUMENT) + .setDefault(16) + .help("Max concurrency for DynamoDB operations"); + + subparser.addArgument("--retries") + .type(Integer.class) + .dest(RETRIES_ARGUMENT) + .setDefault(3) + .help("Maximum number of DynamoDB retries permitted per device"); + } + + @Override + protected void crawlAccounts(final Flux accounts) { + final boolean dryRun = getNamespace().getBoolean(DRY_RUN_ARGUMENT); + final int maxConcurrency = getNamespace().getInt(MAX_CONCURRENCY_ARGUMENT); + final int maxRetries = getNamespace().getInt(RETRIES_ARGUMENT); + + final AccountsManager accountsManager = getCommandDependencies().accountsManager(); + + accounts + .filter(account -> account.getIdentityKey(IdentityType.PNI) == null) + .filter(accountWithoutPniIdentityKey -> { + if (!accountWithoutPniIdentityKey.hasLockedCredentials()) { + log.warn("Account {} is not locked", accountWithoutPniIdentityKey.getIdentifier(IdentityType.ACI)); + return false; + } + + return true; + }) + .flatMap(accountWithoutPniIdentityKey -> { + final String platform = DevicePlatformUtil.getDevicePlatform(accountWithoutPniIdentityKey.getPrimaryDevice()) + .map(Enum::name) + .orElse("unknown"); + + return dryRun + ? Mono.just(platform) + : Mono.fromFuture(() -> accountsManager.delete(accountWithoutPniIdentityKey, AccountsManager.DeletionReason.ADMIN_DELETED)) + .retryWhen(Retry.backoff(maxRetries, Duration.ofSeconds(1)) + .onRetryExhaustedThrow((spec, rs) -> rs.failure())) + .thenReturn(platform) + .onErrorResume(throwable -> { + log.warn("Failed to delete account without PNI identity key: {}", + accountWithoutPniIdentityKey.getIdentifier(IdentityType.ACI), throwable); + + return Mono.empty(); + }); + }, maxConcurrency) + .doOnNext(deletedAccountPlatform -> { + Metrics.counter(REMOVED_ACCOUNT_COUNTER_NAME, + "dryRun", String.valueOf(dryRun), + "platform", deletedAccountPlatform) + .increment(); + }) + .then() + .block(); + } +} diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPniKeysCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPniKeysCommand.java new file mode 100644 index 000000000..ce1e89bc4 --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPniKeysCommand.java @@ -0,0 +1,122 @@ +/* + * Copyright 2025 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.workers; + +import com.google.common.annotations.VisibleForTesting; +import io.micrometer.core.instrument.Metrics; +import io.micrometer.shaded.reactor.util.function.Tuples; +import java.time.Duration; +import net.sourceforge.argparse4j.inf.Subparser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.whispersystems.textsecuregcm.identity.IdentityType; +import org.whispersystems.textsecuregcm.metrics.DevicePlatformUtil; +import org.whispersystems.textsecuregcm.metrics.MetricsUtil; +import org.whispersystems.textsecuregcm.storage.Account; +import org.whispersystems.textsecuregcm.storage.AccountsManager; +import org.whispersystems.textsecuregcm.storage.Device; +import org.whispersystems.textsecuregcm.storage.KeysManager; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.util.retry.Retry; + +public class RemoveLinkedDevicesWithoutPniKeysCommand extends AbstractSinglePassCrawlAccountsCommand { + + @VisibleForTesting + static final String DRY_RUN_ARGUMENT = "dry-run"; + + @VisibleForTesting + static final String MAX_CONCURRENCY_ARGUMENT = "max-concurrency"; + + @VisibleForTesting + static final String RETRIES_ARGUMENT = "retries"; + + private static final String REMOVED_DEVICE_COUNTER_NAME = + MetricsUtil.name(RemoveLinkedDevicesWithoutPqKeysCommand.class, "removedDevice"); + + private static final Logger log = LoggerFactory.getLogger(RemoveLinkedDevicesWithoutPniKeysCommand.class); + + public RemoveLinkedDevicesWithoutPniKeysCommand() { + super("remove-linked-devices-without-pni-keys", "Removes linked devices that do not have PNI signed pre-keys"); + } + + @Override + public void configure(final Subparser subparser) { + super.configure(subparser); + + subparser.addArgument("--dry-run") + .type(Boolean.class) + .dest(DRY_RUN_ARGUMENT) + .required(false) + .setDefault(true) + .help("If true, don’t actually modify accounts with expired linked devices"); + + subparser.addArgument("--max-concurrency") + .type(Integer.class) + .dest(MAX_CONCURRENCY_ARGUMENT) + .setDefault(16) + .help("Max concurrency for DynamoDB operations"); + + subparser.addArgument("--retries") + .type(Integer.class) + .dest(RETRIES_ARGUMENT) + .setDefault(3) + .help("Maximum number of DynamoDB retries permitted per device"); + } + + @Override + protected void crawlAccounts(final Flux accounts) { + final boolean dryRun = getNamespace().getBoolean(DRY_RUN_ARGUMENT); + final int maxConcurrency = getNamespace().getInt(MAX_CONCURRENCY_ARGUMENT); + final int maxRetries = getNamespace().getInt(RETRIES_ARGUMENT); + + final AccountsManager accountsManager = getCommandDependencies().accountsManager(); + final KeysManager keysManager = getCommandDependencies().keysManager(); + + accounts + .filter(account -> account.getDevices().size() > 1) + .flatMap(account -> Flux.fromIterable(account.getDevices()) + .filter(device -> !device.isPrimary()) + .map(device -> Tuples.of(account, device))) + .flatMap(accountAndDevice -> { + final Account account = accountAndDevice.getT1(); + final Device device = accountAndDevice.getT2(); + + return Mono.fromFuture( + () -> keysManager.getEcSignedPreKey(account.getIdentifier(IdentityType.PNI), device.getId())) + .retryWhen(Retry.backoff(maxRetries, Duration.ofSeconds(1)) + .onRetryExhaustedThrow((spec, rs) -> rs.failure())) + .onErrorResume(throwable -> { + log.warn("Failed to get PNI signed pre-key presence for account/device: {}:{}", + account.getIdentifier(IdentityType.ACI), device.getId()); + return Mono.empty(); + }) + .map(maybeEcSignedPreKey -> Tuples.of(account, device, maybeEcSignedPreKey)); + }, maxConcurrency) + .filter(tuple -> tuple.getT3().isEmpty()) + .flatMap(accountAndDevice -> { + final Account account = accountAndDevice.getT1(); + final Device device = accountAndDevice.getT2(); + + return dryRun + ? Mono.just(device) + : Mono.fromFuture(() -> accountsManager.removeDevice(account, device.getId())) + .retryWhen(Retry.backoff(maxRetries, Duration.ofSeconds(1)) + .onRetryExhaustedThrow((spec, rs) -> rs.failure())) + .onErrorResume(throwable -> { + log.warn("Failed to remove device: {}:{}", account.getIdentifier(IdentityType.ACI), device.getId()); + return Mono.empty(); + }) + .then(Mono.just(device)); + }, maxConcurrency) + .doOnNext(removedDevice -> Metrics.counter(REMOVED_DEVICE_COUNTER_NAME, + "dryRun", String.valueOf(dryRun), + "platform", DevicePlatformUtil.getDevicePlatform(removedDevice).map(Enum::name).orElse("unknown")) + .increment()) + .then() + .block(); + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPniIdentityKeysCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPniIdentityKeysCommandTest.java new file mode 100644 index 000000000..0880d9b27 --- /dev/null +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPniIdentityKeysCommandTest.java @@ -0,0 +1,119 @@ +/* + * Copyright 2025 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.workers; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.mockito.Mockito.when; + +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.function.Consumer; +import net.sourceforge.argparse4j.inf.Namespace; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.signal.libsignal.protocol.IdentityKey; +import org.whispersystems.textsecuregcm.identity.IdentityType; +import org.whispersystems.textsecuregcm.storage.Account; +import org.whispersystems.textsecuregcm.storage.AccountsManager; +import org.whispersystems.textsecuregcm.storage.Device; +import org.whispersystems.textsecuregcm.storage.KeysManager; +import reactor.core.publisher.Flux; + +class LockAccountsWithoutPniIdentityKeysCommandTest { + + private AccountsManager accountsManager; + + private static class TestLockAccountsWithoutPniIdentityKeysCommand extends LockAccountsWithoutPniIdentityKeysCommand { + + private final CommandDependencies commandDependencies; + private final Namespace namespace; + + TestLockAccountsWithoutPniIdentityKeysCommand(final AccountsManager accountsManager, + final boolean dryRun) { + + commandDependencies = new CommandDependencies(accountsManager, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null); + + namespace = new Namespace(Map.of( + LockAccountsWithoutPqKeysCommand.DRY_RUN_ARGUMENT, dryRun, + LockAccountsWithoutPqKeysCommand.MAX_CONCURRENCY_ARGUMENT, 16, + LockAccountsWithoutPqKeysCommand.RETRIES_ARGUMENT, 3)); + } + + @Override + protected CommandDependencies getCommandDependencies() { + return commandDependencies; + } + + @Override + protected Namespace getNamespace() { + return namespace; + } + } + + @BeforeEach + void setUp() { + accountsManager = mock(AccountsManager.class); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void crawlAccounts(final boolean dryRun) { + final Account accountWithPniIdentityKey = mock(Account.class); + when(accountWithPniIdentityKey.getIdentityKey(IdentityType.PNI)).thenReturn(mock(IdentityKey.class)); + when(accountWithPniIdentityKey.getPrimaryDevice()).thenReturn(mock(Device.class)); + + final Account accountWithoutPniIdentityKey = mock(Account.class); + when(accountWithoutPniIdentityKey.getIdentityKey(IdentityType.PNI)).thenReturn(null); + when(accountWithoutPniIdentityKey.getPrimaryDevice()).thenReturn(mock(Device.class)); + + when(accountsManager.updateAsync(any(), any())).thenAnswer(invocation -> { + final Account account = invocation.getArgument(0); + final Consumer updater = invocation.getArgument(1); + + updater.accept(account); + + return CompletableFuture.completedFuture(account); + }); + + final LockAccountsWithoutPniIdentityKeysCommand lockAccountsWithoutPniIdentityKeysCommand = + new TestLockAccountsWithoutPniIdentityKeysCommand(accountsManager, dryRun); + + lockAccountsWithoutPniIdentityKeysCommand.crawlAccounts( + Flux.just(accountWithPniIdentityKey, accountWithoutPniIdentityKey)); + + if (!dryRun) { + verify(accountsManager).updateAsync(eq(accountWithoutPniIdentityKey), any()); + verify(accountWithoutPniIdentityKey).lockAuthTokenHash(); + } + + verifyNoMoreInteractions(accountsManager); + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPniIdentityKeysCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPniIdentityKeysCommandTest.java new file mode 100644 index 000000000..76d45ffaa --- /dev/null +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPniIdentityKeysCommandTest.java @@ -0,0 +1,112 @@ +/* + * Copyright 2025 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.workers; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.mockito.Mockito.when; + +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import net.sourceforge.argparse4j.inf.Namespace; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.signal.libsignal.protocol.IdentityKey; +import org.whispersystems.textsecuregcm.identity.IdentityType; +import org.whispersystems.textsecuregcm.storage.Account; +import org.whispersystems.textsecuregcm.storage.AccountsManager; +import org.whispersystems.textsecuregcm.storage.Device; +import reactor.core.publisher.Flux; + +class RemoveAccountsWithoutPniIdentityKeysCommandTest { + + private AccountsManager accountsManager; + + private static class TestRemoveAccountsWithoutPniIdentityKeysCommand extends RemoveAccountsWithoutPniIdentityKeysCommand { + + private final CommandDependencies commandDependencies; + private final Namespace namespace; + + TestRemoveAccountsWithoutPniIdentityKeysCommand(final AccountsManager accountsManager, + final boolean dryRun) { + + commandDependencies = new CommandDependencies(accountsManager, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null); + + namespace = new Namespace(Map.of( + LockAccountsWithoutPqKeysCommand.DRY_RUN_ARGUMENT, dryRun, + LockAccountsWithoutPqKeysCommand.MAX_CONCURRENCY_ARGUMENT, 16, + LockAccountsWithoutPqKeysCommand.RETRIES_ARGUMENT, 3)); + } + + @Override + protected CommandDependencies getCommandDependencies() { + return commandDependencies; + } + + @Override + protected Namespace getNamespace() { + return namespace; + } + } + + @BeforeEach + void setUp() { + accountsManager = mock(AccountsManager.class); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void crawlAccounts(final boolean dryRun) { + final Account accountWithPniIdentityKey = mock(Account.class); + when(accountWithPniIdentityKey.getIdentityKey(IdentityType.PNI)).thenReturn(mock(IdentityKey.class)); + when(accountWithPniIdentityKey.getPrimaryDevice()).thenReturn(mock(Device.class)); + + final Account lockedAccountWithoutPniIdentityKey = mock(Account.class); + when(lockedAccountWithoutPniIdentityKey.hasLockedCredentials()).thenReturn(true); + when(lockedAccountWithoutPniIdentityKey.getIdentityKey(IdentityType.PNI)).thenReturn(null); + when(lockedAccountWithoutPniIdentityKey.getPrimaryDevice()).thenReturn(mock(Device.class)); + + final Account unlockedAccountWithoutPniIdentityKey = mock(Account.class); + when(unlockedAccountWithoutPniIdentityKey.hasLockedCredentials()).thenReturn(false); + when(unlockedAccountWithoutPniIdentityKey.getIdentityKey(IdentityType.PNI)).thenReturn(null); + when(unlockedAccountWithoutPniIdentityKey.getPrimaryDevice()).thenReturn(mock(Device.class)); + + when(accountsManager.delete(any(), any())).thenReturn(CompletableFuture.completedFuture(null)); + + final RemoveAccountsWithoutPniIdentityKeysCommand removeAccountsWithoutPniIdentityKeysCommand = + new TestRemoveAccountsWithoutPniIdentityKeysCommand(accountsManager, dryRun); + + removeAccountsWithoutPniIdentityKeysCommand.crawlAccounts(Flux.just( + accountWithPniIdentityKey, lockedAccountWithoutPniIdentityKey, unlockedAccountWithoutPniIdentityKey)); + + if (!dryRun) { + verify(accountsManager).delete(lockedAccountWithoutPniIdentityKey, AccountsManager.DeletionReason.ADMIN_DELETED); + } + + verifyNoMoreInteractions(accountsManager); + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPniKeysCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPniKeysCommandTest.java new file mode 100644 index 000000000..7e3e8bc6d --- /dev/null +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPniKeysCommandTest.java @@ -0,0 +1,142 @@ +/* + * Copyright 2025 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.workers; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyByte; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.mockito.Mockito.when; + +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import net.sourceforge.argparse4j.inf.Namespace; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.whispersystems.textsecuregcm.entities.ECSignedPreKey; +import org.whispersystems.textsecuregcm.identity.IdentityType; +import org.whispersystems.textsecuregcm.storage.Account; +import org.whispersystems.textsecuregcm.storage.AccountsManager; +import org.whispersystems.textsecuregcm.storage.Device; +import org.whispersystems.textsecuregcm.storage.KeysManager; +import reactor.core.publisher.Flux; + +class RemoveLinkedDevicesWithoutPniKeysCommandTest { + + private AccountsManager accountsManager; + private KeysManager keysManager; + + private static class TestRemoveLinkedDevicesWithoutPniKeysCommand extends RemoveLinkedDevicesWithoutPniKeysCommand { + + private final CommandDependencies commandDependencies; + private final Namespace namespace; + + TestRemoveLinkedDevicesWithoutPniKeysCommand(final AccountsManager accountsManager, + final KeysManager keysManager, + final boolean dryRun) { + + commandDependencies = new CommandDependencies(accountsManager, + null, + null, + null, + null, + keysManager, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null); + + namespace = new Namespace(Map.of( + RemoveLinkedDevicesWithoutPqKeysCommand.DRY_RUN_ARGUMENT, dryRun, + RemoveLinkedDevicesWithoutPqKeysCommand.MAX_CONCURRENCY_ARGUMENT, 16, + RemoveLinkedDevicesWithoutPqKeysCommand.RETRIES_ARGUMENT, 3)); + } + + @Override + protected CommandDependencies getCommandDependencies() { + return commandDependencies; + } + + @Override + protected Namespace getNamespace() { + return namespace; + } + } + + @BeforeEach + void setUp() { + accountsManager = mock(AccountsManager.class); + keysManager = mock(KeysManager.class); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void crawlAccounts(final boolean dryRun) { + final Account accountWithoutLinkedDevices = mock(Account.class); + when(accountWithoutLinkedDevices.getDevices()).thenReturn(List.of(mock(Device.class))); + + final Device primaryDevice = mock(Device.class); + when(primaryDevice.isPrimary()).thenReturn(true); + + final byte deviceIdWithPniKey = Device.PRIMARY_ID + 1; + + final Device linkedDeviceWithPniKey = mock(Device.class); + when(linkedDeviceWithPniKey.isPrimary()).thenReturn(false); + when(linkedDeviceWithPniKey.getId()).thenReturn(deviceIdWithPniKey); + + final UUID pniWithLinkedDeviceWithPniKey = UUID.randomUUID(); + + final Account accountWithLinkedDeviceWithPniKey = mock(Account.class); + when(accountWithLinkedDeviceWithPniKey.getIdentifier(IdentityType.PNI)).thenReturn(pniWithLinkedDeviceWithPniKey); + when(accountWithLinkedDeviceWithPniKey.getDevices()).thenReturn(List.of(primaryDevice, linkedDeviceWithPniKey)); + + final byte deviceIdWithoutPniKey = deviceIdWithPniKey + 1; + + final Device linkedDeviceWithoutPniKey = mock(Device.class); + when(linkedDeviceWithoutPniKey.isPrimary()).thenReturn(false); + when(linkedDeviceWithoutPniKey.getId()).thenReturn(deviceIdWithoutPniKey); + + final Account accountWithLinkedDeviceWithoutPniKey = mock(Account.class); + when(accountWithLinkedDeviceWithoutPniKey.getDevices()).thenReturn(List.of(primaryDevice, linkedDeviceWithoutPniKey)); + + when(accountsManager.removeDevice(any(), anyByte())).thenAnswer(invocation -> { + final Account account = invocation.getArgument(0); + return CompletableFuture.completedFuture(account); + }); + + when(keysManager.getEcSignedPreKey(any(), anyByte())) + .thenReturn(CompletableFuture.completedFuture(Optional.empty())); + + when(keysManager.getEcSignedPreKey(pniWithLinkedDeviceWithPniKey, deviceIdWithPniKey)) + .thenReturn(CompletableFuture.completedFuture(Optional.of(mock(ECSignedPreKey.class)))); + + final RemoveLinkedDevicesWithoutPniKeysCommand removeLinkedDevicesWithoutPniKeysCommand = + new TestRemoveLinkedDevicesWithoutPniKeysCommand(accountsManager, keysManager, dryRun); + + removeLinkedDevicesWithoutPniKeysCommand.crawlAccounts(Flux.just( + accountWithoutLinkedDevices, accountWithLinkedDeviceWithPniKey, accountWithLinkedDeviceWithoutPniKey)); + + if (!dryRun) { + verify(accountsManager).removeDevice(accountWithLinkedDeviceWithoutPniKey, deviceIdWithoutPniKey); + } + + verifyNoMoreInteractions(accountsManager); + } +}