Add a temporary method for storing signed EC pre-keys if and only if another key has not already been stored

This commit is contained in:
Jon Chambers
2023-06-20 11:22:09 -04:00
committed by Jon Chambers
parent 97a7469432
commit a3e82dfae8
2 changed files with 51 additions and 0 deletions

View File

@@ -7,17 +7,27 @@ package org.whispersystems.textsecuregcm.storage;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import org.signal.libsignal.protocol.InvalidKeyException;
import org.signal.libsignal.protocol.ecc.ECPublicKey;
import org.whispersystems.textsecuregcm.entities.ECSignedPreKey;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import org.whispersystems.textsecuregcm.util.ExceptionUtils;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
public class RepeatedUseECSignedPreKeyStore extends RepeatedUseSignedPreKeyStore<ECSignedPreKey> {
private final DynamoDbAsyncClient dynamoDbAsyncClient;
private final String tableName;
public RepeatedUseECSignedPreKeyStore(final DynamoDbAsyncClient dynamoDbAsyncClient, final String tableName) {
super(dynamoDbAsyncClient, tableName);
this.dynamoDbAsyncClient = dynamoDbAsyncClient;
this.tableName = tableName;
}
@Override
@@ -43,4 +53,21 @@ public class RepeatedUseECSignedPreKeyStore extends RepeatedUseSignedPreKeyStore
throw new IllegalArgumentException(e);
}
}
public CompletableFuture<Boolean> storeIfAbsent(final UUID identifier, final long deviceId, final ECSignedPreKey signedPreKey) {
return dynamoDbAsyncClient.putItem(PutItemRequest.builder()
.tableName(tableName)
.item(getItemFromPreKey(identifier, deviceId, signedPreKey))
.conditionExpression("attribute_not_exists(#public_key)")
.expressionAttributeNames(Map.of("#public_key", ATTR_PUBLIC_KEY))
.build())
.thenApply(ignored -> true)
.exceptionally(throwable -> {
if (ExceptionUtils.unwrap(throwable) instanceof ConditionalCheckFailedException) {
return false;
}
throw ExceptionUtils.wrap(throwable);
});
}
}