Add commands for removing devices without PQ keys

service/src/test/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPqKeysCommandTest.java

@@ -0,0 +1,131 @@
+/*
+ * 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.ArgumentMatchers.eq;
+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.Optional;
+import java.util.UUID;
+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.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
+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 LockAccountsWithoutPqKeysCommandTest {
+
+  private AccountsManager accountsManager;
+  private KeysManager keysManager;
+
+  private static class TestLockAccountsWithoutPqKeysCommand extends LockAccountsWithoutPqKeysCommand {
+
+    private final CommandDependencies commandDependencies;
+    private final Namespace namespace;
+
+    TestLockAccountsWithoutPqKeysCommand(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(
+          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);
+    keysManager = mock(KeysManager.class);
+  }
+
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  void crawlAccounts(final boolean dryRun) {
+    final UUID accountIdentifierWithPqKeys = UUID.randomUUID();
+
+    final Account accountWithPqKeys = mock(Account.class);
+    when(accountWithPqKeys.getIdentifier(IdentityType.ACI)).thenReturn(accountIdentifierWithPqKeys);
+
+    final Account accountWithoutPqKeys = mock(Account.class);
+    when(accountWithoutPqKeys.getIdentifier(IdentityType.ACI)).thenReturn(UUID.randomUUID());
+    when(accountWithoutPqKeys.getPrimaryDevice()).thenReturn(mock(Device.class));
+
+    when(keysManager.getLastResort(any(), anyByte())).thenReturn(CompletableFuture.completedFuture(Optional.empty()));
+    when(keysManager.getLastResort(accountIdentifierWithPqKeys, Device.PRIMARY_ID))
+        .thenReturn(CompletableFuture.completedFuture(Optional.of(mock(KEMSignedPreKey.class))));
+
+    when(accountsManager.delete(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
+    when(accountsManager.updateAsync(any(), any())).thenAnswer(invocation -> {
+      final Account account = invocation.getArgument(0);
+      final Consumer<Account> updater = invocation.getArgument(1);
+
+      updater.accept(account);
+
+      return CompletableFuture.completedFuture(account);
+    });
+
+    final LockAccountsWithoutPqKeysCommand lockAccountsWithoutPqKeysCommand =
+        new TestLockAccountsWithoutPqKeysCommand(accountsManager, keysManager, dryRun);
+
+    lockAccountsWithoutPqKeysCommand.crawlAccounts(Flux.just(accountWithPqKeys, accountWithoutPqKeys));
+
+    if (dryRun) {
+      verify(accountsManager, never()).updateAsync(any(), any());
+    } else {
+      verify(accountsManager).updateAsync(eq(accountWithoutPqKeys), any());
+      verifyNoMoreInteractions(accountsManager);
+
+      verify(accountWithoutPqKeys).lockAuthTokenHash();
+    }
+  }
+}

service/src/test/java/org/whispersystems/textsecuregcm/workers/PqKeysUtilTest.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.workers;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyByte;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.storage.KeysManager;
+import reactor.core.publisher.Flux;
+
+class PqKeysUtilTest {
+
+  private KeysManager keysManager;
+  private PqKeysUtil pqKeysUtil;
+
+  @BeforeEach
+  void setUp() {
+    keysManager = mock(KeysManager.class);
+    pqKeysUtil = new PqKeysUtil(keysManager, 16, 3);
+  }
+
+  @Test
+  void getAccountsWithoutPqKeys() {
+    final UUID aciWithPqKeys = UUID.randomUUID();
+
+    final Account accountWithPqKeys = mock(Account.class);
+    when(accountWithPqKeys.getIdentifier(IdentityType.ACI)).thenReturn(aciWithPqKeys);
+
+    final Account accountWithoutPqKeys = mock(Account.class);
+    when(accountWithoutPqKeys.getIdentifier(IdentityType.ACI)).thenReturn(UUID.randomUUID());
+
+    when(keysManager.getLastResort(any(), anyByte())).thenReturn(CompletableFuture.completedFuture(Optional.empty()));
+    when(keysManager.getLastResort(aciWithPqKeys, Device.PRIMARY_ID))
+        .thenReturn(CompletableFuture.completedFuture(Optional.of(mock(KEMSignedPreKey.class))));
+
+    assertEquals(List.of(accountWithoutPqKeys),
+        Flux.just(accountWithPqKeys, accountWithoutPqKeys)
+            .transform(pqKeysUtil::getAccountsWithoutPqKeys)
+            .collectList()
+            .block());
+  }
+}

service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPqKeysCommandTest.java

@@ -0,0 +1,131 @@
+/*
+ * 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.ArgumentMatchers.eq;
+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.when;
+
+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.KEMSignedPreKey;
+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 RemoveAccountsWithoutPqKeysCommandTest {
+
+  private AccountsManager accountsManager;
+  private KeysManager keysManager;
+
+  private static class TestRemoveAccountsWithoutPqKeysCommand extends RemoveAccountsWithoutPqKeysCommand {
+
+    private final CommandDependencies commandDependencies;
+    private final Namespace namespace;
+
+    TestRemoveAccountsWithoutPqKeysCommand(final AccountsManager accountsManager,
+        final KeysManager keysManager,
+        final int maxAccounts,
+        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(
+          RemoveAccountsWithoutPqKeysCommand.DRY_RUN_ARGUMENT, dryRun,
+          RemoveAccountsWithoutPqKeysCommand.MAX_CONCURRENCY_ARGUMENT, 16,
+          RemoveAccountsWithoutPqKeysCommand.RETRIES_ARGUMENT, 3,
+          RemoveAccountsWithoutPqKeysCommand.MAX_ACCOUNTS_ARGUMENT, maxAccounts));
+    }
+
+    @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 UUID accountIdentifierWithPqKeys = UUID.randomUUID();
+
+    final Account accountWithPqKeys = mock(Account.class);
+    when(accountWithPqKeys.getIdentifier(IdentityType.ACI)).thenReturn(accountIdentifierWithPqKeys);
+    when(accountWithPqKeys.getPrimaryDevice()).thenReturn(mock(Device.class));
+    when(accountWithPqKeys.hasLockedCredentials()).thenReturn(true);
+
+    when(keysManager.getLastResort(any(), anyByte())).thenReturn(CompletableFuture.completedFuture(Optional.empty()));
+    when(keysManager.getLastResort(accountIdentifierWithPqKeys, Device.PRIMARY_ID))
+        .thenReturn(CompletableFuture.completedFuture(Optional.of(mock(KEMSignedPreKey.class))));
+
+    when(accountsManager.delete(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
+
+    final int maxAccounts = 5;
+
+    final RemoveAccountsWithoutPqKeysCommand removeAccountsWithoutPqKeysCommand =
+        new TestRemoveAccountsWithoutPqKeysCommand(accountsManager, keysManager, maxAccounts, dryRun);
+
+    removeAccountsWithoutPqKeysCommand.crawlAccounts(Flux.concat(
+        Flux.just(accountWithPqKeys),
+        Flux.generate(sink -> {
+          final Account accountWithoutPqKeys = mock(Account.class);
+          when(accountWithoutPqKeys.getIdentifier(IdentityType.ACI)).thenReturn(UUID.randomUUID());
+          when(accountWithoutPqKeys.getPrimaryDevice()).thenReturn(mock(Device.class));
+          when(accountWithoutPqKeys.hasLockedCredentials()).thenReturn(true);
+
+          sink.next(accountWithoutPqKeys);
+        })));
+
+    if (dryRun) {
+      verify(accountsManager, never()).delete(any(), any());
+    } else {
+      verify(accountsManager, times(maxAccounts)).delete(any(), eq(AccountsManager.DeletionReason.ADMIN_DELETED));
+      verify(accountsManager, never()).delete(eq(accountWithPqKeys), any());
+    }
+  }
+}

service/src/test/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPqKeysCommandTest.java

@@ -0,0 +1,129 @@
+/*
+ * 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.never;
+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.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.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 RemoveLinkedDevicesWithoutPqKeysCommandTest {
+
+  private AccountsManager accountsManager;
+  private KeysManager keysManager;
+
+  private static class TestRemoveLinkedDevicesWithoutPqKeysCommand extends RemoveLinkedDevicesWithoutPqKeysCommand {
+
+    private final CommandDependencies commandDependencies;
+    private final Namespace namespace;
+
+    TestRemoveLinkedDevicesWithoutPqKeysCommand(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 UUID accountIdentifier = UUID.randomUUID();
+    final byte deviceIdWithPqKeys = Device.PRIMARY_ID + 1;
+    final byte deviceIdWithoutPqKeys = deviceIdWithPqKeys + 1;
+
+    final Device primaryDevice = mock(Device.class);
+    when(primaryDevice.isPrimary()).thenReturn(true);
+    when(primaryDevice.getId()).thenReturn(Device.PRIMARY_ID);
+
+    final Device linkedDeviceWithPqKeys = mock(Device.class);
+    when(linkedDeviceWithPqKeys.isPrimary()).thenReturn(false);
+    when(linkedDeviceWithPqKeys.getId()).thenReturn(deviceIdWithPqKeys);
+
+    final Device linkedDeviceWithoutPqKeys = mock(Device.class);
+    when(linkedDeviceWithoutPqKeys.isPrimary()).thenReturn(false);
+    when(linkedDeviceWithoutPqKeys.getId()).thenReturn(deviceIdWithoutPqKeys);
+
+    final Account account = mock(Account.class);
+    when(account.getIdentifier(IdentityType.ACI)).thenReturn(accountIdentifier);
+    when(account.getDevices()).thenReturn(List.of(primaryDevice, linkedDeviceWithPqKeys, linkedDeviceWithoutPqKeys));
+
+    when(keysManager.getPqEnabledDevices(accountIdentifier))
+        .thenReturn(CompletableFuture.completedFuture(List.of(deviceIdWithPqKeys)));
+
+    when(accountsManager.removeDevice(any(), anyByte()))
+        .thenAnswer(invocation -> CompletableFuture.completedFuture(invocation.getArgument(0)));
+
+    final RemoveLinkedDevicesWithoutPqKeysCommand removeLinkedDevicesWithoutPqKeysCommand =
+        new TestRemoveLinkedDevicesWithoutPqKeysCommand(accountsManager, keysManager, dryRun);
+
+    removeLinkedDevicesWithoutPqKeysCommand.crawlAccounts(Flux.just(account));
+
+    if (dryRun) {
+      verify(accountsManager, never()).removeDevice(any(), anyByte());
+    } else {
+      verify(accountsManager).removeDevice(account, deviceIdWithoutPqKeys);
+      verifyNoMoreInteractions(accountsManager);
+    }
+  }
+}