Add commands for removing devices without PQ keys

2026-04-19 16:09:02 +01:00 · 2025-04-28 15:45:27 -04:00
parent 91a36f4421
commit cdb651b68f
9 changed files with 837 additions and 0 deletions
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -266,13 +266,16 @@ 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.LockAccountsWithoutPqKeysCommand;
 import org.whispersystems.textsecuregcm.workers.MessagePersisterServiceCommand;
 import org.whispersystems.textsecuregcm.workers.NotifyIdleDevicesCommand;
 import org.whispersystems.textsecuregcm.workers.ProcessScheduledJobsServiceCommand;
+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.RemoveLinkedDevicesWithoutPqKeysCommand;
 import org.whispersystems.textsecuregcm.workers.ScheduledApnPushNotificationSenderServiceCommand;
 import org.whispersystems.textsecuregcm.workers.ServerVersionCommand;
 import org.whispersystems.textsecuregcm.workers.SetRequestLoggingEnabledTask;
@@ -335,6 +338,10 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
    bootstrap.addCommand(new ProcessScheduledJobsServiceCommand("process-idle-device-notification-jobs",
        "Processes scheduled jobs to send notifications to idle devices",
        new IdleDeviceNotificationSchedulerFactory()));
+
+    bootstrap.addCommand(new RemoveLinkedDevicesWithoutPqKeysCommand());
+    bootstrap.addCommand(new LockAccountsWithoutPqKeysCommand());
+    bootstrap.addCommand(new RemoveAccountsWithoutPqKeysCommand());
  }

  @Override
--- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPqKeysCommand.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/LockAccountsWithoutPqKeysCommand.java
@@ -0,0 +1,103 @@
+/*
+ * 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 LockAccountsWithoutPqKeysCommand 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 LOCKED_ACCOUNT_COUNTER_NAME =
+      MetricsUtil.name(LockAccountsWithoutPqKeysCommand.class, "lockedAccount");
+
+  private static final Logger log = LoggerFactory.getLogger(LockAccountsWithoutPqKeysCommand.class);
+
+  public LockAccountsWithoutPqKeysCommand() {
+    super("lock-accounts-without-pq-keys", "Locks accounts with primary devices that don't have PQ 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<Account> 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 PqKeysUtil pqKeysUtil = new PqKeysUtil(getCommandDependencies().keysManager(), maxConcurrency, maxRetries);
+
+    accounts
+        .transform(pqKeysUtil::getAccountsWithoutPqKeys)
+        .flatMap(accountWithoutPqKeys -> {
+          final String platform = DevicePlatformUtil.getDevicePlatform(accountWithoutPqKeys.getPrimaryDevice())
+              .map(Enum::name)
+              .orElse("unknown");
+
+          return dryRun
+              ? Mono.just(platform)
+              : Mono.fromFuture(() -> accountsManager.updateAsync(accountWithoutPqKeys, 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 PQ keys {}", accountWithoutPqKeys.getIdentifier(IdentityType.ACI), throwable);
+                    return Mono.empty();
+                  });
+        })
+        .doOnNext(deletedAccountPlatform -> {
+          Metrics.counter(LOCKED_ACCOUNT_COUNTER_NAME,
+                  "dryRun", String.valueOf(dryRun),
+                  "platform", deletedAccountPlatform)
+              .increment();
+        })
+        .then()
+        .block();
+  }
+}
--- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/PqKeysUtil.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/PqKeysUtil.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.workers;
+
+import java.time.Duration;
+import java.util.Optional;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+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;
+import reactor.core.publisher.Mono;
+import reactor.util.retry.Retry;
+
+class PqKeysUtil {
+
+  private final KeysManager keysManager;
+  private final int maxConcurrency;
+  private final int maxRetries;
+
+  private static final Logger log = LoggerFactory.getLogger(PqKeysUtil.class);
+
+  PqKeysUtil(final KeysManager keysManager, final int maxConcurrency, final int maxRetries) {
+    this.keysManager = keysManager;
+    this.maxConcurrency = maxConcurrency;
+    this.maxRetries = maxRetries;
+  }
+
+  public Flux<Account> getAccountsWithoutPqKeys(final Flux<Account> accounts) {
+    return accounts.flatMap(account -> Mono.fromFuture(
+                    () -> keysManager.getLastResort(account.getIdentifier(IdentityType.ACI), Device.PRIMARY_ID))
+                .retryWhen(Retry.backoff(maxRetries, Duration.ofSeconds(1))
+                    .onRetryExhaustedThrow((spec, rs) -> rs.failure()))
+                .onErrorResume(throwable -> {
+                  log.warn("Failed to get last-resort key for {}", account.getIdentifier(IdentityType.ACI), throwable);
+                  return Mono.empty();
+                })
+                .filter(Optional::isEmpty)
+                .map(ignored -> account),
+            maxConcurrency);
+  }
+}
--- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPqKeysCommand.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/RemoveAccountsWithoutPqKeysCommand.java
@@ -0,0 +1,121 @@
+/*
+ * 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 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 reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.util.retry.Retry;
+
+public class RemoveAccountsWithoutPqKeysCommand 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";
+
+  @VisibleForTesting
+  static final String MAX_ACCOUNTS_ARGUMENT = "max-accounts";
+
+  private static final String REMOVED_ACCOUNT_COUNTER_NAME =
+      MetricsUtil.name(RemoveAccountsWithoutPqKeysCommand.class, "removedAccount");
+
+  private static final Logger log = LoggerFactory.getLogger(RemoveAccountsWithoutPqKeysCommand.class);
+
+  public RemoveAccountsWithoutPqKeysCommand() {
+    super("remove-accounts-without-pq-keys", "Removes accounts with primary devices that don't have PQ 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");
+
+    subparser.addArgument("--max-accounts")
+        .type(Integer.class)
+        .required(true)
+        .dest(MAX_ACCOUNTS_ARGUMENT)
+        .help("Maximum number of accounts to remove per run");
+  }
+
+  @Override
+  protected void crawlAccounts(final Flux<Account> 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 int maxAccounts = getNamespace().getInt(MAX_ACCOUNTS_ARGUMENT);
+
+    final AccountsManager accountsManager = getCommandDependencies().accountsManager();
+    final PqKeysUtil pqKeysUtil = new PqKeysUtil(getCommandDependencies().keysManager(), maxConcurrency, maxRetries);
+
+    accounts
+        .transform(pqKeysUtil::getAccountsWithoutPqKeys)
+        .take(maxAccounts)
+        .filter(accountWithoutPqKeys -> {
+          if (!accountWithoutPqKeys.hasLockedCredentials()) {
+            log.warn("Account {} is not locked", accountWithoutPqKeys.getIdentifier(IdentityType.ACI));
+          }
+
+          return accountWithoutPqKeys.hasLockedCredentials();
+        })
+        .flatMap(accountWithoutPqKeys -> {
+          final String platform = DevicePlatformUtil.getDevicePlatform(accountWithoutPqKeys.getPrimaryDevice())
+              .map(Enum::name)
+              .orElse("unknown");
+
+          return dryRun
+              ? Mono.just(platform)
+              : Mono.fromFuture(() -> accountsManager.delete(accountWithoutPqKeys, AccountsManager.DeletionReason.ADMIN_DELETED))
+                  .retryWhen(Retry.backoff(maxRetries, Duration.ofSeconds(1))
+                      .onRetryExhaustedThrow((spec, rs) -> rs.failure()))
+                  .thenReturn(platform)
+                  .onErrorResume(throwable -> {
+                    log.warn("Failed to remove account without PQ keys {}", accountWithoutPqKeys.getIdentifier(IdentityType.ACI), throwable);
+                    return Mono.empty();
+                  });
+        })
+        .doOnNext(deletedAccountPlatform -> {
+          Metrics.counter(REMOVED_ACCOUNT_COUNTER_NAME,
+                  "dryRun", String.valueOf(dryRun),
+                  "platform", deletedAccountPlatform)
+              .increment();
+        })
+        .then()
+        .block();
+  }
+}
--- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPqKeysCommand.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/RemoveLinkedDevicesWithoutPqKeysCommand.java
@@ -0,0 +1,110 @@
+/*
+ * 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 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.KeysManager;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.util.function.Tuples;
+import reactor.util.retry.Retry;
+
+public class RemoveLinkedDevicesWithoutPqKeysCommand 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(RemoveLinkedDevicesWithoutPqKeysCommand.class);
+
+  public RemoveLinkedDevicesWithoutPqKeysCommand() {
+    super("remove-linked-devices-without-pq-keys", "Removes linked devices that don't have PQ 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<Account> 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
+        .flatMap(
+            account -> Mono.fromFuture(() -> keysManager.getPqEnabledDevices(account.getIdentifier(IdentityType.ACI)))
+                .retryWhen(Retry.backoff(maxRetries, Duration.ofSeconds(1))
+                    .onRetryExhaustedThrow((spec, rs) -> rs.failure()))
+                .onErrorResume(throwable -> {
+                  log.warn("Failed to get PQ key presence for account: {}", account.getIdentifier(IdentityType.ACI));
+                  return Mono.empty();
+                })
+                .flatMapMany(pqEnabledDeviceIds -> Flux.fromIterable(account.getDevices())
+                    .filter(device -> !device.isPrimary())
+                    .filter(device -> !pqEnabledDeviceIds.contains(device.getId()))
+                    .map(device -> Tuples.of(account, device))), maxConcurrency)
+        .flatMap(accountAndDevice -> dryRun
+            ? Mono.just(accountAndDevice.getT2())
+            : Mono.fromFuture(() -> accountsManager.removeDevice(accountAndDevice.getT1(), accountAndDevice.getT2().getId()))
+                .retryWhen(Retry.backoff(maxRetries, Duration.ofSeconds(1))
+                    .onRetryExhaustedThrow((spec, rs) -> rs.failure()))
+                .onErrorResume(throwable -> {
+                  log.warn("Failed to remove linked device without PQ keys: {}:{}",
+                      accountAndDevice.getT1().getIdentifier(IdentityType.ACI), accountAndDevice.getT2().getId());
+
+                  return Mono.empty();
+                })
+                .map(ignored -> accountAndDevice.getT2()), 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();
+  }
+}