Lock account and send notification when someone passes phone verification but fails reglock

service/src/test/java/org/whispersystems/textsecuregcm/auth/RegistrationLockVerificationManagerTest.java

@@ -11,10 +11,15 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import java.util.List;
 import java.util.UUID;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
@@ -23,15 +28,19 @@ import javax.ws.rs.WebApplicationException;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
-import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
 import org.whispersystems.textsecuregcm.entities.PhoneVerificationRequest;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
+import org.whispersystems.textsecuregcm.push.NotPushRegisteredException;
+import org.whispersystems.textsecuregcm.push.PushNotificationManager;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.storage.RegistrationRecoveryPasswordsManager;
+import org.whispersystems.textsecuregcm.tests.util.AccountsHelper;
 import org.whispersystems.textsecuregcm.util.Pair;
 
 class RegistrationLockVerificationManagerTest {
@@ -40,9 +49,12 @@ class RegistrationLockVerificationManagerTest {
   private final ClientPresenceManager clientPresenceManager = mock(ClientPresenceManager.class);
   private final ExternalServiceCredentialsGenerator backupServiceCredentialsGeneraor = mock(
       ExternalServiceCredentialsGenerator.class);
+  private final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager = mock(
+      RegistrationRecoveryPasswordsManager.class);
+  private static PushNotificationManager pushNotificationManager = mock(PushNotificationManager.class);
   private final RateLimiters rateLimiters = mock(RateLimiters.class);
   private final RegistrationLockVerificationManager registrationLockVerificationManager = new RegistrationLockVerificationManager(
-      accountsManager, clientPresenceManager, backupServiceCredentialsGeneraor, rateLimiters);
+      accountsManager, clientPresenceManager, backupServiceCredentialsGeneraor, registrationRecoveryPasswordsManager, pushNotificationManager, rateLimiters);
 
   private final RateLimiter pinLimiter = mock(RateLimiter.class);
 
@@ -51,22 +63,32 @@ class RegistrationLockVerificationManagerTest {
 
   @BeforeEach
   void setUp() {
+    clearInvocations(pushNotificationManager);
     when(rateLimiters.getPinLimiter()).thenReturn(pinLimiter);
     when(backupServiceCredentialsGeneraor.generateForUuid(any()))
         .thenReturn(mock(ExternalServiceCredentials.class));
 
+    final Device device = mock(Device.class);
+    when(device.getId()).thenReturn(Device.MASTER_ID);
+
+    AccountsHelper.setupMockUpdate(accountsManager);
+
     account = mock(Account.class);
     when(account.getUuid()).thenReturn(UUID.randomUUID());
     when(account.getNumber()).thenReturn("+18005551212");
+    when(account.getDevices()).thenReturn(List.of(device));
+
     existingRegistrationLock = mock(StoredRegistrationLock.class);
     when(account.getRegistrationLock()).thenReturn(existingRegistrationLock);
   }
 
   @ParameterizedTest
-  @EnumSource
-  void testErrors(RegistrationLockError error) throws Exception {
+  @MethodSource
+  void testErrors(RegistrationLockError error, boolean alreadyLocked) throws Exception {
 
     when(existingRegistrationLock.getStatus()).thenReturn(StoredRegistrationLock.Status.REQUIRED);
+    when(account.hasLockedCredentials()).thenReturn(alreadyLocked);
+    doThrow(new NotPushRegisteredException()).when(pushNotificationManager).sendAttemptLoginNotification(any(), any());
 
     final String submittedRegistrationLock = "reglock";
 
@@ -76,6 +98,16 @@ class RegistrationLockVerificationManagerTest {
         yield new Pair<>(WebApplicationException.class, e -> {
           if (e instanceof WebApplicationException wae) {
             assertEquals(RegistrationLockVerificationManager.FAILURE_HTTP_STATUS, wae.getResponse().getStatus());
+            verify(registrationRecoveryPasswordsManager).removeForNumber(account.getNumber());
+            verify(clientPresenceManager).disconnectAllPresences(account.getUuid(), List.of(Device.MASTER_ID));
+            try {
+              verify(pushNotificationManager).sendAttemptLoginNotification(any(), eq("failedRegistrationLock"));
+            } catch (NotPushRegisteredException npre) {}
+            if (alreadyLocked) {
+              verify(account, never()).lockAuthTokenHash();
+            } else {
+              verify(account).lockAuthTokenHash();
+            }
           } else {
             fail("Exception was not of expected type");
           }
@@ -85,6 +117,12 @@ class RegistrationLockVerificationManagerTest {
         when(existingRegistrationLock.verify(any())).thenReturn(true);
         doThrow(RateLimitExceededException.class).when(pinLimiter).validate(anyString());
         yield new Pair<>(RateLimitExceededException.class, ignored -> {
+          verify(account, never()).lockAuthTokenHash();
+          try {
+            verify(pushNotificationManager, never()).sendAttemptLoginNotification(any(), eq("failedRegistrationLock"));
+          } catch (NotPushRegisteredException npre) {}
+          verify(registrationRecoveryPasswordsManager, never()).removeForNumber(account.getNumber());
+          verify(clientPresenceManager, never()).disconnectAllPresences(account.getUuid(), List.of(Device.MASTER_ID));
         });
       }
     };
@@ -97,6 +135,14 @@ class RegistrationLockVerificationManagerTest {
     exceptionType.second().accept(e);
   }
 
+  static Stream<Arguments> testErrors() {
+    return Stream.of(
+        Arguments.of(RegistrationLockError.MISMATCH, true),
+        Arguments.of(RegistrationLockError.MISMATCH, false),
+        Arguments.of(RegistrationLockError.RATE_LIMITED, false)
+    );
+  }
+
   @ParameterizedTest
   @MethodSource
   void testSuccess(final StoredRegistrationLock.Status status, @Nullable final String submittedRegistrationLock) {
@@ -109,6 +155,10 @@ class RegistrationLockVerificationManagerTest {
         () -> registrationLockVerificationManager.verifyRegistrationLock(account, submittedRegistrationLock,
             "Signal-Android/4.68.3", RegistrationLockVerificationManager.Flow.REGISTRATION,
             PhoneVerificationRequest.VerificationType.SESSION));
+
+    verify(account, never()).lockAuthTokenHash();
+    verify(registrationRecoveryPasswordsManager, never()).removeForNumber(account.getNumber());
+    verify(clientPresenceManager, never()).disconnectAllPresences(account.getUuid(), List.of(Device.MASTER_ID));
   }
 
   static Stream<Arguments> testSuccess() {

service/src/test/java/org/whispersystems/textsecuregcm/controllers/AccountControllerTest.java

@@ -202,7 +202,8 @@ class AccountControllerTest {
       BACKUP_CFG);
 
   private static final RegistrationLockVerificationManager registrationLockVerificationManager = new RegistrationLockVerificationManager(
-      accountsManager, clientPresenceManager, backupCredentialsGenerator, rateLimiters);
+      accountsManager, clientPresenceManager, backupCredentialsGenerator, registrationRecoveryPasswordsManager,
+      pushNotificationManager, rateLimiters);
   private static final RegistrationCaptchaManager registrationCaptchaManager = new RegistrationCaptchaManager(
       captchaChecker, rateLimiters, Map.of(TEST_NUMBER, 123456), dynamicConfigurationManager);
 

service/src/test/java/org/whispersystems/textsecuregcm/push/PushNotificationManagerTest.java

@@ -9,6 +9,7 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.when;
@@ -137,6 +138,38 @@ class PushNotificationManagerTest {
     verify(apnSender).sendNotification(new PushNotification(deviceToken, PushNotification.TokenType.APN, PushNotification.NotificationType.RATE_LIMIT_CHALLENGE, challengeToken, account, device, true));
   }
 
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  void sendAttemptLoginNotification(final boolean isApn) throws NotPushRegisteredException {
+    final Account account = mock(Account.class);
+    final Device device = mock(Device.class);
+
+    final String deviceToken = "token";
+
+    when(device.getId()).thenReturn(Device.MASTER_ID);
+    if (isApn) {
+      when(device.getApnId()).thenReturn(deviceToken);
+      when(apnSender.sendNotification(any()))
+          .thenReturn(CompletableFuture.completedFuture(new SendPushNotificationResult(true, null, false)));
+    } else {
+      when(device.getGcmId()).thenReturn(deviceToken);
+      when(fcmSender.sendNotification(any()))
+          .thenReturn(CompletableFuture.completedFuture(new SendPushNotificationResult(true, null, false)));
+    }
+    when(account.getDevice(Device.MASTER_ID)).thenReturn(Optional.of(device));
+
+    pushNotificationManager.sendAttemptLoginNotification(account, "someContext");
+
+    if (isApn){
+      verify(apnSender).sendNotification(new PushNotification(deviceToken, PushNotification.TokenType.APN,
+          PushNotification.NotificationType.ATTEMPT_LOGIN_NOTIFICATION_HIGH_PRIORITY, "someContext", account, device, true));
+      verify(apnPushNotificationScheduler).scheduleBackgroundNotification(account, device);
+    } else {
+      verify(fcmSender, times(1)).sendNotification(new PushNotification(deviceToken, PushNotification.TokenType.FCM,
+          PushNotification.NotificationType.ATTEMPT_LOGIN_NOTIFICATION_HIGH_PRIORITY, "someContext", account, device, true));
+    }
+  }
+
   @ParameterizedTest
   @ValueSource(booleans = {true, false})
   void testSendNotificationFcm(final boolean urgent) {

service/src/test/java/org/whispersystems/textsecuregcm/tests/util/AccountsHelper.java

@@ -132,6 +132,7 @@ public class AccountsHelper {
           case "getIdentityKey" -> when(updatedAccount.getIdentityKey()).thenAnswer(stubbing);
           case "getBadges" -> when(updatedAccount.getBadges()).thenAnswer(stubbing);
           case "getLastSeen" -> when(updatedAccount.getLastSeen()).thenAnswer(stubbing);
+          case "hasLockedCredentials" -> when(updatedAccount.hasLockedCredentials()).thenAnswer(stubbing);
           default -> throw new IllegalArgumentException("unsupported method: Account#" + stubbing.getInvocation().getMethod().getName());
         }
       }