Add uak as a top level attribute

This commit is contained in:
Ravi Khadiwala 2022-03-01 14:58:56 -06:00 committed by ravi-signal
parent 8b4355b21d
commit faa6ae284a
2 changed files with 39 additions and 14 deletions

View File

@ -60,6 +60,8 @@ public class Accounts extends AbstractDynamoDbStore {
static final String ATTR_CANONICALLY_DISCOVERABLE = "C";
// username; string
static final String ATTR_USERNAME = "N";
// unidentified access key; byte[] or null
static final String ATTR_UAK = "UAK";
private final DynamoDbClient client;
@ -147,6 +149,11 @@ public class Accounts extends AbstractDynamoDbStore {
ATTR_VERSION, AttributeValues.fromInt(account.getVersion()),
ATTR_CANONICALLY_DISCOVERABLE, AttributeValues.fromBool(account.shouldBeVisibleInDirectory())));
// Add the UAK if it's in the account
account.getUnidentifiedAccessKey()
.map(AttributeValues::fromByteArray)
.ifPresent(uak -> item.put(ATTR_UAK, uak));
TransactWriteItem accountPut = TransactWriteItem.builder()
.put(Put.builder()
.conditionExpression("attribute_not_exists(#number) OR #number = :number")
@ -453,21 +460,36 @@ public class Accounts extends AbstractDynamoDbStore {
final UpdateItemRequest updateItemRequest;
try {
// username, e164, and pni cannot be modified through this method
Map<String, String> attrNames = new HashMap<>(Map.of(
"#number", ATTR_ACCOUNT_E164,
"#data", ATTR_ACCOUNT_DATA,
"#cds", ATTR_CANONICALLY_DISCOVERABLE,
"#version", ATTR_VERSION));
Map<String, AttributeValue> attrValues = new HashMap<>(Map.of(
":data", AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
":cds", AttributeValues.fromBool(account.shouldBeVisibleInDirectory()),
":version", AttributeValues.fromInt(account.getVersion()),
":version_increment", AttributeValues.fromInt(1)));
final String updateExpression;
if (account.getUnidentifiedAccessKey().isPresent()) {
// if it's present in the account, also set the uak
attrNames.put("#uak", ATTR_UAK);
attrValues.put(":uak", AttributeValues.fromByteArray(account.getUnidentifiedAccessKey().get()));
updateExpression = "SET #data = :data, #cds = :cds, #uak = :uak ADD #version :version_increment";
} else {
updateExpression = "SET #data = :data, #cds = :cds ADD #version :version_increment";
}
updateItemRequest = UpdateItemRequest.builder()
.tableName(accountsTableName)
.key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
.updateExpression("SET #data = :data, #cds = :cds ADD #version :version_increment")
.conditionExpression("attribute_exists(#number) AND #version = :version")
.expressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164,
"#data", ATTR_ACCOUNT_DATA,
"#cds", ATTR_CANONICALLY_DISCOVERABLE,
"#version", ATTR_VERSION))
.expressionAttributeValues(Map.of(
":data", AttributeValues.fromByteArray(SystemMapper.getMapper().writeValueAsBytes(account)),
":cds", AttributeValues.fromBool(account.shouldBeVisibleInDirectory()),
":version", AttributeValues.fromInt(account.getVersion()),
":version_increment", AttributeValues.fromInt(1)))
.build();
.tableName(accountsTableName)
.key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid())))
.updateExpression(updateExpression)
.conditionExpression("attribute_exists(#number) AND #version = :version")
.expressionAttributeNames(attrNames)
.expressionAttributeValues(attrValues)
.build();
} catch (JsonProcessingException e) {
throw new IllegalArgumentException(e);
}

View File

@ -917,6 +917,9 @@ class AccountsTest {
assertThat(AttributeValues.getBool(get.item(), Accounts.ATTR_CANONICALLY_DISCOVERABLE,
!canonicallyDiscoverable)).isEqualTo(canonicallyDiscoverable);
assertThat(AttributeValues.getByteArray(get.item(), Accounts.ATTR_UAK, null))
.isEqualTo(expecting.getUnidentifiedAccessKey().orElse(null));
Account result = accounts.fromItem(get.item());
verifyStoredState(number, uuid, pni, result, expecting);
} else {