Skip to content

Commit 057d1f0

Browse files
jon-signaljon-signal
committed
Remove bulk "set repeated-use signed pre-keys" methods because they were only ever used for single devices
1 parent 25c3f55 commit 057d1f0

File tree

9 files changed

+47
-109
lines changed

9 files changed

+47
-109
lines changed

‎service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,7 @@ public CompletableFuture<Response> setKeys(@Auth final DisabledPermittedAuthenti
164164

165165
if (setKeysRequest.pqLastResortPreKey() != null) {
166166
storeFutures.add(
167-
keys.storePqLastResort(identifier, Map.of(device.getId(), setKeysRequest.pqLastResortPreKey())));
167+
keys.storePqLastResort(identifier, device.getId(), setKeysRequest.pqLastResortPreKey()));
168168
}
169169

170170
return CompletableFuture.allOf(storeFutures.toArray(EMPTY_FUTURE_ARRAY));

‎service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcService.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -200,7 +200,7 @@ public Mono<SetPreKeyResponse> setEcSignedPreKey(final SetEcSignedPreKeyRequest
200200
final UUID identifier = account.getIdentifier(identityType);
201201

202202
return Flux.merge(
203-
Mono.fromFuture(() -> keysManager.storeEcSignedPreKeys(identifier, Map.of(authenticatedDevice.deviceId(), signedPreKey))),
203+
Mono.fromFuture(() -> keysManager.storeEcSignedPreKeys(identifier, authenticatedDevice.deviceId(), signedPreKey)),
204204
Mono.fromFuture(() -> accountsManager.updateDeviceAsync(account, authenticatedDevice.deviceId(), deviceUpdater)))
205205
.then();
206206
}));
@@ -217,7 +217,7 @@ public Mono<SetPreKeyResponse> setKemLastResortPreKey(final SetKemLastResortPreK
217217
final UUID identifier =
218218
account.getIdentifier(IdentityTypeUtil.fromGrpcIdentityType(request.getIdentityType()));
219219

220-
return Mono.fromFuture(() -> keysManager.storePqLastResort(identifier, Map.of(authenticatedDevice.deviceId(), lastResortKey)));
220+
return Mono.fromFuture(() -> keysManager.storePqLastResort(identifier, authenticatedDevice.deviceId(), lastResortKey));
221221
}));
222222
}
223223

‎service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysManager.java

+4-5
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
import com.google.common.annotations.VisibleForTesting;
99
import java.util.ArrayList;
1010
import java.util.List;
11-
import java.util.Map;
1211
import java.util.Optional;
1312
import java.util.UUID;
1413
import java.util.concurrent.CompletableFuture;
@@ -100,9 +99,9 @@ public List<TransactWriteItem> buildWriteItemsForRemovedDevice(final UUID accoun
10099
return writeItems;
101100
}
102101

103-
public CompletableFuture<Void> storeEcSignedPreKeys(final UUID identifier, final Map<Byte, ECSignedPreKey> keys) {
102+
public CompletableFuture<Void> storeEcSignedPreKeys(final UUID identifier, final byte deviceId, final ECSignedPreKey ecSignedPreKey) {
104103
if (dynamicConfigurationManager.getConfiguration().getEcPreKeyMigrationConfiguration().storeEcSignedPreKeys()) {
105-
return ecSignedPreKeys.store(identifier, keys);
104+
return ecSignedPreKeys.store(identifier, deviceId, ecSignedPreKey);
106105
} else {
107106
return CompletableFuture.completedFuture(null);
108107
}
@@ -113,8 +112,8 @@ public CompletableFuture<Boolean> storeEcSignedPreKeyIfAbsent(final UUID identif
113112
return ecSignedPreKeys.storeIfAbsent(identifier, deviceId, signedPreKey);
114113
}
115114

116-
public CompletableFuture<Void> storePqLastResort(final UUID identifier, final Map<Byte, KEMSignedPreKey> keys) {
117-
return pqLastResortKeys.store(identifier, keys);
115+
public CompletableFuture<Void> storePqLastResort(final UUID identifier, final byte deviceId, final KEMSignedPreKey lastResortKey) {
116+
return pqLastResortKeys.store(identifier, deviceId, lastResortKey);
118117
}
119118

120119
public CompletableFuture<Void> storeEcOneTimePreKeys(final UUID identifier, final byte deviceId,

‎service/src/main/java/org/whispersystems/textsecuregcm/storage/RepeatedUseSignedPreKeyStore.java

-37
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
2424
import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
2525
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
26-
import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
2726

2827
/**
2928
* A repeated-use signed pre-key store manages storage for pre-keys that may be used more than once. Generally, these
@@ -45,7 +44,6 @@ public abstract class RepeatedUseSignedPreKeyStore<K extends SignedPreKey<?>> {
4544
static final String ATTR_SIGNATURE = "S";
4645

4746
private final Timer storeSingleKeyTimer = Metrics.timer(MetricsUtil.name(getClass(), "storeSingleKey"));
48-
private final Timer storeKeyBatchTimer = Metrics.timer(MetricsUtil.name(getClass(), "storeKeyBatch"));
4947

5048
private final String findKeyTimerName = MetricsUtil.name(getClass(), "findKey");
5149

@@ -74,41 +72,6 @@ public CompletableFuture<Void> store(final UUID identifier, final byte deviceId,
7472
.thenRun(() -> sample.stop(storeSingleKeyTimer));
7573
}
7674

77-
/**
78-
* Stores repeated-use pre-keys for a collection of devices associated with a single account/identity, displacing any
79-
* previously-stored repeated-use pre-keys for the targeted devices. Note that this method is transactional; either
80-
* all keys will be stored or none will.
81-
*
82-
* @param identifier the identifier for the account/identity with which the target devices are associated
83-
* @param signedPreKeysByDeviceId a map of device identifiers to pre-keys
84-
*
85-
* @return a future that completes once all keys have been stored
86-
*/
87-
public CompletableFuture<Void> store(final UUID identifier, final Map<Byte, K> signedPreKeysByDeviceId) {
88-
if (signedPreKeysByDeviceId.isEmpty()) {
89-
return CompletableFuture.completedFuture(null);
90-
}
91-
92-
final Timer.Sample sample = Timer.start();
93-
94-
return dynamoDbAsyncClient.transactWriteItems(TransactWriteItemsRequest.builder()
95-
.transactItems(signedPreKeysByDeviceId.entrySet().stream()
96-
.map(entry -> {
97-
final byte deviceId = entry.getKey();
98-
final K signedPreKey = entry.getValue();
99-
100-
return TransactWriteItem.builder()
101-
.put(Put.builder()
102-
.tableName(tableName)
103-
.item(getItemFromPreKey(identifier, deviceId, signedPreKey))
104-
.build())
105-
.build();
106-
})
107-
.toList())
108-
.build())
109-
.thenRun(() -> sample.stop(storeKeyBatchTimer));
110-
}
111-
11275
TransactWriteItem buildTransactWriteItemForInsertion(final UUID identifier, final byte deviceId, final K preKey) {
11376
return TransactWriteItem.builder()
11477
.put(Put.builder()

‎service/src/test/java/org/whispersystems/textsecuregcm/controllers/KeysControllerTest.java

+4-4
Original file line numberDiff line numberDiff line change
@@ -254,11 +254,11 @@ void setup() {
254254
when(KEYS.storeKemOneTimePreKeys(any(), anyByte(), any()))
255255
.thenReturn(CompletableFutureTestUtil.almostCompletedFuture(null));
256256

257-
when(KEYS.storePqLastResort(any(), any()))
257+
when(KEYS.storePqLastResort(any(), anyByte(), any()))
258258
.thenReturn(CompletableFutureTestUtil.almostCompletedFuture(null));
259259

260260
when(KEYS.getEcSignedPreKey(any(), anyByte())).thenReturn(CompletableFuture.completedFuture(Optional.empty()));
261-
when(KEYS.storeEcSignedPreKeys(any(), any())).thenReturn(CompletableFutureTestUtil.almostCompletedFuture(null));
261+
when(KEYS.storeEcSignedPreKeys(any(), anyByte(), any())).thenReturn(CompletableFutureTestUtil.almostCompletedFuture(null));
262262

263263
when(KEYS.takeEC(EXISTS_UUID, sampleDeviceId)).thenReturn(
264264
CompletableFuture.completedFuture(Optional.of(SAMPLE_KEY)));
@@ -827,7 +827,7 @@ void putKeysPqTestV2() {
827827
ArgumentCaptor<List<KEMSignedPreKey>> pqCaptor = ArgumentCaptor.forClass(List.class);
828828
verify(KEYS).storeEcOneTimePreKeys(eq(AuthHelper.VALID_UUID), eq(SAMPLE_DEVICE_ID), ecCaptor.capture());
829829
verify(KEYS).storeKemOneTimePreKeys(eq(AuthHelper.VALID_UUID), eq(SAMPLE_DEVICE_ID), pqCaptor.capture());
830-
verify(KEYS).storePqLastResort(AuthHelper.VALID_UUID, Map.of(SAMPLE_DEVICE_ID, pqLastResortPreKey));
830+
verify(KEYS).storePqLastResort(AuthHelper.VALID_UUID, SAMPLE_DEVICE_ID, pqLastResortPreKey);
831831

832832
assertThat(ecCaptor.getValue()).containsExactly(preKey);
833833
assertThat(pqCaptor.getValue()).containsExactly(pqPreKey);
@@ -966,7 +966,7 @@ void putKeysByPhoneNumberIdentifierPqTestV2() {
966966
ArgumentCaptor<List<KEMSignedPreKey>> pqCaptor = ArgumentCaptor.forClass(List.class);
967967
verify(KEYS).storeEcOneTimePreKeys(eq(AuthHelper.VALID_PNI), eq(SAMPLE_DEVICE_ID), ecCaptor.capture());
968968
verify(KEYS).storeKemOneTimePreKeys(eq(AuthHelper.VALID_PNI), eq(SAMPLE_DEVICE_ID), pqCaptor.capture());
969-
verify(KEYS).storePqLastResort(AuthHelper.VALID_PNI, Map.of(SAMPLE_DEVICE_ID, pqLastResortPreKey));
969+
verify(KEYS).storePqLastResort(AuthHelper.VALID_PNI, SAMPLE_DEVICE_ID, pqLastResortPreKey);
970970

971971
assertThat(ecCaptor.getValue()).containsExactly(preKey);
972972
assertThat(pqCaptor.getValue()).containsExactly(pqPreKey);

‎service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcServiceTest.java

+5-5
Original file line numberDiff line numberDiff line change
@@ -304,7 +304,7 @@ void setSignedPreKey(final org.signal.chat.common.IdentityType identityType) {
304304
return CompletableFuture.completedFuture(account);
305305
});
306306

307-
when(keysManager.storeEcSignedPreKeys(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
307+
when(keysManager.storeEcSignedPreKeys(any(), anyByte(), any())).thenReturn(CompletableFuture.completedFuture(null));
308308

309309
final ECKeyPair identityKeyPair = switch (IdentityTypeUtil.fromGrpcIdentityType(identityType)) {
310310
case ACI -> ACI_IDENTITY_KEY_PAIR;
@@ -326,12 +326,12 @@ void setSignedPreKey(final org.signal.chat.common.IdentityType identityType) {
326326
switch (identityType) {
327327
case IDENTITY_TYPE_ACI -> {
328328
verify(authenticatedDevice).setSignedPreKey(signedPreKey);
329-
verify(keysManager).storeEcSignedPreKeys(AUTHENTICATED_ACI, Map.of(AUTHENTICATED_DEVICE_ID, signedPreKey));
329+
verify(keysManager).storeEcSignedPreKeys(AUTHENTICATED_ACI, AUTHENTICATED_DEVICE_ID, signedPreKey);
330330
}
331331

332332
case IDENTITY_TYPE_PNI -> {
333333
verify(authenticatedDevice).setPhoneNumberIdentitySignedPreKey(signedPreKey);
334-
verify(keysManager).storeEcSignedPreKeys(AUTHENTICATED_PNI, Map.of(AUTHENTICATED_DEVICE_ID, signedPreKey));
334+
verify(keysManager).storeEcSignedPreKeys(AUTHENTICATED_PNI, AUTHENTICATED_DEVICE_ID, signedPreKey);
335335
}
336336
}
337337
}
@@ -387,7 +387,7 @@ private static Stream<Arguments> setSignedPreKeyWithError() {
387387
@ParameterizedTest
388388
@EnumSource(value = org.signal.chat.common.IdentityType.class, names = {"IDENTITY_TYPE_ACI", "IDENTITY_TYPE_PNI"})
389389
void setLastResortPreKey(final org.signal.chat.common.IdentityType identityType) {
390-
when(keysManager.storePqLastResort(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
390+
when(keysManager.storePqLastResort(any(), anyByte(), any())).thenReturn(CompletableFuture.completedFuture(null));
391391

392392
final ECKeyPair identityKeyPair = switch (IdentityTypeUtil.fromGrpcIdentityType(identityType)) {
393393
case ACI -> ACI_IDENTITY_KEY_PAIR;
@@ -412,7 +412,7 @@ void setLastResortPreKey(final org.signal.chat.common.IdentityType identityType)
412412
case IDENTITY_TYPE_UNSPECIFIED, UNRECOGNIZED -> throw new AssertionError("Bad identity type");
413413
};
414414

415-
verify(keysManager).storePqLastResort(expectedIdentifier, Map.of(AUTHENTICATED_DEVICE_ID, lastResortPreKey));
415+
verify(keysManager).storePqLastResort(expectedIdentifier, AUTHENTICATED_DEVICE_ID, lastResortPreKey);
416416
}
417417

418418
@ParameterizedTest

‎service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java

+6-6
Original file line numberDiff line numberDiff line change
@@ -1284,7 +1284,7 @@ void testChangePhoneNumberWithPqKeysExistingAccount() throws InterruptedExceptio
12841284
final Account existingAccount = AccountsHelper.generateTestAccount(targetNumber, existingAccountUuid, targetPni, new ArrayList<>(), new byte[UnidentifiedAccessUtil.UNIDENTIFIED_ACCESS_KEY_LENGTH]);
12851285
when(accounts.getByE164(targetNumber)).thenReturn(Optional.of(existingAccount));
12861286
when(keysManager.getPqEnabledDevices(uuid)).thenReturn(CompletableFuture.completedFuture(List.of(Device.PRIMARY_ID, deviceId3)));
1287-
when(keysManager.storePqLastResort(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
1287+
when(keysManager.storePqLastResort(any(), anyByte(), any())).thenReturn(CompletableFuture.completedFuture(null));
12881288

12891289
final List<Device> devices = List.of(
12901290
DevicesHelper.createDevice(Device.PRIMARY_ID, 0L, 101),
@@ -1381,7 +1381,7 @@ void testPniUpdate() throws MismatchedDevicesException {
13811381
final IdentityKey pniIdentityKey = new IdentityKey(Curve.generateKeyPair().getPublicKey());
13821382

13831383
when(keysManager.getPqEnabledDevices(any())).thenReturn(CompletableFuture.completedFuture(Collections.emptyList()));
1384-
when(keysManager.storeEcSignedPreKeys(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
1384+
when(keysManager.storeEcSignedPreKeys(any(), anyByte(), any())).thenReturn(CompletableFuture.completedFuture(null));
13851385

13861386
final Account updatedAccount = accountsManager.updatePniKeys(account, pniIdentityKey, newSignedKeys, null, newRegistrationIds);
13871387

@@ -1437,8 +1437,8 @@ void testPniPqUpdate() throws MismatchedDevicesException {
14371437

14381438
when(keysManager.getPqEnabledDevices(oldPni)).thenReturn(
14391439
CompletableFuture.completedFuture(List.of(Device.PRIMARY_ID)));
1440-
when(keysManager.storeEcSignedPreKeys(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
1441-
when(keysManager.storePqLastResort(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
1440+
when(keysManager.storeEcSignedPreKeys(any(), anyByte(), any())).thenReturn(CompletableFuture.completedFuture(null));
1441+
when(keysManager.storePqLastResort(any(), anyByte(), any())).thenReturn(CompletableFuture.completedFuture(null));
14421442

14431443
Map<Byte, ECSignedPreKey> oldSignedPreKeys = account.getDevices().stream()
14441444
.collect(Collectors.toMap(Device::getId, d -> d.getSignedPreKey(IdentityType.ACI)));
@@ -1500,8 +1500,8 @@ void testPniNonPqToPqUpdate() throws MismatchedDevicesException {
15001500
UUID oldPni = account.getPhoneNumberIdentifier();
15011501

15021502
when(keysManager.getPqEnabledDevices(oldPni)).thenReturn(CompletableFuture.completedFuture(List.of()));
1503-
when(keysManager.storeEcSignedPreKeys(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
1504-
when(keysManager.storePqLastResort(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
1503+
when(keysManager.storeEcSignedPreKeys(any(), anyByte(), any())).thenReturn(CompletableFuture.completedFuture(null));
1504+
when(keysManager.storePqLastResort(any(), anyByte(), any())).thenReturn(CompletableFuture.completedFuture(null));
15051505

15061506
Map<Byte, ECSignedPreKey> oldSignedPreKeys = account.getDevices().stream()
15071507
.collect(Collectors.toMap(Device::getId, d -> d.getSignedPreKey(IdentityType.ACI)));

0 commit comments

Comments
 (0)