From 236b0496d342ee3bb2350bc2ece33332446bb8a7 Mon Sep 17 00:00:00 2001 From: Ravi Khadiwala Date: Wed, 4 Dec 2024 14:15:46 -0600 Subject: [PATCH] Write a set field in IssuedReceiptsManager --- .../storage/IssuedReceiptsManager.java | 15 +++++++- .../storage/IssuedReceiptsManagerTest.java | 38 +++++++++++++++++-- 2 files changed, 47 insertions(+), 6 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManager.java index 1a19c065c..2fb4340fe 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManager.java @@ -17,8 +17,11 @@ import java.security.InvalidKeyException; import java.security.NoSuchAlgorithmException; import java.time.Duration; import java.time.Instant; +import java.util.Collections; +import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.UUID; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.function.Consumer; @@ -27,6 +30,8 @@ import javax.crypto.Mac; import javax.crypto.spec.SecretKeySpec; import org.signal.libsignal.zkgroup.receipts.ReceiptCredentialRequest; import org.whispersystems.textsecuregcm.subscriptions.PaymentProvider; +import org.whispersystems.textsecuregcm.util.AttributeValues; +import software.amazon.awssdk.core.SdkBytes; import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException; @@ -39,6 +44,8 @@ public class IssuedReceiptsManager { public static final String KEY_ISSUED_RECEIPT_TAG = "B"; // B public static final String KEY_EXPIRATION = "E"; // N + public static final String KEY_ISSUED_RECEIPT_TAG_SET = "T"; // BS + private final String table; private final Duration expiration; private final DynamoDbAsyncClient dynamoDbAsyncClient; @@ -79,6 +86,7 @@ public class IssuedReceiptsManager { } else { key = s(processor.name() + "_" + processorItemId); } + final byte[] tag = generateIssuedReceiptTag(request); UpdateItemRequest updateItemRequest = UpdateItemRequest.builder() .tableName(table) .key(Map.of(KEY_PROCESSOR_ITEM_ID, key)) @@ -86,13 +94,16 @@ public class IssuedReceiptsManager { .returnValues(ReturnValue.NONE) .updateExpression("SET " + "#tag = if_not_exists(#tag, :tag), " - + "#exp = if_not_exists(#exp, :exp)") + + "#exp = if_not_exists(#exp, :exp) " + + "ADD #tags :singletonTag") .expressionAttributeNames(Map.of( "#key", KEY_PROCESSOR_ITEM_ID, "#tag", KEY_ISSUED_RECEIPT_TAG, + "#tags", KEY_ISSUED_RECEIPT_TAG_SET, "#exp", KEY_EXPIRATION)) .expressionAttributeValues(Map.of( - ":tag", b(generateIssuedReceiptTag(request)), + ":tag", b(tag), + ":singletonTag", AttributeValue.fromBs(List.of(SdkBytes.fromByteArray(tag))), ":exp", n(now.plus(expiration).getEpochSecond()))) .build(); return dynamoDbAsyncClient.updateItem(updateItemRequest).handle((updateItemResponse, throwable) -> { diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManagerTest.java index 14d564d8c..ba17a3708 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManagerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManagerTest.java @@ -12,7 +12,11 @@ import static org.mockito.Mockito.when; import jakarta.ws.rs.ClientErrorException; import java.time.Duration; import java.time.Instant; +import java.util.List; +import java.util.Map; +import java.util.Set; import java.util.concurrent.CompletableFuture; +import java.util.stream.Collectors; import org.assertj.core.api.Condition; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -20,7 +24,13 @@ import org.junit.jupiter.api.extension.RegisterExtension; import org.signal.libsignal.zkgroup.receipts.ReceiptCredentialRequest; import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema.Tables; import org.whispersystems.textsecuregcm.subscriptions.PaymentProvider; +import org.whispersystems.textsecuregcm.util.AttributeValues; import org.whispersystems.textsecuregcm.util.TestRandomUtil; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.GetItemRequest; +import software.amazon.awssdk.services.dynamodb.model.GetItemResponse; class IssuedReceiptsManagerTest { @@ -29,12 +39,10 @@ class IssuedReceiptsManagerTest { @RegisterExtension static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(Tables.ISSUED_RECEIPTS); - ReceiptCredentialRequest receiptCredentialRequest; IssuedReceiptsManager issuedReceiptsManager; @BeforeEach void beforeEach() { - receiptCredentialRequest = mock(ReceiptCredentialRequest.class); issuedReceiptsManager = new IssuedReceiptsManager( Tables.ISSUED_RECEIPTS.tableName(), Duration.ofDays(90), @@ -45,12 +53,18 @@ class IssuedReceiptsManagerTest { @Test void testRecordIssuance() { Instant now = Instant.ofEpochSecond(NOW_EPOCH_SECONDS); - byte[] request1 = TestRandomUtil.nextBytes(20); - when(receiptCredentialRequest.serialize()).thenReturn(request1); + final ReceiptCredentialRequest receiptCredentialRequest = randomReceiptCredentialRequest(); CompletableFuture future = issuedReceiptsManager.recordIssuance("item-1", PaymentProvider.STRIPE, receiptCredentialRequest, now); assertThat(future).succeedsWithin(Duration.ofSeconds(3)); + final Map item = getItem("item-1").item(); + final Set tagSet = item.get(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG_SET).bs() + .stream() + .map(SdkBytes::asByteArray) + .collect(Collectors.toSet()); + assertThat(tagSet).containsExactly(item.get(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG).b().asByteArray()); + // same request should succeed future = issuedReceiptsManager.recordIssuance("item-1", PaymentProvider.STRIPE, receiptCredentialRequest, now); @@ -74,4 +88,20 @@ class IssuedReceiptsManagerTest { now); assertThat(future).succeedsWithin(Duration.ofSeconds(3)); } + + + private GetItemResponse getItem(final String itemId) { + final DynamoDbClient client = DYNAMO_DB_EXTENSION.getDynamoDbClient(); + return client.getItem(GetItemRequest.builder() + .tableName(Tables.ISSUED_RECEIPTS.tableName()) + .key(Map.of(IssuedReceiptsManager.KEY_PROCESSOR_ITEM_ID, AttributeValues.s(itemId))) + .build()); + } + + private static ReceiptCredentialRequest randomReceiptCredentialRequest() { + final ReceiptCredentialRequest request = mock(ReceiptCredentialRequest.class); + final byte[] bytes = TestRandomUtil.nextBytes(20); + when(request.serialize()).thenReturn(bytes); + return request; + } }