Add a feature flag to disable account normalization

This commit is contained in:
Ravi Khadiwala 2022-03-09 12:28:09 -06:00 committed by ravi-signal
parent 9cb098ad8a
commit 1571f14815
11 changed files with 81 additions and 24 deletions

@ -1 +1 @@
Subproject commit eaa72b4f4b3bd09f809520dc305478a85606f86a
Subproject commit fcd3d03d8ec5d4b1612e0c2cd618e4ac22728ecb

View File

@ -339,7 +339,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
config.getAppConfig().getConfigurationName(),
DynamicConfiguration.class);
Accounts accounts = new Accounts(dynamoDbClient,
Accounts accounts = new Accounts(dynamicConfigurationManager, dynamoDbClient,
config.getDynamoDbTables().getAccounts().getTableName(),
config.getDynamoDbTables().getAccounts().getPhoneNumberTableName(),
config.getDynamoDbTables().getAccounts().getPhoneNumberIdentifierTableName(),

View File

@ -52,6 +52,10 @@ public class DynamicConfiguration {
@Valid
private DynamicPushLatencyConfiguration pushLatency = new DynamicPushLatencyConfiguration(Collections.emptyMap());
@JsonProperty
@Valid
private DynamicUakMigrationConfiguration uakMigrationConfiguration = new DynamicUakMigrationConfiguration();
public Optional<DynamicExperimentEnrollmentConfiguration> getExperimentEnrollmentConfiguration(
final String experimentName) {
return Optional.ofNullable(experiments.get(experimentName));
@ -103,4 +107,6 @@ public class DynamicConfiguration {
return pushLatency;
}
public DynamicUakMigrationConfiguration getUakMigrationConfiguration() { return uakMigrationConfiguration; }
}

View File

@ -0,0 +1,12 @@
package org.whispersystems.textsecuregcm.configuration.dynamic;
import com.fasterxml.jackson.annotation.JsonProperty;
public class DynamicUakMigrationConfiguration {
@JsonProperty
private boolean enabled = true;
public boolean isEnabled() {
return enabled;
}
}

View File

@ -25,6 +25,7 @@ import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import org.whispersystems.textsecuregcm.util.SystemMapper;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
@ -70,6 +71,7 @@ public class Accounts extends AbstractDynamoDbStore {
// unidentified access key; byte[] or null
static final String ATTR_UAK = "UAK";
private DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager;
private final DynamoDbClient client;
private final String phoneNumberConstraintTableName;
@ -98,12 +100,13 @@ public class Accounts extends AbstractDynamoDbStore {
private static final Logger log = LoggerFactory.getLogger(Accounts.class);
public Accounts(DynamoDbClient client, String accountsTableName, String phoneNumberConstraintTableName,
public Accounts(final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager,
DynamoDbClient client, String accountsTableName, String phoneNumberConstraintTableName,
String phoneNumberIdentifierConstraintTableName, final String usernamesConstraintTableName,
final int scanPageSize) {
super(client);
this.dynamicConfigurationManager = dynamicConfigurationManager;
this.client = client;
this.phoneNumberConstraintTableName = phoneNumberConstraintTableName;
this.phoneNumberIdentifierConstraintTableName = phoneNumberIdentifierConstraintTableName;
@ -656,6 +659,11 @@ public class Accounts extends AbstractDynamoDbStore {
}
}
if (!this.dynamicConfigurationManager.getConfiguration().getUakMigrationConfiguration().isEnabled()) {
log.debug("Account normalization is disabled, skipping normalization for {} accounts", accountsToNormalize.size());
return allAccounts;
}
final int BATCH_SIZE = 25; // dynamodb max batch size
final String updateUakStatement = String.format("UPDATE %s SET %s = ? WHERE %s = ?", accountsTableName, ATTR_UAK, KEY_ACCOUNT_UUID);
for (List<Account> toNormalize : Lists.partition(accountsToNormalize, BATCH_SIZE)) {

View File

@ -135,7 +135,7 @@ public class AssignUsernameCommand extends EnvironmentCommand<WhisperServerConfi
VerificationCodeStore pendingAccounts = new VerificationCodeStore(dynamoDbClient,
configuration.getDynamoDbTables().getPendingAccounts().getTableName());
Accounts accounts = new Accounts(dynamoDbClient,
Accounts accounts = new Accounts(dynamicConfigurationManager, dynamoDbClient,
configuration.getDynamoDbTables().getAccounts().getTableName(),
configuration.getDynamoDbTables().getAccounts().getPhoneNumberTableName(),
configuration.getDynamoDbTables().getAccounts().getPhoneNumberIdentifierTableName(),

View File

@ -138,7 +138,7 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
VerificationCodeStore pendingAccounts = new VerificationCodeStore(dynamoDbClient,
configuration.getDynamoDbTables().getPendingAccounts().getTableName());
Accounts accounts = new Accounts(dynamoDbClient,
Accounts accounts = new Accounts(dynamicConfigurationManager, dynamoDbClient,
configuration.getDynamoDbTables().getAccounts().getTableName(),
configuration.getDynamoDbTables().getAccounts().getPhoneNumberTableName(),
configuration.getDynamoDbTables().getAccounts().getPhoneNumberIdentifierTableName(),

View File

@ -141,7 +141,7 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand<WhisperSer
VerificationCodeStore pendingAccounts = new VerificationCodeStore(dynamoDbClient,
configuration.getDynamoDbTables().getPendingAccounts().getTableName());
Accounts accounts = new Accounts(dynamoDbClient,
Accounts accounts = new Accounts(dynamicConfigurationManager, dynamoDbClient,
configuration.getDynamoDbTables().getAccounts().getTableName(),
configuration.getDynamoDbTables().getAccounts().getPhoneNumberTableName(),
configuration.getDynamoDbTables().getAccounts().getPhoneNumberIdentifierTableName(),

View File

@ -144,14 +144,6 @@ class AccountsManagerChangeNumberIntegrationTest {
ACCOUNTS_DYNAMO_EXTENSION.getDynamoDbClient().createTable(createPhoneNumberIdentifierTableRequest);
}
final Accounts accounts = new Accounts(
ACCOUNTS_DYNAMO_EXTENSION.getDynamoDbClient(),
ACCOUNTS_DYNAMO_EXTENSION.getTableName(),
NUMBERS_TABLE_NAME,
PNI_ASSIGNMENT_TABLE_NAME,
USERNAMES_TABLE_NAME,
SCAN_PAGE_SIZE);
{
@SuppressWarnings("unchecked") final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager =
mock(DynamicConfigurationManager.class);
@ -159,6 +151,15 @@ class AccountsManagerChangeNumberIntegrationTest {
DynamicConfiguration dynamicConfiguration = new DynamicConfiguration();
when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration);
final Accounts accounts = new Accounts(
dynamicConfigurationManager,
ACCOUNTS_DYNAMO_EXTENSION.getDynamoDbClient(),
ACCOUNTS_DYNAMO_EXTENSION.getTableName(),
NUMBERS_TABLE_NAME,
PNI_ASSIGNMENT_TABLE_NAME,
USERNAMES_TABLE_NAME,
SCAN_PAGE_SIZE);
deletedAccounts = new DeletedAccounts(DELETED_ACCOUNTS_DYNAMO_EXTENSION.getDynamoDbClient(),
DELETED_ACCOUNTS_DYNAMO_EXTENSION.getTableName(),
NEEDS_RECONCILIATION_INDEX_NAME);

View File

@ -39,6 +39,7 @@ import org.junit.jupiter.api.extension.RegisterExtension;
import org.mockito.ArgumentCaptor;
import org.mockito.stubbing.Answer;
import org.whispersystems.textsecuregcm.auth.AuthenticationCredentials;
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
import org.whispersystems.textsecuregcm.entities.AccountAttributes;
import org.whispersystems.textsecuregcm.entities.SignedPreKey;
import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
@ -118,7 +119,12 @@ class AccountsManagerConcurrentModificationIntegrationTest {
dynamoDbExtension.getDynamoDbClient().createTable(createPhoneNumberIdentifierTableRequest);
}
@SuppressWarnings("unchecked") final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager =
mock(DynamicConfigurationManager.class);
when(dynamicConfigurationManager.getConfiguration()).thenReturn(new DynamicConfiguration());
accounts = new Accounts(
dynamicConfigurationManager,
dynamoDbExtension.getDynamoDbClient(),
dynamoDbExtension.getTableName(),
NUMBERS_TABLE_NAME,

View File

@ -38,6 +38,7 @@ import org.junit.jupiter.api.extension.RegisterExtension;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
import org.whispersystems.textsecuregcm.entities.SignedPreKey;
import org.whispersystems.textsecuregcm.util.AttributeValues;
import org.whispersystems.textsecuregcm.util.SystemMapper;
@ -79,6 +80,7 @@ class AccountsTest {
.build())
.build();
private DynamicConfigurationManager<DynamicConfiguration> mockDynamicConfigManager;
private Accounts accounts;
@BeforeEach
@ -128,7 +130,14 @@ class AccountsTest {
dynamoDbExtension.getDynamoDbClient().createTable(createUsernamesTableRequest);
@SuppressWarnings("unchecked") DynamicConfigurationManager<DynamicConfiguration> m = mock(DynamicConfigurationManager.class);
mockDynamicConfigManager = m;
when(mockDynamicConfigManager.getConfiguration())
.thenReturn(new DynamicConfiguration());
this.accounts = new Accounts(
mockDynamicConfigManager,
dynamoDbExtension.getDynamoDbClient(),
dynamoDbExtension.getTableName(),
NUMBER_CONSTRAINT_TABLE_NAME,
@ -372,7 +381,7 @@ class AccountsTest {
void testUpdateWithMockTransactionConflictException() {
final DynamoDbClient dynamoDbClient = mock(DynamoDbClient.class);
accounts = new Accounts(dynamoDbClient,
accounts = new Accounts(mockDynamicConfigManager, dynamoDbClient,
dynamoDbExtension.getTableName(), NUMBER_CONSTRAINT_TABLE_NAME, PNI_CONSTRAINT_TABLE_NAME, USERNAME_CONSTRAINT_TABLE_NAME, SCAN_PAGE_SIZE);
when(dynamoDbClient.updateItem(any(UpdateItemRequest.class)))
@ -510,7 +519,7 @@ class AccountsTest {
when(client.updateItem(any(UpdateItemRequest.class)))
.thenThrow(RuntimeException.class);
Accounts accounts = new Accounts(client, ACCOUNTS_TABLE_NAME, NUMBER_CONSTRAINT_TABLE_NAME,
Accounts accounts = new Accounts(mockDynamicConfigManager, client, ACCOUNTS_TABLE_NAME, NUMBER_CONSTRAINT_TABLE_NAME,
PNI_CONSTRAINT_TABLE_NAME, USERNAME_CONSTRAINT_TABLE_NAME, SCAN_PAGE_SIZE);
Account account = generateAccount("+14151112222", UUID.randomUUID(), UUID.randomUUID());
@ -807,10 +816,21 @@ class AccountsTest {
assertThat(item).doesNotContainKey(Accounts.ATTR_UAK);
}
@Test
void testAddMissingUakAttribute() {
@ParameterizedTest
@ValueSource(booleans = {true, false})
void testAddMissingUakAttribute(boolean normalizeDisabled) throws JsonProcessingException {
final UUID accountIdentifier = UUID.randomUUID();
if (normalizeDisabled) {
final DynamicConfiguration config = DynamicConfigurationManager.parseConfiguration("""
captcha:
scoreFloor: 1.0
uakMigrationConfiguration:
enabled: false
""", DynamicConfiguration.class).orElseThrow();
when(mockDynamicConfigManager.getConfiguration()).thenReturn(config);
}
final Account account = generateAccount("+18005551234", accountIdentifier, UUID.randomUUID());
accounts.create(account);
@ -822,20 +842,24 @@ class AccountsTest {
.updateExpression("REMOVE #uak").build());
// crawling should return 1 account, and fix the discrepancy between
// the json blob and the top level attributes
// the json blob and the top level attributes if normalization is enabled
final AccountCrawlChunk allFromStart = accounts.getAllFromStart(1);
assertThat(allFromStart.getAccounts()).hasSize(1);
assertThat(allFromStart.getAccounts().get(0).getUuid()).isEqualTo(accountIdentifier);
// check that the attribute now exists at top level
// check whether normalization happened
final Map<String, AttributeValue> item = dynamoDbExtension.getDynamoDbClient()
.getItem(GetItemRequest.builder()
.tableName(ACCOUNTS_TABLE_NAME)
.key(Map.of(Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(accountIdentifier)))
.consistentRead(true)
.build()).item();
assertThat(item).containsEntry(Accounts.ATTR_UAK,
AttributeValues.fromByteArray(account.getUnidentifiedAccessKey().get()));
if (normalizeDisabled) {
assertThat(item).doesNotContainKey(Accounts.ATTR_UAK);
} else {
assertThat(item).containsEntry(Accounts.ATTR_UAK,
AttributeValues.fromByteArray(account.getUnidentifiedAccessKey().get()));
}
}
@ParameterizedTest
@ -975,7 +999,7 @@ class AccountsTest {
assertThat(AttributeValues.getByteArray(get.item(), Accounts.ATTR_UAK, null))
.isEqualTo(expecting.getUnidentifiedAccessKey().orElse(null));
Account result = accounts.fromItem(get.item());
Account result = Accounts.fromItem(get.item());
verifyStoredState(number, uuid, pni, result, expecting);
} else {
throw new AssertionError("No data");