Remove commands for removing accounts/devices without PQ or PNI key material

This commit is contained in:
Jon Chambers 2025-05-05 15:04:46 -04:00 committed by Jon Chambers
parent dee3723d97
commit e43487155f
13 changed files with 0 additions and 1459 deletions

View File

@ -266,19 +266,13 @@ 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;
import org.whispersystems.textsecuregcm.workers.SetRequestLoggingEnabledTask;
@ -341,14 +335,6 @@ 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());
bootstrap.addCommand(new RemoveLinkedDevicesWithoutPniKeysCommand());
bootstrap.addCommand(new LockAccountsWithoutPniIdentityKeysCommand());
bootstrap.addCommand(new RemoveAccountsWithoutPniIdentityKeysCommand());
}
@Override

View File

@ -1,105 +0,0 @@
/*
* 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 LockAccountsWithoutPniIdentityKeysCommand 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(LockAccountsWithoutPniIdentityKeysCommand.class, "lockedAccount");
private static final Logger log = LoggerFactory.getLogger(LockAccountsWithoutPniIdentityKeysCommand.class);
public LockAccountsWithoutPniIdentityKeysCommand() {
super("lock-accounts-without-pni-identity-keys", "Locks 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, dont 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();
accounts
.filter(account -> !account.hasLockedCredentials())
.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();
}
}

View File

@ -1,103 +0,0 @@
/*
* 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, dont 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();
}
}

View File

@ -1,112 +0,0 @@
/*
* 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(RemoveAccountsWithoutPniIdentityKeysCommand.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, dont 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();
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();
}
}

View File

@ -1,121 +0,0 @@
/*
* 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, dont 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();
}
}

View File

@ -1,123 +0,0 @@
/*
* 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(RemoveLinkedDevicesWithoutPniKeysCommand.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, dont 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
.filter(account -> !account.hasLockedCredentials())
.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();
}
}

View File

@ -1,111 +0,0 @@
/*
* 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, dont 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
.filter(account -> account.getDevices().size() > 1)
.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();
}
}

View File

@ -1,120 +0,0 @@
/*
* 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,
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<Account> 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);
}
}

View File

@ -1,132 +0,0 @@
/*
* 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,
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();
}
}
}

View File

@ -1,113 +0,0 @@
/*
* 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,
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);
}
}

View File

@ -1,132 +0,0 @@
/*
* 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,
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());
}
}
}

View File

@ -1,143 +0,0 @@
/*
* 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,
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);
}
}

View File

@ -1,130 +0,0 @@
/*
* 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,
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);
}
}
}