diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index de5317998..3912ec94e 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java @@ -272,6 +272,7 @@ import org.whispersystems.textsecuregcm.workers.BackupUsageRecalculationCommand; import org.whispersystems.textsecuregcm.workers.CertificateCommand; import org.whispersystems.textsecuregcm.workers.CheckDynamicConfigurationCommand; import org.whispersystems.textsecuregcm.workers.DeleteUserCommand; +import org.whispersystems.textsecuregcm.workers.EncryptDeviceCreationTimestampCommand; import org.whispersystems.textsecuregcm.workers.IdleDeviceNotificationSchedulerFactory; import org.whispersystems.textsecuregcm.workers.MessagePersisterServiceCommand; import org.whispersystems.textsecuregcm.workers.NotifyIdleDevicesCommand; @@ -346,6 +347,7 @@ public class WhisperServerService extends Application accounts) { + final boolean isDryRun = getNamespace().getBoolean(DRY_RUN_ARGUMENT); + final Counter encryptedTimestampCounter = + Metrics.counter(ENCRYPTED_CREATION_TIMESTAMP_COUNTER_NAME, "dryRun", String.valueOf(isDryRun)); + final Counter processedAccountCounter = + Metrics.counter(PROCESSED_ACCOUNT_COUNTER_NAME, "dryRun", String.valueOf(isDryRun)); + accounts + .flatMap(account -> + Flux.fromIterable(account.getDevices()) + .flatMap(device -> { + final byte[] createdAtCiphertext = EncryptDeviceCreationTimestampUtil.encrypt( + device.getCreated(), account.getIdentityKey(IdentityType.ACI), + device.getId(), device.getRegistrationId(IdentityType.ACI)); + + final Mono encryptTimestampMono = isDryRun + ? Mono.empty() + : Mono.fromFuture(() -> getCommandDependencies().accountsManager().updateDeviceAsync( + account, device.getId(), d -> d.setCreatedAtCiphertext(createdAtCiphertext)) + .thenRun(Util.NOOP)); + return encryptTimestampMono + .doOnSuccess(_ -> encryptedTimestampCounter.increment()) + .retryWhen(Retry.backoff(3, Duration.ofSeconds(1)).maxBackoff(Duration.ofSeconds(4))) + .onErrorResume(throwable -> { + log.warn("Failed to encrypt creation timestamp on device {}, account {}", device.getId(), account.getUuid(), throwable); + return Mono.empty(); + }); + }, MAX_CONCURRENCY), + MAX_CONCURRENCY) + .doOnComplete(processedAccountCounter::increment) + .then() + .block(); + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/EncryptDeviceCreationTimestampCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/EncryptDeviceCreationTimestampCommandTest.java new file mode 100644 index 000000000..cfe08d5e3 --- /dev/null +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/EncryptDeviceCreationTimestampCommandTest.java @@ -0,0 +1,106 @@ +/* + * Copyright 2025 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.workers; + +import net.sourceforge.argparse4j.inf.Namespace; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.mockito.MockedStatic; +import org.signal.libsignal.protocol.IdentityKey; +import org.signal.libsignal.protocol.ecc.ECKeyPair; +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.util.EncryptDeviceCreationTimestampUtil; +import org.whispersystems.textsecuregcm.util.TestRandomUtil; +import reactor.core.publisher.Flux; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyByte; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class EncryptDeviceCreationTimestampCommandTest { + + private static class TestEncryptDeviceCreationTimestampCommand extends EncryptDeviceCreationTimestampCommand { + + private final CommandDependencies commandDependencies; + private final Namespace namespace; + + public TestEncryptDeviceCreationTimestampCommand(final AccountsManager accountsManager, final boolean isDryRun) { + super(); + + commandDependencies = mock(CommandDependencies.class); + when(commandDependencies.accountsManager()).thenReturn(accountsManager); + + namespace = mock(Namespace.class); + when(namespace.getBoolean(EncryptDeviceCreationTimestampCommand.DRY_RUN_ARGUMENT)).thenReturn(isDryRun); + } + + @Override + protected CommandDependencies getCommandDependencies() { + return commandDependencies; + } + + @Override + protected Namespace getNamespace() { + return namespace; + } + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void crawlAccounts(final boolean isDryRun) { + final Account account = mock(Account.class); + final Device device = mock(Device.class); + + final IdentityKey identityKey = new IdentityKey(ECKeyPair.generate().getPublicKey()); + final byte deviceId = (byte) 1; + final long createdAt = System.currentTimeMillis(); + final int registrationId = 123; + + when(account.getDevices()).thenReturn(List.of(device)); + when(account.getIdentityKey(IdentityType.ACI)).thenReturn(identityKey); + when(device.getCreated()).thenReturn(createdAt); + when(device.getId()).thenReturn(deviceId); + when(device.getRegistrationId(IdentityType.ACI)).thenReturn(registrationId); + + final AccountsManager accountsManager = mock(AccountsManager.class); + when(accountsManager.updateDeviceAsync(any(), anyByte(), any())).thenReturn(CompletableFuture.completedFuture(null)); + + final EncryptDeviceCreationTimestampCommand encryptDeviceCreationTimestampCommand = + new TestEncryptDeviceCreationTimestampCommand(accountsManager, isDryRun); + + try (MockedStatic mockUtil = mockStatic(EncryptDeviceCreationTimestampUtil.class)) { + mockUtil.when(() -> EncryptDeviceCreationTimestampUtil.encrypt(anyLong(), any(), anyByte(), anyInt())) + .thenReturn(TestRandomUtil.nextBytes(56)); + + encryptDeviceCreationTimestampCommand.crawlAccounts(Flux.just(account)); + + mockUtil.verify(() -> EncryptDeviceCreationTimestampUtil.encrypt( + eq(createdAt), + eq(identityKey), + eq(deviceId), + eq(registrationId) + )); + + if (isDryRun) { + verify(accountsManager, never()).updateDeviceAsync(any(), anyByte(), any()); + } else { + verify(accountsManager).updateDeviceAsync(eq(account), eq(deviceId), any()); + } + } + } +}