mirror of
https://github.com/signalapp/Signal-Server
synced 2026-04-21 13:38:04 +01:00
Add a crawler for orphaned prekey pages
This commit is contained in:
Ravi Khadiwala
committed by
ravi-signal
ravi-signal
parent
2bb14892af
commit
aaa36fd8f5
@@ -9,37 +9,27 @@ import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertNotEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
import static org.testcontainers.containers.localstack.LocalStackContainer.Service.S3;
|
||||
|
||||
import java.util.Comparator;
|
||||
import java.util.HashSet;
|
||||
import java.util.List;
|
||||
import java.util.Objects;
|
||||
import java.util.Optional;
|
||||
import java.util.Set;
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.ThreadLocalRandom;
|
||||
import java.util.stream.Collectors;
|
||||
import java.util.stream.IntStream;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.RepeatedTest;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.RegisterExtension;
|
||||
import org.signal.libsignal.protocol.ecc.Curve;
|
||||
import org.signal.libsignal.protocol.ecc.ECKeyPair;
|
||||
import org.testcontainers.containers.localstack.LocalStackContainer;
|
||||
import org.testcontainers.junit.jupiter.Container;
|
||||
import org.testcontainers.junit.jupiter.Testcontainers;
|
||||
import org.testcontainers.utility.DockerImageName;
|
||||
import org.whispersystems.textsecuregcm.entities.KEMSignedPreKey;
|
||||
import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
|
||||
import org.whispersystems.textsecuregcm.util.TestRandomUtil;
|
||||
import reactor.core.publisher.Flux;
|
||||
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
|
||||
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
import software.amazon.awssdk.services.s3.S3AsyncClient;
|
||||
import software.amazon.awssdk.services.s3.model.CreateBucketRequest;
|
||||
import software.amazon.awssdk.core.async.AsyncRequestBody;
|
||||
import software.amazon.awssdk.services.s3.model.ListObjectsV2Request;
|
||||
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
|
||||
import software.amazon.awssdk.services.s3.model.S3Object;
|
||||
|
||||
class PagedSingleUseKEMPreKeyStoreTest {
|
||||
@@ -77,7 +67,7 @@ class PagedSingleUseKEMPreKeyStoreTest {
|
||||
assertDoesNotThrow(() -> keyStore.store(accountIdentifier, deviceId, preKeys).join());
|
||||
|
||||
final List<KEMSignedPreKey> sortedPreKeys = preKeys.stream()
|
||||
.sorted(Comparator.comparing(preKey -> preKey.keyId()))
|
||||
.sorted(Comparator.comparing(KEMSignedPreKey::keyId))
|
||||
.toList();
|
||||
|
||||
assertEquals(Optional.of(sortedPreKeys.get(0)), keyStore.take(accountIdentifier, deviceId).join());
|
||||
@@ -91,18 +81,18 @@ class PagedSingleUseKEMPreKeyStoreTest {
|
||||
|
||||
final List<KEMSignedPreKey> preKeys1 = generateRandomPreKeys();
|
||||
keyStore.store(accountIdentifier, deviceId, preKeys1).join();
|
||||
List<String> oldPages = listPages(accountIdentifier).stream().map(S3Object::key).collect(Collectors.toList());
|
||||
List<String> oldPages = listPages(accountIdentifier).stream().map(S3Object::key).toList();
|
||||
assertEquals(1, oldPages.size());
|
||||
|
||||
final List<KEMSignedPreKey> preKeys2 = generateRandomPreKeys();
|
||||
keyStore.store(accountIdentifier, deviceId, preKeys2).join();
|
||||
List<String> newPages = listPages(accountIdentifier).stream().map(S3Object::key).collect(Collectors.toList());
|
||||
List<String> newPages = listPages(accountIdentifier).stream().map(S3Object::key).toList();
|
||||
assertEquals(1, newPages.size());
|
||||
|
||||
assertNotEquals(oldPages.getFirst(), newPages.getFirst());
|
||||
|
||||
assertEquals(
|
||||
preKeys2.stream().sorted(Comparator.comparing(preKey -> preKey.keyId())).toList(),
|
||||
preKeys2.stream().sorted(Comparator.comparing(KEMSignedPreKey::keyId)).toList(),
|
||||
|
||||
IntStream.range(0, preKeys2.size())
|
||||
.mapToObj(i -> keyStore.take(accountIdentifier, deviceId).join())
|
||||
@@ -122,7 +112,7 @@ class PagedSingleUseKEMPreKeyStoreTest {
|
||||
assertDoesNotThrow(() -> keyStore.store(accountIdentifier, deviceId, preKeys).join());
|
||||
|
||||
final List<KEMSignedPreKey> sortedPreKeys = preKeys.stream()
|
||||
.sorted(Comparator.comparing(preKey -> preKey.keyId()))
|
||||
.sorted(Comparator.comparing(KEMSignedPreKey::keyId))
|
||||
.toList();
|
||||
|
||||
for (int i = 0; i < KEY_COUNT; i++) {
|
||||
@@ -171,7 +161,7 @@ class PagedSingleUseKEMPreKeyStoreTest {
|
||||
|
||||
final List<S3Object> pages = listPages(accountIdentifier);
|
||||
assertEquals(1, pages.size());
|
||||
assertTrue(pages.get(0).key().startsWith("%s/%s".formatted(accountIdentifier, deviceId + 1)));
|
||||
assertTrue(pages.getFirst().key().startsWith("%s/%s".formatted(accountIdentifier, deviceId + 1)));
|
||||
}
|
||||
|
||||
@Test
|
||||
@@ -194,6 +184,66 @@ class PagedSingleUseKEMPreKeyStoreTest {
|
||||
assertEquals(0, listPages(accountIdentifier).size());
|
||||
}
|
||||
|
||||
@Test
|
||||
void listPages() {
|
||||
final UUID aci1 = UUID.randomUUID();
|
||||
final UUID aci2 = new UUID(aci1.getMostSignificantBits(), aci1.getLeastSignificantBits() + 1);
|
||||
final byte deviceId = 1;
|
||||
|
||||
keyStore.store(aci1, deviceId, generateRandomPreKeys()).join();
|
||||
keyStore.store(aci1, (byte) (deviceId + 1), generateRandomPreKeys()).join();
|
||||
keyStore.store(aci2, deviceId, generateRandomPreKeys()).join();
|
||||
|
||||
List<DeviceKEMPreKeyPages> stored = keyStore.listStoredPages(1).collectList().block();
|
||||
assertEquals(3, stored.size());
|
||||
for (DeviceKEMPreKeyPages pages : stored) {
|
||||
assertEquals(1, pages.pageIdToLastModified().size());
|
||||
}
|
||||
|
||||
assertEquals(List.of(aci1, aci1, aci2), stored.stream().map(DeviceKEMPreKeyPages::identifier).toList());
|
||||
assertEquals(
|
||||
List.of(deviceId, (byte) (deviceId + 1), deviceId),
|
||||
stored.stream().map(DeviceKEMPreKeyPages::deviceId).toList());
|
||||
}
|
||||
|
||||
@Test
|
||||
void listPagesWithOrphans() {
|
||||
final UUID aci1 = UUID.randomUUID();
|
||||
final UUID aci2 = new UUID(aci1.getMostSignificantBits(), aci1.getLeastSignificantBits() + 1);
|
||||
final byte deviceId = 1;
|
||||
|
||||
// Two orphans
|
||||
keyStore.store(aci1, deviceId, generateRandomPreKeys()).join();
|
||||
writeOrphanedS3Object(aci1, deviceId);
|
||||
writeOrphanedS3Object(aci1, deviceId);
|
||||
|
||||
// No orphans
|
||||
keyStore.store(aci1, (byte) (deviceId + 1), generateRandomPreKeys()).join();
|
||||
|
||||
// One orphan
|
||||
keyStore.store(aci2, deviceId, generateRandomPreKeys()).join();
|
||||
writeOrphanedS3Object(aci2, deviceId);
|
||||
|
||||
// Orphan with no database record
|
||||
writeOrphanedS3Object(aci2, (byte) (deviceId + 2));
|
||||
|
||||
List<DeviceKEMPreKeyPages> stored = keyStore.listStoredPages(1).collectList().block();
|
||||
assertEquals(4, stored.size());
|
||||
|
||||
assertEquals(
|
||||
List.of(3, 1, 2, 1),
|
||||
stored.stream().map(s -> s.pageIdToLastModified().size()).toList());
|
||||
}
|
||||
|
||||
private void writeOrphanedS3Object(final UUID identifier, final byte deviceId) {
|
||||
S3_EXTENSION.getS3Client()
|
||||
.putObject(PutObjectRequest.builder()
|
||||
.bucket(BUCKET_NAME)
|
||||
.key("%s/%s/%s".formatted(identifier, deviceId, UUID.randomUUID())).build(),
|
||||
AsyncRequestBody.fromBytes(TestRandomUtil.nextBytes(10)))
|
||||
.join();
|
||||
}
|
||||
|
||||
private List<S3Object> listPages(final UUID identifier) {
|
||||
return Flux.from(S3_EXTENSION.getS3Client().listObjectsV2Paginator(ListObjectsV2Request.builder()
|
||||
.bucket(BUCKET_NAME)
|
||||
|
||||
@@ -0,0 +1,138 @@
|
||||
/*
|
||||
* 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.anyInt;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.times;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.verifyNoMoreInteractions;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
import io.dropwizard.core.setup.Environment;
|
||||
import java.time.Clock;
|
||||
import java.time.Duration;
|
||||
import java.time.Instant;
|
||||
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.assertj.core.api.Assertions;
|
||||
import org.junit.Test;
|
||||
import org.junit.jupiter.params.ParameterizedTest;
|
||||
import org.junit.jupiter.params.provider.ValueSource;
|
||||
import org.junitpioneer.jupiter.cartesian.CartesianTest;
|
||||
import org.whispersystems.textsecuregcm.WhisperServerConfiguration;
|
||||
import org.whispersystems.textsecuregcm.storage.DeviceKEMPreKeyPages;
|
||||
import org.whispersystems.textsecuregcm.storage.KeysManager;
|
||||
import org.whispersystems.textsecuregcm.util.TestClock;
|
||||
import reactor.core.publisher.Flux;
|
||||
|
||||
public class RemoveOrphanedPreKeyPagesCommandTest {
|
||||
|
||||
|
||||
@ParameterizedTest
|
||||
@ValueSource(booleans = {true, false})
|
||||
public void removeStalePages(boolean dryRun) throws Exception {
|
||||
final TestClock clock = TestClock.pinned(Instant.EPOCH.plus(Duration.ofSeconds(10)));
|
||||
final KeysManager keysManager = mock(KeysManager.class);
|
||||
|
||||
final UUID currentPage = UUID.randomUUID();
|
||||
final UUID freshOrphanedPage = UUID.randomUUID();
|
||||
final UUID staleOrphanedPage = UUID.randomUUID();
|
||||
|
||||
when(keysManager.listStoredKEMPreKeyPages(anyInt())).thenReturn(Flux.fromIterable(List.of(
|
||||
new DeviceKEMPreKeyPages(UUID.randomUUID(), (byte) 1, Optional.of(currentPage), Map.of(
|
||||
currentPage, Instant.EPOCH,
|
||||
staleOrphanedPage, Instant.EPOCH.plus(Duration.ofSeconds(4)),
|
||||
freshOrphanedPage, Instant.EPOCH.plus(Duration.ofSeconds(5)))))));
|
||||
|
||||
when(keysManager.pruneDeadPage(any(), anyByte(), any()))
|
||||
.thenReturn(CompletableFuture.completedFuture(null));
|
||||
|
||||
runCommand(clock, Duration.ofSeconds(5), dryRun, keysManager);
|
||||
verify(keysManager, times(dryRun ? 0 : 1))
|
||||
.pruneDeadPage(any(), eq((byte) 1), eq(staleOrphanedPage));
|
||||
verify(keysManager, times(1)).listStoredKEMPreKeyPages(anyInt());
|
||||
verifyNoMoreInteractions(keysManager);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void noCurrentPage() throws Exception {
|
||||
final TestClock clock = TestClock.pinned(Instant.EPOCH.plus(Duration.ofSeconds(10)));
|
||||
final KeysManager keysManager = mock(KeysManager.class);
|
||||
|
||||
final UUID freshOrphanedPage = UUID.randomUUID();
|
||||
final UUID staleOrphanedPage = UUID.randomUUID();
|
||||
|
||||
when(keysManager.listStoredKEMPreKeyPages(anyInt())).thenReturn(Flux.fromIterable(List.of(
|
||||
new DeviceKEMPreKeyPages(UUID.randomUUID(), (byte) 1, Optional.empty(), Map.of(
|
||||
staleOrphanedPage, Instant.EPOCH.plus(Duration.ofSeconds(4)),
|
||||
freshOrphanedPage, Instant.EPOCH.plus(Duration.ofSeconds(5)))))));
|
||||
|
||||
when(keysManager.pruneDeadPage(any(), anyByte(), any()))
|
||||
.thenReturn(CompletableFuture.completedFuture(null));
|
||||
|
||||
runCommand(clock, Duration.ofSeconds(5), false, keysManager);
|
||||
verify(keysManager, times(1))
|
||||
.pruneDeadPage(any(), eq((byte) 1), eq(staleOrphanedPage));
|
||||
verify(keysManager, times(1)).listStoredKEMPreKeyPages(anyInt());
|
||||
verifyNoMoreInteractions(keysManager);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void noPages() throws Exception {
|
||||
final TestClock clock = TestClock.pinned(Instant.EPOCH);
|
||||
final KeysManager keysManager = mock(KeysManager.class);
|
||||
when(keysManager.listStoredKEMPreKeyPages(anyInt())).thenReturn(Flux.empty());
|
||||
runCommand(clock, Duration.ofSeconds(5), false, keysManager);
|
||||
verify(keysManager).listStoredKEMPreKeyPages(anyInt());
|
||||
verifyNoMoreInteractions(keysManager);
|
||||
}
|
||||
|
||||
private enum PageStatus {NO_CURRENT, MATCH_CURRENT, MISMATCH_CURRENT}
|
||||
|
||||
@CartesianTest
|
||||
void shouldDeletePage(
|
||||
@CartesianTest.Enum final PageStatus pageStatus,
|
||||
@CartesianTest.Values(booleans = {false, true}) final boolean isOld) {
|
||||
final Optional<UUID> currentPage = pageStatus == PageStatus.NO_CURRENT
|
||||
? Optional.empty()
|
||||
: Optional.of(UUID.randomUUID());
|
||||
final UUID page = switch (pageStatus) {
|
||||
case MATCH_CURRENT -> currentPage.orElseThrow();
|
||||
case NO_CURRENT, MISMATCH_CURRENT -> UUID.randomUUID();
|
||||
};
|
||||
|
||||
final Instant threshold = Instant.EPOCH.plus(Duration.ofSeconds(10));
|
||||
final Instant lastModified = isOld ? threshold.minus(Duration.ofSeconds(1)) : threshold;
|
||||
|
||||
final boolean shouldDelete = pageStatus != PageStatus.MATCH_CURRENT && isOld;
|
||||
Assertions.assertThat(RemoveOrphanedPreKeyPagesCommand.shouldDeletePage(currentPage, page, threshold, lastModified))
|
||||
.isEqualTo(shouldDelete);
|
||||
}
|
||||
|
||||
|
||||
private void runCommand(final Clock clock, final Duration minimumOrphanAge, final boolean dryRun,
|
||||
final KeysManager keysManager) throws Exception {
|
||||
final CommandDependencies commandDependencies = mock(CommandDependencies.class);
|
||||
when(commandDependencies.keysManager()).thenReturn(keysManager);
|
||||
|
||||
final Namespace namespace = mock(Namespace.class);
|
||||
when(namespace.getBoolean(RemoveOrphanedPreKeyPagesCommand.DRY_RUN_ARGUMENT)).thenReturn(dryRun);
|
||||
when(namespace.getInt(RemoveOrphanedPreKeyPagesCommand.CONCURRENCY_ARGUMENT)).thenReturn(2);
|
||||
when(namespace.getString(RemoveOrphanedPreKeyPagesCommand.MINIMUM_ORPHAN_AGE_ARGUMENT))
|
||||
.thenReturn(minimumOrphanAge.toString());
|
||||
|
||||
final RemoveOrphanedPreKeyPagesCommand command = new RemoveOrphanedPreKeyPagesCommand(clock);
|
||||
command.run(mock(Environment.class), namespace, mock(WhisperServerConfiguration.class), commandDependencies);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user