Write a set field in IssuedReceiptsManager

Author: ravi-signal

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

@@ -17,8 +17,11 @@
 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.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,20 +86,24 @@ public CompletableFuture<Void> recordIssuance(
     } 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))
         .conditionExpression("attribute_not_exists(#key) OR #tag = :tag")
         .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) -> {

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

@@ -12,15 +12,25 @@
 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;
 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 @@ void beforeEach() {
   @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<Void> future = issuedReceiptsManager.recordIssuance("item-1", PaymentProvider.STRIPE,
         receiptCredentialRequest, now);
     assertThat(future).succeedsWithin(Duration.ofSeconds(3));
 
+    final Map<String, AttributeValue> item = getItem("item-1").item();
+    final Set<byte[]> 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 @@ void testRecordIssuance() {
         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;
+  }
 }