Signal/Signal-Server
1
0
Fork 0
mirror of https://github.com/signalapp/Signal-Server synced 2026-04-20 06:08:03 +01:00
Code Issues Packages Projects Releases Wiki Activity

Store and retrieve one-time pre-keys by UUID

Browse Source
This commit is contained in:
Jon Chambers
2021-11-03 18:21:40 -04:00
committed by Jon Chambers
parent 5e1334e8de
commit 3a4c5a2bfb
4 changed files with 87 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
View File

@@ -11,6 +11,7 @@ import io.dropwizard.auth.Auth;
import io.micrometer.core.instrument.Metrics;
import io.micrometer.core.instrument.Tags;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
@@ -77,7 +78,7 @@ public class KeysController {
@GET
@Produces(MediaType.APPLICATION_JSON)
public PreKeyCount getStatus(@Auth AuthenticatedAccount auth) {
int count = keys.getCount(auth.getAccount(), auth.getAuthenticatedDevice().getId());
int count = keys.getCount(auth.getAccount().getUuid(), auth.getAuthenticatedDevice().getId());
if (count > 0) {
count = count - 1;
@@ -109,7 +110,7 @@ public class KeysController {
});
}
keys.store(account, device.getId(), preKeys.getPreKeys());
keys.store(account.getUuid(), device.getId(), preKeys.getPreKeys());
}
@Timed
@@ -211,18 +212,26 @@ public class KeysController {
}
private Map<Long, PreKey> getLocalKeys(Account destination, String deviceIdSelector) {
try {
if (deviceIdSelector.equals("*")) {
return keys.take(destination);
final Map<Long, PreKey> preKeys;
if (deviceIdSelector.equals("*")) {
preKeys = new HashMap<>();
for (final Device device : destination.getDevices()) {
keys.take(destination.getUuid(), device.getId()).ifPresent(preKey -> preKeys.put(device.getId(), preKey));
}
} else {
try {
long deviceId = Long.parseLong(deviceIdSelector);
long deviceId = Long.parseLong(deviceIdSelector);
return keys.take(destination, deviceId)
.map(preKey -> Map.of(deviceId, preKey))
.orElse(Collections.emptyMap());
} catch (NumberFormatException e) {
throw new WebApplicationException(Response.status(422).build());
preKeys = keys.take(destination.getUuid(), deviceId)
.map(preKey -> Map.of(deviceId, preKey))
.orElse(Collections.emptyMap());
} catch (NumberFormatException e) {
throw new WebApplicationException(Response.status(422).build());
}
}
return preKeys;
}
}

28
service/src/main/java/org/whispersystems/textsecuregcm/storage/Keys.java
View File

@@ -13,7 +13,6 @@ import io.micrometer.core.instrument.Metrics;
import io.micrometer.core.instrument.Timer;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
@@ -42,7 +41,6 @@ public class Keys extends AbstractDynamoDbStore {
private static final Timer STORE_KEYS_TIMER = Metrics.timer(name(Keys.class, "storeKeys"));
private static final Timer TAKE_KEY_FOR_DEVICE_TIMER = Metrics.timer(name(Keys.class, "takeKeyForDevice"));
private static final Timer TAKE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(Keys.class, "takeKeyForAccount"));
private static final Timer GET_KEY_COUNT_TIMER = Metrics.timer(name(Keys.class, "getKeyCount"));
private static final Timer DELETE_KEYS_FOR_DEVICE_TIMER = Metrics.timer(name(Keys.class, "deleteKeysForDevice"));
private static final Timer DELETE_KEYS_FOR_ACCOUNT_TIMER = Metrics.timer(name(Keys.class, "deleteKeysForAccount"));
@@ -54,16 +52,16 @@ public class Keys extends AbstractDynamoDbStore {
this.tableName = tableName;
}
public void store(final Account account, final long deviceId, final List<PreKey> keys) {
public void store(final UUID identifier, final long deviceId, final List<PreKey> keys) {
STORE_KEYS_TIMER.record(() -> {
delete(account.getUuid(), deviceId);
delete(identifier, deviceId);
writeInBatches(keys, batch -> {
List<WriteRequest> items = new ArrayList<>();
for (final PreKey preKey : batch) {
items.add(WriteRequest.builder()
.putRequest(PutRequest.builder()
.item(getItemFromPreKey(account.getUuid(), deviceId, preKey))
.item(getItemFromPreKey(identifier, deviceId, preKey))
.build())
.build());
}
@@ -72,9 +70,9 @@ public class Keys extends AbstractDynamoDbStore {
});
}
public Optional<PreKey> take(final Account account, final long deviceId) {
public Optional<PreKey> take(final UUID identifier, final long deviceId) {
return TAKE_KEY_FOR_DEVICE_TIMER.record(() -> {
final AttributeValue partitionKey = getPartitionKey(account.getUuid());
final AttributeValue partitionKey = getPartitionKey(identifier);
QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
@@ -113,26 +111,14 @@ public class Keys extends AbstractDynamoDbStore {
});
}
public Map<Long, PreKey> take(final Account account) {
return TAKE_KEYS_FOR_ACCOUNT_TIMER.record(() -> {
final Map<Long, PreKey> preKeysByDeviceId = new HashMap<>();
for (final Device device : account.getDevices()) {
take(account, device.getId()).ifPresent(preKey -> preKeysByDeviceId.put(device.getId(), preKey));
}
return preKeysByDeviceId;
});
}
public int getCount(final Account account, final long deviceId) {
public int getCount(final UUID identifier, final long deviceId) {
return GET_KEY_COUNT_TIMER.record(() -> {
QueryRequest queryRequest = QueryRequest.builder()
.tableName(tableName)
.keyConditionExpression("#uuid = :uuid AND begins_with (#sort, :sortprefix)")
.expressionAttributeNames(Map.of("#uuid", KEY_ACCOUNT_UUID, "#sort", KEY_DEVICE_ID_KEY_ID))
.expressionAttributeValues(Map.of(
":uuid", getPartitionKey(account.getUuid()),
":uuid", getPartitionKey(identifier),
":sortprefix", getSortKeyPrefix(deviceId)))
.select(Select.COUNT)
.consistentRead(false)