Filter expired items from Dynamo

service/src/main/java/org/whispersystems/textsecuregcm/storage/PushChallengeDynamoDb.java

@@ -5,11 +5,11 @@
 
 package org.whispersystems.textsecuregcm.storage;
 
+import com.google.common.annotations.VisibleForTesting;
 import java.time.Clock;
 import java.time.Duration;
 import java.util.Map;
 import java.util.UUID;
-import com.google.common.annotations.VisibleForTesting;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
@@ -29,7 +29,8 @@ public class PushChallengeDynamoDb extends AbstractDynamoDbStore {
   static final String ATTR_TTL = "T";
 
   private static final Map<String, String> UUID_NAME_MAP = Map.of("#uuid", KEY_ACCOUNT_UUID);
-  private static final Map<String, String> CHALLENGE_TOKEN_NAME_MAP = Map.of("#challenge", ATTR_CHALLENGE_TOKEN);
+  private static final Map<String, String> CHALLENGE_TOKEN_NAME_MAP = Map.of("#challenge", ATTR_CHALLENGE_TOKEN, "#ttl",
+      ATTR_TTL);
 
   public PushChallengeDynamoDb(final DynamoDbClient dynamoDB, final String tableName) {
     this(dynamoDB, tableName, Clock.systemUTC());
@@ -87,9 +88,10 @@ public class PushChallengeDynamoDb extends AbstractDynamoDbStore {
       db().deleteItem(DeleteItemRequest.builder()
           .tableName(tableName)
           .key(Map.of(KEY_ACCOUNT_UUID, AttributeValues.fromUUID(accountUuid)))
-          .conditionExpression("#challenge = :challenge")
+          .conditionExpression("#challenge = :challenge AND #ttl >= :currentTime")
           .expressionAttributeNames(CHALLENGE_TOKEN_NAME_MAP)
-          .expressionAttributeValues(Map.of(":challenge", AttributeValues.fromByteArray(challengeToken)))
+          .expressionAttributeValues(Map.of(":challenge", AttributeValues.fromByteArray(challengeToken),
+              ":currentTime", AttributeValues.fromLong(clock.instant().getEpochSecond())))
           .build());
       return true;
     } catch (final ConditionalCheckFailedException e) {

service/src/main/java/org/whispersystems/textsecuregcm/storage/VerificationCodeStore.java

@@ -5,10 +5,15 @@
 
 package org.whispersystems.textsecuregcm.storage;
 
+import static com.codahale.metrics.MetricRegistry.name;
+
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.google.common.annotations.VisibleForTesting;
 import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Timer;
+import java.time.Instant;
+import java.util.Map;
+import java.util.Optional;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.StoredVerificationCode;
@@ -19,11 +24,6 @@ import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
 import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
-import java.time.Instant;
-import java.util.Map;
-import java.util.Optional;
-
-import static com.codahale.metrics.MetricRegistry.name;
 
 public class VerificationCodeStore {
 
@@ -83,7 +83,8 @@ public class VerificationCodeStore {
 
       try {
         return response.hasItem()
-            ? Optional.of(SystemMapper.getMapper().readValue(response.item().get(ATTR_STORED_CODE).s(), StoredVerificationCode.class))
+            ? filterMaybeExpiredCode(
+            SystemMapper.getMapper().readValue(response.item().get(ATTR_STORED_CODE).s(), StoredVerificationCode.class))
             : Optional.empty();
       } catch (final JsonProcessingException e) {
         log.error("Failed to parse stored verification code", e);
@@ -92,6 +93,16 @@ public class VerificationCodeStore {
     });
   }
 
+  private Optional<StoredVerificationCode> filterMaybeExpiredCode(StoredVerificationCode storedVerificationCode) {
+    // It's possible for DynamoDB to return items after their expiration time (although it is very unlikely for small
+    // tables)
+    if (getExpirationTimestamp(storedVerificationCode) < Instant.now().getEpochSecond()) {
+      return Optional.empty();
+    }
+
+    return Optional.of(storedVerificationCode);
+  }
+
   public void remove(final String number) {
     removeTimer.record(() -> {
       dynamoDbClient.deleteItem(DeleteItemRequest.builder()

service/src/test/java/org/whispersystems/textsecuregcm/storage/PushChallengeDynamoDbTest.java

@@ -62,6 +62,8 @@ class PushChallengeDynamoDbTest {
     assertFalse(pushChallengeDynamoDb.remove(uuid, token));
     assertTrue(pushChallengeDynamoDb.add(uuid, token, Duration.ofMinutes(1)));
     assertTrue(pushChallengeDynamoDb.remove(uuid, token));
+    assertTrue(pushChallengeDynamoDb.add(uuid, token, Duration.ofMinutes(-1)));
+    assertFalse(pushChallengeDynamoDb.remove(uuid, token));
   }
 
   @Test

service/src/test/java/org/whispersystems/textsecuregcm/storage/VerificationCodeStoreTest.java

@@ -5,18 +5,20 @@
 
 package org.whispersystems.textsecuregcm.storage;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Objects;
+import java.util.Optional;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.whispersystems.textsecuregcm.auth.StoredVerificationCode;
 import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
 import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
-import java.util.Objects;
-import java.util.Optional;
-
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 class VerificationCodeStoreTest {
 
@@ -26,6 +28,10 @@ class VerificationCodeStoreTest {
 
   private static final String PHONE_NUMBER = "+14151112222";
 
+  private static final long VALID_TIMESTAMP = Instant.now().toEpochMilli();
+  private static final long EXPIRED_TIMESTAMP = Instant.now().minus(StoredVerificationCode.EXPIRATION).minus(
+      Duration.ofHours(1)).toEpochMilli();
+
   @RegisterExtension
   static final DynamoDbExtension DYNAMO_DB_EXTENSION = DynamoDbExtension.builder()
       .tableName(TABLE_NAME)
@@ -45,8 +51,8 @@ class VerificationCodeStoreTest {
   void testStoreAndFind() {
     assertEquals(Optional.empty(), verificationCodeStore.findForNumber(PHONE_NUMBER));
 
-    final StoredVerificationCode originalCode = new StoredVerificationCode("1234", 1111, "abcd", "0987");
+    final StoredVerificationCode originalCode = new StoredVerificationCode("1234", VALID_TIMESTAMP, "abcd", "0987");
-    final StoredVerificationCode secondCode = new StoredVerificationCode("5678", 2222, "efgh", "7890");
+    final StoredVerificationCode secondCode = new StoredVerificationCode("5678", VALID_TIMESTAMP, "efgh", "7890");
 
     verificationCodeStore.insert(PHONE_NUMBER, originalCode);
     {
@@ -69,11 +75,14 @@ class VerificationCodeStoreTest {
   void testRemove() {
     assertEquals(Optional.empty(), verificationCodeStore.findForNumber(PHONE_NUMBER));
 
-    verificationCodeStore.insert(PHONE_NUMBER, new StoredVerificationCode("1234", 1111, "abcd", "0987"));
+    verificationCodeStore.insert(PHONE_NUMBER, new StoredVerificationCode("1234", VALID_TIMESTAMP, "abcd", "0987"));
     assertTrue(verificationCodeStore.findForNumber(PHONE_NUMBER).isPresent());
 
     verificationCodeStore.remove(PHONE_NUMBER);
     assertFalse(verificationCodeStore.findForNumber(PHONE_NUMBER).isPresent());
+
+    verificationCodeStore.insert(PHONE_NUMBER, new StoredVerificationCode("1234", EXPIRED_TIMESTAMP, "abcd", "0987"));
+    assertFalse(verificationCodeStore.findForNumber(PHONE_NUMBER).isPresent());
   }
 
   private static boolean storedVerificationCodesAreEqual(final StoredVerificationCode first, final StoredVerificationCode second) {