diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/PushChallengeDynamoDb.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/PushChallengeDynamoDb.java index b221258f9..2ad0600d6 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/PushChallengeDynamoDb.java +++ b/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 UUID_NAME_MAP = Map.of("#uuid", KEY_ACCOUNT_UUID); - private static final Map CHALLENGE_TOKEN_NAME_MAP = Map.of("#challenge", ATTR_CHALLENGE_TOKEN); + private static final Map 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) { diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/VerificationCodeStore.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/VerificationCodeStore.java index 7be0675ad..78dc9a839 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/VerificationCodeStore.java +++ b/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 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() diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/PushChallengeDynamoDbTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/PushChallengeDynamoDbTest.java index 183f8edfe..5da07d2c2 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/PushChallengeDynamoDbTest.java +++ b/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 diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/VerificationCodeStoreTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/VerificationCodeStoreTest.java index b3fb86e56..2249d6e05 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/VerificationCodeStoreTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/VerificationCodeStoreTest.java @@ -5,27 +5,33 @@ 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 { private VerificationCodeStore verificationCodeStore; private static final String TABLE_NAME = "verification_code_test"; - + 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 secondCode = new StoredVerificationCode("5678", 2222, "efgh", "7890"); + final StoredVerificationCode originalCode = new StoredVerificationCode("1234", VALID_TIMESTAMP, "abcd", "0987"); + 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) {