Use destination service ID from the envelope when removing views from shared MRM data

This commit is contained in:
Chris Eager
2024-09-13 17:38:18 -05:00
committed by Chris Eager
parent 11691c3122
commit 374fe087bc
8 changed files with 68 additions and 52 deletions

View File

@@ -73,8 +73,7 @@ class MessageSenderTest {
MessageProtos.Envelope.class);
verify(messagesManager).insert(any(), anyByte(), envelopeArgumentCaptor.capture());
verify(messagesManager, never()).removeRecipientViewFromMrmData(any(), anyByte(),
any(MessageProtos.Envelope.class));
verify(messagesManager, never()).removeRecipientViewFromMrmData(anyByte(), any(MessageProtos.Envelope.class));
assertTrue(envelopeArgumentCaptor.getValue().getEphemeral());
@@ -96,7 +95,7 @@ class MessageSenderTest {
}
verify(messagesManager, never()).insert(any(), anyByte(), any());
verify(messagesManager).removeRecipientViewFromMrmData(any(), anyByte(), any(MessageProtos.Envelope.class));
verify(messagesManager).removeRecipientViewFromMrmData(anyByte(), any(MessageProtos.Envelope.class));
verifyNoInteractions(pushNotificationManager);
}

View File

@@ -8,6 +8,7 @@ package org.whispersystems.textsecuregcm.storage;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import io.lettuce.core.RedisCommandExecutionException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
@@ -15,13 +16,13 @@ import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import io.lettuce.core.RedisCommandExecutionException;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
import org.whispersystems.textsecuregcm.util.Pair;
@@ -32,7 +33,7 @@ class MessagesCacheInsertSharedMultiRecipientPayloadAndViewsScriptTest {
@ParameterizedTest
@MethodSource
void testInsert(final int count, final Map<AciServiceIdentifier, List<Byte>> destinations) throws Exception {
void testInsert(final int count, final Map<ServiceIdentifier, List<Byte>> destinations) throws Exception {
final MessagesCacheInsertSharedMultiRecipientPayloadAndViewsScript insertMrmScript = new MessagesCacheInsertSharedMultiRecipientPayloadAndViewsScript(
REDIS_CLUSTER_EXTENSION.getRedisCluster());
@@ -49,7 +50,7 @@ class MessagesCacheInsertSharedMultiRecipientPayloadAndViewsScriptTest {
}
public static List<Arguments> testInsert() {
final Map<AciServiceIdentifier, List<Byte>> singleAccount = Map.of(
final Map<ServiceIdentifier, List<Byte>> singleAccount = Map.of(
new AciServiceIdentifier(UUID.randomUUID()), List.of((byte) 1, (byte) 2));
final List<Arguments> testCases = new ArrayList<>();
@@ -58,7 +59,7 @@ class MessagesCacheInsertSharedMultiRecipientPayloadAndViewsScriptTest {
for (int j = 1000; j <= 30000; j += 1000) {
final Map<Integer, List<Byte>> deviceLists = new HashMap<>();
final Map<AciServiceIdentifier, List<Byte>> manyAccounts = IntStream.range(0, j)
final Map<ServiceIdentifier, List<Byte>> manyAccounts = IntStream.range(0, j)
.mapToObj(i -> {
final int deviceCount = 1 + i % 5;
final List<Byte> devices = deviceLists.computeIfAbsent(deviceCount, count -> IntStream.rangeClosed(1, count)

View File

@@ -22,6 +22,7 @@ import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.MethodSource;
import org.junit.jupiter.params.provider.ValueSource;
import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
import org.whispersystems.textsecuregcm.util.Pair;
import reactor.core.publisher.Flux;
@@ -34,7 +35,7 @@ class MessagesCacheRemoveRecipientViewFromMrmDataScriptTest {
@ParameterizedTest
@MethodSource
void testUpdateSingleKey(final Map<AciServiceIdentifier, List<Byte>> destinations) throws Exception {
void testUpdateSingleKey(final Map<ServiceIdentifier, List<Byte>> destinations) throws Exception {
final MessagesCacheInsertSharedMultiRecipientPayloadAndViewsScript insertMrmScript = new MessagesCacheInsertSharedMultiRecipientPayloadAndViewsScript(
REDIS_CLUSTER_EXTENSION.getRedisCluster());
@@ -48,8 +49,8 @@ class MessagesCacheRemoveRecipientViewFromMrmDataScriptTest {
final long keysRemoved = Objects.requireNonNull(Flux.fromIterable(destinations.entrySet())
.flatMap(e -> Flux.fromStream(e.getValue().stream().map(deviceId -> Tuples.of(e.getKey(), deviceId))))
.flatMap(aciServiceIdentifierByteTuple -> removeRecipientViewFromMrmDataScript.execute(List.of(sharedMrmKey),
aciServiceIdentifierByteTuple.getT1(), aciServiceIdentifierByteTuple.getT2()))
.flatMap(serviceIdentifierByteTuple -> removeRecipientViewFromMrmDataScript.execute(List.of(sharedMrmKey),
serviceIdentifierByteTuple.getT1(), serviceIdentifierByteTuple.getT2()))
.reduce(Long::sum)
.block(Duration.ofSeconds(35)));
@@ -60,18 +61,18 @@ class MessagesCacheRemoveRecipientViewFromMrmDataScriptTest {
assertEquals(0, keyExists);
}
public static List<Map<AciServiceIdentifier, List<Byte>>> testUpdateSingleKey() {
final Map<AciServiceIdentifier, List<Byte>> singleAccount = Map.of(
public static List<Map<ServiceIdentifier, List<Byte>>> testUpdateSingleKey() {
final Map<ServiceIdentifier, List<Byte>> singleAccount = Map.of(
new AciServiceIdentifier(UUID.randomUUID()), List.of((byte) 1, (byte) 2));
final List<Map<AciServiceIdentifier, List<Byte>>> testCases = new ArrayList<>();
final List<Map<ServiceIdentifier, List<Byte>>> testCases = new ArrayList<>();
testCases.add(singleAccount);
// Generate a more, from smallish to very large
for (int j = 1000; j <= 81000; j *= 3) {
final Map<Integer, List<Byte>> deviceLists = new HashMap<>();
final Map<AciServiceIdentifier, List<Byte>> manyAccounts = IntStream.range(0, j)
final Map<ServiceIdentifier, List<Byte>> manyAccounts = IntStream.range(0, j)
.mapToObj(i -> {
final int deviceCount = 1 + i % 5;
final List<Byte> devices = deviceLists.computeIfAbsent(deviceCount, count -> IntStream.rangeClosed(1, count)
@@ -93,7 +94,7 @@ class MessagesCacheRemoveRecipientViewFromMrmDataScriptTest {
void testUpdateManyKeys(int keyCount) throws Exception {
final List<byte[]> sharedMrmKeys = new ArrayList<>(keyCount);
final AciServiceIdentifier aciServiceIdentifier = new AciServiceIdentifier(UUID.randomUUID());
final ServiceIdentifier serviceIdentifier = new AciServiceIdentifier(UUID.randomUUID());
final byte deviceId = 1;
for (int i = 0; i < keyCount; i++) {
@@ -103,7 +104,7 @@ class MessagesCacheRemoveRecipientViewFromMrmDataScriptTest {
final byte[] sharedMrmKey = MessagesCache.getSharedMrmKey(UUID.randomUUID());
insertMrmScript.execute(sharedMrmKey,
MessagesCacheTest.generateRandomMrmMessage(aciServiceIdentifier, deviceId));
MessagesCacheTest.generateRandomMrmMessage(serviceIdentifier, deviceId));
sharedMrmKeys.add(sharedMrmKey);
}
@@ -114,7 +115,7 @@ class MessagesCacheRemoveRecipientViewFromMrmDataScriptTest {
final long keysRemoved = Objects.requireNonNull(Flux.fromIterable(sharedMrmKeys)
.collectMultimap(SlotHash::getSlot)
.flatMapMany(slotsAndKeys -> Flux.fromIterable(slotsAndKeys.values()))
.flatMap(keys -> removeRecipientViewFromMrmDataScript.execute(keys, aciServiceIdentifier, deviceId))
.flatMap(keys -> removeRecipientViewFromMrmDataScript.execute(keys, serviceIdentifier, deviceId))
.reduce(Long::sum)
.block(Duration.ofSeconds(5)));

View File

@@ -71,6 +71,7 @@ import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfigurati
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicMessagesConfiguration;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
import org.whispersystems.textsecuregcm.tests.util.RedisClusterHelper;
@@ -565,11 +566,10 @@ class MessagesCacheTest {
@ParameterizedTest
@ValueSource(booleans = {true, false})
void testMultiRecipientMessage(final boolean sharedMrmKeyPresent) throws Exception {
final UUID destinationUuid = UUID.randomUUID();
final ServiceIdentifier destinationServiceId = new AciServiceIdentifier(UUID.randomUUID());
final byte deviceId = 1;
final SealedSenderMultiRecipientMessage mrm = generateRandomMrmMessage(
new AciServiceIdentifier(destinationUuid), deviceId);
final SealedSenderMultiRecipientMessage mrm = generateRandomMrmMessage(destinationServiceId, deviceId);
final byte[] sharedMrmDataKey;
if (sharedMrmKeyPresent) {
@@ -579,35 +579,35 @@ class MessagesCacheTest {
}
final UUID guid = UUID.randomUUID();
final MessageProtos.Envelope message = generateRandomMessage(guid, true)
final MessageProtos.Envelope message = generateRandomMessage(guid, destinationServiceId, true)
.toBuilder()
// clear some things added by the helper
.clearServerGuid()
// mrm views phase 1: messages have content
.setContent(
ByteString.copyFrom(mrm.messageForRecipient(mrm.getRecipients().get(new ServiceId.Aci(destinationUuid)))))
ByteString.copyFrom(mrm.messageForRecipient(mrm.getRecipients().get(destinationServiceId.toLibsignal()))))
.setSharedMrmKey(ByteString.copyFrom(sharedMrmDataKey))
.build();
messagesCache.insert(guid, destinationUuid, deviceId, message);
messagesCache.insert(guid, destinationServiceId.uuid(), deviceId, message);
assertEquals(sharedMrmKeyPresent ? 1 : 0, (long) REDIS_CLUSTER_EXTENSION.getRedisCluster()
.withBinaryCluster(conn -> conn.sync().exists(sharedMrmDataKey)));
final List<MessageProtos.Envelope> messages = get(destinationUuid, deviceId, 1);
final List<MessageProtos.Envelope> messages = get(destinationServiceId.uuid(), deviceId, 1);
assertEquals(1, messages.size());
assertEquals(guid, UUID.fromString(messages.getFirst().getServerGuid()));
assertFalse(messages.getFirst().hasSharedMrmKey());
final SealedSenderMultiRecipientMessage.Recipient recipient = mrm.getRecipients()
.get(new ServiceId.Aci(destinationUuid));
.get(destinationServiceId.toLibsignal());
assertArrayEquals(mrm.messageForRecipient(recipient), messages.getFirst().getContent().toByteArray());
final Optional<RemovedMessage> removedMessage = messagesCache.remove(destinationUuid, deviceId, guid)
final Optional<RemovedMessage> removedMessage = messagesCache.remove(destinationServiceId.uuid(), deviceId, guid)
.join();
assertTrue(removedMessage.isPresent());
assertEquals(guid, UUID.fromString(removedMessage.get().serverGuid().toString()));
assertTrue(get(destinationUuid, deviceId, 1).isEmpty());
assertTrue(get(destinationServiceId.uuid(), deviceId, 1).isEmpty());
// updating the shared MRM data is purely async, so we just wait for it
assertTimeoutPreemptively(Duration.ofSeconds(1), () -> {
@@ -874,10 +874,17 @@ class MessagesCacheTest {
}
private MessageProtos.Envelope generateRandomMessage(final UUID messageGuid, final boolean sealedSender) {
return generateRandomMessage(messageGuid, sealedSender, serialTimestamp++);
return generateRandomMessage(messageGuid, new AciServiceIdentifier(UUID.randomUUID()), sealedSender,
serialTimestamp++);
}
private MessageProtos.Envelope generateRandomMessage(final UUID messageGuid, final boolean sealedSender,
private MessageProtos.Envelope generateRandomMessage(final UUID messageGuid,
final ServiceIdentifier destinationServiceId, final boolean sealedSender) {
return generateRandomMessage(messageGuid, destinationServiceId, sealedSender, serialTimestamp++);
}
private MessageProtos.Envelope generateRandomMessage(final UUID messageGuid,
final ServiceIdentifier destinationServiceId, final boolean sealedSender,
final long timestamp) {
final MessageProtos.Envelope.Builder envelopeBuilder = MessageProtos.Envelope.newBuilder()
.setClientTimestamp(timestamp)
@@ -885,7 +892,7 @@ class MessagesCacheTest {
.setContent(ByteString.copyFromUtf8(RandomStringUtils.randomAlphanumeric(256)))
.setType(MessageProtos.Envelope.Type.CIPHERTEXT)
.setServerGuid(messageGuid.toString())
.setDestinationServiceId(UUID.randomUUID().toString());
.setDestinationServiceId(destinationServiceId.toServiceIdentifierString());
if (!sealedSender) {
envelopeBuilder.setSourceDevice(random.nextInt(Device.MAXIMUM_DEVICE_ID) + 1)
@@ -896,8 +903,7 @@ class MessagesCacheTest {
}
static SealedSenderMultiRecipientMessage generateRandomMrmMessage(
Map<AciServiceIdentifier, List<Byte>> destinations) {
Map<ServiceIdentifier, List<Byte>> destinations) {
try {
final ByteBuffer prefix = ByteBuffer.allocate(7);
@@ -907,10 +913,10 @@ class MessagesCacheTest {
List<ByteBuffer> recipients = new ArrayList<>(destinations.size());
for (Map.Entry<AciServiceIdentifier, List<Byte>> aciServiceIdentifierAndDeviceIds : destinations.entrySet()) {
for (Map.Entry<ServiceIdentifier, List<Byte>> serviceIdentifierAndDeviceIds : destinations.entrySet()) {
final AciServiceIdentifier destination = aciServiceIdentifierAndDeviceIds.getKey();
final List<Byte> deviceIds = aciServiceIdentifierAndDeviceIds.getValue();
final ServiceIdentifier destination = serviceIdentifierAndDeviceIds.getKey();
final List<Byte> deviceIds = serviceIdentifierAndDeviceIds.getValue();
assert deviceIds.size() < 255;
@@ -946,10 +952,10 @@ class MessagesCacheTest {
}
}
static SealedSenderMultiRecipientMessage generateRandomMrmMessage(AciServiceIdentifier destination,
static SealedSenderMultiRecipientMessage generateRandomMrmMessage(ServiceIdentifier destination,
byte... deviceIds) {
final Map<AciServiceIdentifier, List<Byte>> destinations = new HashMap<>();
final Map<ServiceIdentifier, List<Byte>> destinations = new HashMap<>();
destinations.put(destination, Arrays.asList(ArrayUtils.toObject(deviceIds)));
return generateRandomMrmMessage(destinations);
}