Clear one-time pre-keys when PNIs are assigned

This commit is contained in:
Jon Chambers 2022-02-28 11:21:45 -05:00 committed by Jon Chambers
parent eee6307789
commit 179b4a69eb
2 changed files with 29 additions and 6 deletions

View File

@ -194,6 +194,7 @@ public class AccountsManager {
if (!originalUuid.equals(actualUuid)) { if (!originalUuid.equals(actualUuid)) {
messagesManager.clear(actualUuid); messagesManager.clear(actualUuid);
keys.delete(actualUuid); keys.delete(actualUuid);
keys.delete(account.getPhoneNumberIdentifier());
profilesManager.deleteAll(actualUuid); profilesManager.deleteAll(actualUuid);
} }
@ -221,6 +222,7 @@ public class AccountsManager {
public Account changeNumber(final Account account, final String number) throws InterruptedException { public Account changeNumber(final Account account, final String number) throws InterruptedException {
final String originalNumber = account.getNumber(); final String originalNumber = account.getNumber();
final UUID originalPhoneNumberIdentifier = account.getPhoneNumberIdentifier();
if (originalNumber.equals(number)) { if (originalNumber.equals(number)) {
return account; return account;
@ -261,6 +263,9 @@ public class AccountsManager {
updatedAccount.set(numberChangedAccount); updatedAccount.set(numberChangedAccount);
directoryQueue.changePhoneNumber(numberChangedAccount, originalNumber, number); directoryQueue.changePhoneNumber(numberChangedAccount, originalNumber, number);
keys.delete(phoneNumberIdentifier);
keys.delete(originalPhoneNumberIdentifier);
return displacedUuid; return displacedUuid;
}); });

View File

@ -65,6 +65,8 @@ class AccountsManagerTest {
private ProfilesManager profilesManager; private ProfilesManager profilesManager;
private ReservedUsernames reservedUsernames; private ReservedUsernames reservedUsernames;
private Map<String, UUID> phoneNumberIdentifiersByE164;
private RedisAdvancedClusterCommands<String, String> commands; private RedisAdvancedClusterCommands<String, String> commands;
private AccountsManager accountsManager; private AccountsManager accountsManager;
@ -112,7 +114,7 @@ class AccountsManagerTest {
when(backupClient.deleteBackups(any())).thenReturn(CompletableFuture.completedFuture(null)); when(backupClient.deleteBackups(any())).thenReturn(CompletableFuture.completedFuture(null));
final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class); final PhoneNumberIdentifiers phoneNumberIdentifiers = mock(PhoneNumberIdentifiers.class);
final Map<String, UUID> phoneNumberIdentifiersByE164 = new HashMap<>(); phoneNumberIdentifiersByE164 = new HashMap<>();
when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString())).thenAnswer((Answer<UUID>) invocation -> { when(phoneNumberIdentifiers.getPhoneNumberIdentifier(anyString())).thenAnswer((Answer<UUID>) invocation -> {
final String number = invocation.getArgument(0, String.class); final String number = invocation.getArgument(0, String.class);
@ -532,9 +534,13 @@ class AccountsManagerTest {
final AccountAttributes attributes = new AccountAttributes(false, 0, null, null, true, null); final AccountAttributes attributes = new AccountAttributes(false, 0, null, null, true, null);
accountsManager.create(e164, "password", null, attributes, new ArrayList<>()); accountsManager.create(e164, "password", null, attributes, new ArrayList<>());
verify(accounts).create( assertTrue(phoneNumberIdentifiersByE164.containsKey(e164));
argThat(account -> e164.equals(account.getNumber()) && existingUuid.equals(account.getUuid())));
verify(accounts)
.create(argThat(account -> e164.equals(account.getNumber()) && existingUuid.equals(account.getUuid())));
verify(keys).delete(existingUuid); verify(keys).delete(existingUuid);
verify(keys).delete(phoneNumberIdentifiersByE164.get(e164));
verify(messagesManager).clear(existingUuid); verify(messagesManager).clear(existingUuid);
verify(profilesManager).deleteAll(existingUuid); verify(profilesManager).deleteAll(existingUuid);
} }
@ -646,13 +652,18 @@ class AccountsManagerTest {
final String originalNumber = "+14152222222"; final String originalNumber = "+14152222222";
final String targetNumber = "+14153333333"; final String targetNumber = "+14153333333";
final UUID uuid = UUID.randomUUID(); final UUID uuid = UUID.randomUUID();
final UUID originalPni = UUID.randomUUID();
Account account = new Account(originalNumber, uuid, UUID.randomUUID(), new HashSet<>(), new byte[16]); Account account = new Account(originalNumber, uuid, originalPni, new HashSet<>(), new byte[16]);
account = accountsManager.changeNumber(account, targetNumber); account = accountsManager.changeNumber(account, targetNumber);
assertEquals(targetNumber, account.getNumber()); assertEquals(targetNumber, account.getNumber());
assertTrue(phoneNumberIdentifiersByE164.containsKey(targetNumber));
verify(directoryQueue).changePhoneNumber(argThat(a -> a.getUuid().equals(uuid)), eq(originalNumber), eq(targetNumber)); verify(directoryQueue).changePhoneNumber(argThat(a -> a.getUuid().equals(uuid)), eq(originalNumber), eq(targetNumber));
verify(keys).delete(originalPni);
verify(keys).delete(phoneNumberIdentifiersByE164.get(targetNumber));
} }
@Test @Test
@ -665,6 +676,7 @@ class AccountsManagerTest {
assertEquals(number, account.getNumber()); assertEquals(number, account.getNumber());
verify(deletedAccountsManager, never()).lockAndPut(anyString(), anyString(), any()); verify(deletedAccountsManager, never()).lockAndPut(anyString(), anyString(), any());
verify(directoryQueue, never()).changePhoneNumber(any(), any(), any()); verify(directoryQueue, never()).changePhoneNumber(any(), any(), any());
verify(keys, never()).delete(any());
} }
@Test @Test
@ -676,17 +688,23 @@ class AccountsManagerTest {
final String targetNumber = "+14153333333"; final String targetNumber = "+14153333333";
final UUID existingAccountUuid = UUID.randomUUID(); final UUID existingAccountUuid = UUID.randomUUID();
final UUID uuid = UUID.randomUUID(); final UUID uuid = UUID.randomUUID();
final UUID originalPni = UUID.randomUUID();
final UUID targetPni = UUID.randomUUID();
final Account existingAccount = new Account(targetNumber, existingAccountUuid, UUID.randomUUID(), new HashSet<>(), new byte[16]); final Account existingAccount = new Account(targetNumber, existingAccountUuid, targetPni, new HashSet<>(), new byte[16]);
when(accounts.getByE164(targetNumber)).thenReturn(Optional.of(existingAccount)); when(accounts.getByE164(targetNumber)).thenReturn(Optional.of(existingAccount));
Account account = new Account(originalNumber, uuid, UUID.randomUUID(), new HashSet<>(), new byte[16]); Account account = new Account(originalNumber, uuid, originalPni, new HashSet<>(), new byte[16]);
account = accountsManager.changeNumber(account, targetNumber); account = accountsManager.changeNumber(account, targetNumber);
assertEquals(targetNumber, account.getNumber()); assertEquals(targetNumber, account.getNumber());
assertTrue(phoneNumberIdentifiersByE164.containsKey(targetNumber));
verify(directoryQueue).changePhoneNumber(argThat(a -> a.getUuid().equals(uuid)), eq(originalNumber), eq(targetNumber)); verify(directoryQueue).changePhoneNumber(argThat(a -> a.getUuid().equals(uuid)), eq(originalNumber), eq(targetNumber));
verify(directoryQueue).deleteAccount(existingAccount); verify(directoryQueue).deleteAccount(existingAccount);
verify(keys).delete(originalPni);
verify(keys).delete(targetPni);
} }
@Test @Test