Return a non-stale account from base authenticator when last seen is updated

service/src/main/java/org/whispersystems/textsecuregcm/auth/BaseAccountAuthenticator.java

@@ -83,8 +83,7 @@ public class BaseAccountAuthenticator {
       if (device.get().getAuthenticationCredentials().verify(basicCredentials.getPassword())) {
         succeeded = true;
         account.get().setAuthenticatedDevice(device.get());
-        updateLastSeen(account.get(), device.get());
+        return Optional.of(updateLastSeen(account.get(), device.get()));
-        return account;
       }
 
       return Optional.empty();
@@ -109,7 +108,7 @@ public class BaseAccountAuthenticator {
   }
 
   @VisibleForTesting
-  public void updateLastSeen(Account account, Device device) {
+  public Account updateLastSeen(Account account, Device device) {
     final long lastSeenOffsetSeconds   = Math.abs(account.getUuid().getLeastSignificantBits()) % ChronoUnit.DAYS.getDuration().toSeconds();
     final long todayInMillisWithOffset = Util.todayInMillisGivenOffsetFromNow(clock, Duration.ofSeconds(lastSeenOffsetSeconds).negated());
 
@@ -117,8 +116,10 @@ public class BaseAccountAuthenticator {
       Metrics.summary(DAYS_SINCE_LAST_SEEN_DISTRIBUTION_NAME, IS_PRIMARY_DEVICE_TAG, String.valueOf(device.isMaster()))
           .record(Duration.ofMillis(todayInMillisWithOffset - device.getLastSeen()).toDays());
 
-      accountsManager.updateDevice(account, device.getId(), d -> d.setLastSeen(Util.todayInMillis(clock)));
+      return accountsManager.updateDevice(account, device.getId(), d -> d.setLastSeen(Util.todayInMillis(clock)));
     }
+
+    return account;
   }
 
 }

service/src/test/java/org/whispersystems/textsecuregcm/tests/auth/BaseAccountAuthenticatorTest.java

@@ -65,14 +65,17 @@ class BaseAccountAuthenticatorTest {
     final Device device1 = acct1.getDevices().stream().findFirst().get();
     final Device device2 = acct2.getDevices().stream().findFirst().get();
 
-    baseAccountAuthenticator.updateLastSeen(acct1, device1);
+    final Account updatedAcct1 = baseAccountAuthenticator.updateLastSeen(acct1, device1);
-    baseAccountAuthenticator.updateLastSeen(acct2, device2);
+    final Account updatedAcct2 = baseAccountAuthenticator.updateLastSeen(acct2, device2);
 
     verify(accountsManager, never()).updateDevice(eq(acct1), anyLong(), any());
     verify(accountsManager).updateDevice(eq(acct2), anyLong(), any());
 
     assertThat(device1.getLastSeen()).isEqualTo(yesterday);
     assertThat(device2.getLastSeen()).isEqualTo(today);
+
+    assertThat(acct1).isSameAs(updatedAcct1);
+    assertThat(acct2).isNotSameAs(updatedAcct2);
   }
 
   @Test
@@ -82,14 +85,17 @@ class BaseAccountAuthenticatorTest {
     final Device device1 = acct1.getDevices().stream().findFirst().get();
     final Device device2 = acct2.getDevices().stream().findFirst().get();
 
-    baseAccountAuthenticator.updateLastSeen(acct1, device1);
+    final Account updatedAcct1 = baseAccountAuthenticator.updateLastSeen(acct1, device1);
-    baseAccountAuthenticator.updateLastSeen(acct2, device2);
+    final Account updatedAcct2 = baseAccountAuthenticator.updateLastSeen(acct2, device2);
 
     verify(accountsManager, never()).updateDevice(eq(acct1), anyLong(), any());
     verify(accountsManager, never()).updateDevice(eq(acct2), anyLong(), any());
 
     assertThat(device1.getLastSeen()).isEqualTo(yesterday);
     assertThat(device2.getLastSeen()).isEqualTo(yesterday);
+
+    assertThat(acct1).isSameAs(updatedAcct1);
+    assertThat(acct2).isSameAs(updatedAcct2);
   }
 
   @Test
@@ -99,14 +105,17 @@ class BaseAccountAuthenticatorTest {
     final Device device1 = acct1.getDevices().stream().findFirst().get();
     final Device device2 = acct2.getDevices().stream().findFirst().get();
 
-    baseAccountAuthenticator.updateLastSeen(acct1, device1);
+    final Account updatedAcct1 = baseAccountAuthenticator.updateLastSeen(acct1, device1);
-    baseAccountAuthenticator.updateLastSeen(acct2, device2);
+    final Account updatedAcct2 = baseAccountAuthenticator.updateLastSeen(acct2, device2);
 
     verify(accountsManager).updateDevice(eq(acct1), anyLong(), any());
     verify(accountsManager).updateDevice(eq(acct2), anyLong(), any());
 
     assertThat(device1.getLastSeen()).isEqualTo(today);
     assertThat(device2.getLastSeen()).isEqualTo(today);
+
+    assertThat(updatedAcct1).isNotSameAs(acct1);
+    assertThat(updatedAcct2).isNotSameAs(acct2);
   }
 
   @Test

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

@@ -26,11 +26,22 @@ import org.whispersystems.textsecuregcm.util.SystemMapper;
 public class AccountsHelper {
 
   public static void setupMockUpdate(final AccountsManager mockAccountsManager) {
+    setupMockUpdate(mockAccountsManager, true);
+  }
+
+  /**
+   * Only for use by {@link AuthHelper}
+   */
+  public static void setupMockUpdateForAuthHelper(final AccountsManager mockAccountsManager) {
+    setupMockUpdate(mockAccountsManager, false);
+  }
+
+  private static void setupMockUpdate(final AccountsManager mockAccountsManager, final boolean markStale) {
     when(mockAccountsManager.update(any(), any())).thenAnswer(answer -> {
       final Account account = answer.getArgument(0, Account.class);
       answer.getArgument(1, Consumer.class).accept(account);
 
-      return copyAndMarkStale(account);
+      return markStale ? copyAndMarkStale(account) : account;
     });
 
     when(mockAccountsManager.updateDevice(any(), anyLong(), any())).thenAnswer(answer -> {
@@ -38,7 +49,7 @@ public class AccountsHelper {
       final Long deviceId = answer.getArgument(1, Long.class);
       account.getDevice(deviceId).ifPresent(answer.getArgument(2, Consumer.class));
 
-      return copyAndMarkStale(account);
+      return markStale ? copyAndMarkStale(account) : account;
     });
   }
 

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

@@ -8,6 +8,7 @@ package org.whispersystems.textsecuregcm.tests.util;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.when;
 
 import com.google.common.collect.ImmutableMap;
@@ -141,6 +142,8 @@ public class AuthHelper {
 
     when(VALID_ACCOUNT.getIdentityKey()).thenReturn(VALID_IDENTITY);
 
+    reset(ACCOUNTS_MANAGER);
+
     when(ACCOUNTS_MANAGER.get(VALID_NUMBER)).thenReturn(Optional.of(VALID_ACCOUNT));
     when(ACCOUNTS_MANAGER.get(VALID_UUID)).thenReturn(Optional.of(VALID_ACCOUNT));
     when(ACCOUNTS_MANAGER.get(argThat((ArgumentMatcher<AmbiguousIdentifier>) identifier -> identifier != null && identifier.hasNumber() && identifier.getNumber().equals(VALID_NUMBER)))).thenReturn(Optional.of(VALID_ACCOUNT));
@@ -161,6 +164,8 @@ public class AuthHelper {
     when(ACCOUNTS_MANAGER.get(argThat((ArgumentMatcher<AmbiguousIdentifier>) identifier -> identifier != null && identifier.hasNumber() && identifier.getNumber().equals(UNDISCOVERABLE_NUMBER)))).thenReturn(Optional.of(UNDISCOVERABLE_ACCOUNT));
     when(ACCOUNTS_MANAGER.get(argThat((ArgumentMatcher<AmbiguousIdentifier>) identifier -> identifier != null && identifier.hasUuid() && identifier.getUuid().equals(UNDISCOVERABLE_UUID)))).thenReturn(Optional.of(UNDISCOVERABLE_ACCOUNT));
 
+    AccountsHelper.setupMockUpdateForAuthHelper(ACCOUNTS_MANAGER);
+
     for (TestAccount testAccount : TEST_ACCOUNTS) {
       testAccount.setup(ACCOUNTS_MANAGER);
     }