Relax issuing constraints in IssuedReceiptManager

signalapp · Dec 10, 2024 · 18c9b17 · 18c9b17
1 parent 1970741
commit 18c9b17
Show file tree

Hide file tree

Showing 8 changed files with 120 additions and 21 deletions.
diff --git a/service/config/sample.yml b/service/config/sample.yml
@@ -127,6 +127,11 @@ dynamoDbTables:
     tableName: Example_IssuedReceipts
     expiration: P30D # Duration of time until rows expire
     generator: abcdefg12345678= # random base64-encoded binary sequence
+    maxIssuedReceiptsPerPaymentId:
+      STRIPE: 1
+      BRAINTREE: 1
+      GOOGLE_PLAY_BILLING: 1
+      APPLE_APP_STORE: 1
   ecKeys:
     tableName: Example_Keys
   ecSignedPreKeys:

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -647,7 +647,8 @@ public void run(WhisperServerConfiguration config, Environment environment) thro
         config.getDynamoDbTables().getIssuedReceipts().getTableName(),
         config.getDynamoDbTables().getIssuedReceipts().getExpiration(),
         dynamoDbAsyncClient,
-        config.getDynamoDbTables().getIssuedReceipts().getGenerator());
+        config.getDynamoDbTables().getIssuedReceipts().getGenerator(),
+        config.getDynamoDbTables().getIssuedReceipts().getmaxIssuedReceiptsPerPaymentId());
     OneTimeDonationsManager oneTimeDonationsManager = new OneTimeDonationsManager(
         config.getDynamoDbTables().getOnetimeDonations().getTableName(), config.getDynamoDbTables().getOnetimeDonations().getExpiration(), dynamoDbAsyncClient);
     RedeemedReceiptsManager redeemedReceiptsManager = new RedeemedReceiptsManager(clock,

diff --git a/...java/org/whispersystems/textsecuregcm/configuration/IssuedReceiptsTableConfiguration.java b/...java/org/whispersystems/textsecuregcm/configuration/IssuedReceiptsTableConfiguration.java
@@ -7,22 +7,37 @@
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 import jakarta.validation.constraints.NotEmpty;
+import org.whispersystems.textsecuregcm.subscriptions.PaymentProvider;
+import org.whispersystems.textsecuregcm.util.EnumMapUtil;
 import java.time.Duration;
+import java.util.EnumMap;
+import java.util.Map;
 
 public class IssuedReceiptsTableConfiguration extends DynamoDbTables.TableWithExpiration {
 
   private final byte[] generator;
 
+  /**
+   * The maximum number of issued receipts the issued receipt manager should issue for a particular itemId
+   */
+  private final EnumMap<PaymentProvider, Integer> maxIssuedReceiptsPerPaymentId;
+
   public IssuedReceiptsTableConfiguration(
       @JsonProperty("tableName") final String tableName,
       @JsonProperty("expiration") final Duration expiration,
-      @JsonProperty("generator") final byte[] generator) {
+      @JsonProperty("generator") final byte[] generator,
+      @JsonProperty("maxIssuedReceiptsPerPaymentId") final Map<PaymentProvider, Integer> maxIssuedReceiptsPerPaymentId) {
     super(tableName, expiration);
     this.generator = generator;
+    this.maxIssuedReceiptsPerPaymentId = EnumMapUtil.toCompleteEnumMap(PaymentProvider.class, maxIssuedReceiptsPerPaymentId);
   }
 
   @NotEmpty
   public byte[] getGenerator() {
     return generator;
   }
+
+  public EnumMap<PaymentProvider, Integer> getmaxIssuedReceiptsPerPaymentId() {
+    return maxIssuedReceiptsPerPaymentId;
+  }
 }
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManager.java
@@ -9,6 +9,7 @@
 import static org.whispersystems.textsecuregcm.util.AttributeValues.n;
 import static org.whispersystems.textsecuregcm.util.AttributeValues.s;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Throwables;
 import jakarta.ws.rs.ClientErrorException;
 import jakarta.ws.rs.core.Response.Status;
@@ -17,6 +18,7 @@
 import java.security.NoSuchAlgorithmException;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.EnumMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -46,16 +48,19 @@ public class IssuedReceiptsManager {
   private final Duration expiration;
   private final DynamoDbAsyncClient dynamoDbAsyncClient;
   private final byte[] receiptTagGenerator;
+  private final EnumMap<PaymentProvider, Integer> maxIssuedReceiptsPerPaymentId;
 
   public IssuedReceiptsManager(
       @Nonnull String table,
       @Nonnull Duration expiration,
       @Nonnull DynamoDbAsyncClient dynamoDbAsyncClient,
-      @Nonnull byte[] receiptTagGenerator) {
+      @Nonnull byte[] receiptTagGenerator,
+      @Nonnull EnumMap<PaymentProvider, Integer> maxIssuedReceiptsPerPaymentId) {
     this.table = Objects.requireNonNull(table);
     this.expiration = Objects.requireNonNull(expiration);
     this.dynamoDbAsyncClient = Objects.requireNonNull(dynamoDbAsyncClient);
     this.receiptTagGenerator = Objects.requireNonNull(receiptTagGenerator);
+    this.maxIssuedReceiptsPerPaymentId = Objects.requireNonNull(maxIssuedReceiptsPerPaymentId);
   }
 
   /**
@@ -74,19 +79,12 @@ public CompletableFuture<Void> recordIssuance(
       ReceiptCredentialRequest request,
       Instant now) {
 
-    final AttributeValue key;
-    if (processor == PaymentProvider.STRIPE) {
-      // As the first processor, Stripe’s IDs were not prefixed. Its item IDs have documented prefixes (`il_`, `pi_`)
-      // that will not collide with `SubscriptionProcessor` names
-      key = s(processorItemId);
-    } else {
-      key = s(processor.name() + "_" + processorItemId);
-    }
+    final AttributeValue key = dynamoDbKey(processor, 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")
+        .conditionExpression("attribute_not_exists(#key) OR contains(#tags, :tag) OR size(#tags) < :maxTags")
         .returnValues(ReturnValue.NONE)
         .updateExpression("SET "
             + "#tag = if_not_exists(#tag, :tag), "
@@ -100,7 +98,8 @@ public CompletableFuture<Void> recordIssuance(
         .expressionAttributeValues(Map.of(
             ":tag", b(tag),
             ":singletonTag", AttributeValue.fromBs(List.of(SdkBytes.fromByteArray(tag))),
-            ":exp", n(now.plus(expiration).getEpochSecond())))
+            ":exp", n(now.plus(expiration).getEpochSecond()),
+            ":maxTags", n(maxIssuedReceiptsPerPaymentId.get(processor))))
         .build();
     return dynamoDbAsyncClient.updateItem(updateItemRequest).handle((updateItemResponse, throwable) -> {
       if (throwable != null) {
@@ -115,7 +114,20 @@ public CompletableFuture<Void> recordIssuance(
     });
   }
 
-  private byte[] generateIssuedReceiptTag(ReceiptCredentialRequest request) {
+  @VisibleForTesting
+  static AttributeValue dynamoDbKey(final PaymentProvider processor, String processorItemId) {
+    if (processor == PaymentProvider.STRIPE) {
+      // As the first processor, Stripe’s IDs were not prefixed. Its item IDs have documented prefixes (`il_`, `pi_`)
+      // that will not collide with `SubscriptionProcessor` names
+      return s(processorItemId);
+    } else {
+      return s(processor.name() + "_" + processorItemId);
+    }
+  }
+
+
+  @VisibleForTesting
+  byte[] generateIssuedReceiptTag(ReceiptCredentialRequest request) {
     return generateHmac("issuedReceiptTag", mac -> mac.update(request.serialize()));
   }
 

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/util/EnumMapUtil.java b/service/src/main/java/org/whispersystems/textsecuregcm/util/EnumMapUtil.java
@@ -7,6 +7,7 @@
 
 import java.util.Arrays;
 import java.util.EnumMap;
+import java.util.Map;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -21,4 +22,13 @@ public static <E extends Enum<E>, V> EnumMap<E, V> toEnumMap(final Class<E> enum
             },
             () -> new EnumMap<>(enumClass)));
   }
+
+  public static <E extends Enum<E>, V> EnumMap<E, V> toCompleteEnumMap(final Class<E> enumClass, final Map<E, V> map) {
+    for (E e : enumClass.getEnumConstants()) {
+      if (!map.containsKey(e)) {
+        throw new IllegalArgumentException("Missing enum key: " + e);
+      }
+    }
+    return new EnumMap<>(map);
+  }
 }
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java
@@ -268,7 +268,8 @@ static CommandDependencies build(
         configuration.getDynamoDbTables().getIssuedReceipts().getTableName(),
         configuration.getDynamoDbTables().getIssuedReceipts().getExpiration(),
         dynamoDbAsyncClient,
-        configuration.getDynamoDbTables().getIssuedReceipts().getGenerator());
+        configuration.getDynamoDbTables().getIssuedReceipts().getGenerator(),
+        configuration.getDynamoDbTables().getIssuedReceipts().getmaxIssuedReceiptsPerPaymentId());
 
     APNSender apnSender = new APNSender(apnSenderExecutor, configuration.getApnConfiguration());
     FcmSender fcmSender = new FcmSender(fcmSenderExecutor, configuration.getFcmConfiguration().credentials().value());

diff --git a/...ice/src/test/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManagerTest.java b/...ice/src/test/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManagerTest.java
@@ -12,18 +12,22 @@
 import jakarta.ws.rs.ClientErrorException;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.EnumMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 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.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
 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;
@@ -38,15 +42,22 @@ class IssuedReceiptsManagerTest {
   @RegisterExtension
   static final DynamoDbExtension DYNAMO_DB_EXTENSION = new DynamoDbExtension(Tables.ISSUED_RECEIPTS);
 
-  IssuedReceiptsManager issuedReceiptsManager;
+  private static EnumMap<PaymentProvider, Integer> MAX_TAGS_MAP = new EnumMap<>(Map.of(
+      PaymentProvider.STRIPE, 1,
+      PaymentProvider.BRAINTREE, 2,
+      PaymentProvider.GOOGLE_PLAY_BILLING, 3,
+      PaymentProvider.APPLE_APP_STORE, 4));
+
+  private IssuedReceiptsManager issuedReceiptsManager;
 
   @BeforeEach
   void beforeEach() {
     issuedReceiptsManager = new IssuedReceiptsManager(
         Tables.ISSUED_RECEIPTS.tableName(),
         Duration.ofDays(90),
         DYNAMO_DB_EXTENSION.getDynamoDbAsyncClient(),
-        TestRandomUtil.nextBytes(16));
+        TestRandomUtil.nextBytes(16),
+        MAX_TAGS_MAP);
   }
 
   @Test
@@ -57,7 +68,7 @@ void testRecordIssuance() {
         receiptCredentialRequest, now);
     assertThat(future).succeedsWithin(Duration.ofSeconds(3));
 
-    final Map<String, AttributeValue> item = getItem("item-1").item();
+    final Map<String, AttributeValue> item = getItem(PaymentProvider.STRIPE, "item-1").item();
     final Set<byte[]> tagSet = item.get(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG_SET).bs()
         .stream()
         .map(SdkBytes::asByteArray)
@@ -88,12 +99,51 @@ void testRecordIssuance() {
     assertThat(future).succeedsWithin(Duration.ofSeconds(3));
   }
 
+  @ParameterizedTest
+  @EnumSource(PaymentProvider.class)
+  void testIssueMax(PaymentProvider processor) {
+    final Instant now = Instant.ofEpochSecond(NOW_EPOCH_SECONDS);
+
+    final int maxTags = MAX_TAGS_MAP.get(processor);
+    final List<ReceiptCredentialRequest> requests = IntStream.range(0, maxTags)
+        .mapToObj(i -> randomReceiptCredentialRequest())
+        .toList();
+    for (int i = 0; i < maxTags; i++) {
+      // Should be allowed to insert up to maxTags
+        assertThat(issuedReceiptsManager.recordIssuance("item-1", processor, requests.get(i), now))
+            .succeedsWithin(Duration.ofSeconds(3));
+      for (int j = 0; j < i; j++) {
+        // Also should be allowed to repeat any previous tag
+        assertThat(issuedReceiptsManager.recordIssuance("item-1", processor, requests.get(j), now))
+            .succeedsWithin(Duration.ofSeconds(3));
+      }
+    }
+
+    assertThat(getItem(processor, "item-1").item().get(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG_SET).bs()
+        .stream()
+        .map(SdkBytes::asByteArray)
+        .collect(Collectors.toSet()))
+        .containsExactlyInAnyOrder(requests.stream()
+            .map(issuedReceiptsManager::generateIssuedReceiptTag)
+            .toArray(byte[][]::new));
+
+    // Should not be allowed to insert past maxTags
+    assertThat(issuedReceiptsManager.recordIssuance("item-1", processor, randomReceiptCredentialRequest(), now))
+        .failsWithin(Duration.ofSeconds(3))
+        .withThrowableOfType(Throwable.class)
+        .havingCause()
+        .isExactlyInstanceOf(ClientErrorException.class)
+        .has(new Condition<>(
+            e -> e instanceof ClientErrorException && ((ClientErrorException) e).getResponse().getStatus() == 409,
+            "status 409"));
+  }
+
 
-  private GetItemResponse getItem(final String itemId) {
+  private GetItemResponse getItem(final PaymentProvider processor, 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)))
+        .key(Map.of(IssuedReceiptsManager.KEY_PROCESSOR_ITEM_ID, IssuedReceiptsManager.dynamoDbKey(processor, itemId)))
         .build());
   }
 

diff --git a/service/src/test/resources/config/test.yml b/service/src/test/resources/config/test.yml
@@ -124,6 +124,11 @@ dynamoDbTables:
     tableName: issued_receipts_test
     expiration: P30D # Duration of time until rows expire
     generator: abcdefg12345678= # random base64-encoded binary sequence
+    maxIssuedReceiptsPerPaymentId:
+      STRIPE: 1
+      BRAINTREE: 1
+      GOOGLE_PLAY_BILLING: 1
+      APPLE_APP_STORE: 1
   ecKeys:
     tableName: keys_test
   ecSignedPreKeys: